import asyncio
import hashlib
import tarfile
import tempfile
import time
from io import BytesIO
from pathlib import Path
from collections.abc import Callable
from typing import Any

from pydantic import BaseModel
from storage3.exceptions import StorageApiError

from harbor.models.dataset.manifest import DatasetManifest
from harbor.models.dataset.paths import DatasetPaths
from harbor.models.task.config import TaskConfig
from harbor.models.task.paths import TaskPaths
from harbor.publisher.packager import Packager
from harbor.registry.db import RegistryDB
from harbor.storage.supabase import SupabaseStorage

ARCHIVE_FILENAME = "archive.harbor"
PACKAGE_DIR = "packages"


class PublishResult(BaseModel):
    name: str
    content_hash: str
    archive_path: str
    file_count: int
    archive_size_bytes: int
    build_time_sec: float
    upload_time_sec: float
    rpc_time_sec: float = 0.0
    skipped: bool = False
    revision: int | None = None
    tags: list[str] | None = None
    db_skipped: bool = False


class BatchPublishResult(BaseModel):
    results: list[PublishResult]
    total_time_sec: float


class DatasetPublishResult(BaseModel):
    name: str
    content_hash: str
    revision: int
    task_count: int
    file_count: int
    skipped: bool = False
    db_skipped: bool = False
    rpc_time_sec: float = 0.0
    tags: list[str] | None = None


class FilePublishResult(BaseModel):
    content_hash: str
    remote_path: str
    file_size_bytes: int
    upload_time_sec: float
    skipped: bool = False


class Publisher:
    def __init__(self) -> None:
        self.storage = SupabaseStorage()
        self.registry_db = RegistryDB()

    @staticmethod
    def _create_archive(task_dir: Path, files: list[Path], dest: Path) -> None:
        with tarfile.open(dest, "w:gz") as tar:
            for f in files:
                rel = f.relative_to(task_dir).as_posix()
                data = f.read_bytes()
                info = tarfile.TarInfo(name=rel)
                info.size = len(data)
                info.uid = 0
                info.gid = 0
                info.uname = ""
                info.gname = ""
                info.mtime = 0
                info.mode = 0o644
                tar.addfile(info, BytesIO(data))

    async def publish_file(
        self, package_name: str, file_path: Path
    ) -> FilePublishResult:
        data = file_path.read_bytes()
        content_hash = Packager.compute_file_hash(file_path)
        remote_path = f"{PACKAGE_DIR}/{package_name}/{content_hash}/{file_path.name}"
        skipped = False
        upload_start = time.monotonic()
        try:
            await self.storage.upload_file(file_path, remote_path)
        except StorageApiError as exc:
            if exc.status in (409, "409"):
                skipped = True
            elif exc.status in (403, "403"):
                raise PermissionError(
                    "You don't have permission to publish to this package. "
                    "Either the organization doesn't exist or you are not a member with publish rights."
                ) from exc
            else:
                raise
        upload_time = time.monotonic() - upload_start
        return FilePublishResult(
            content_hash=content_hash,
            remote_path=remote_path,
            file_size_bytes=len(data),
            upload_time_sec=round(upload_time, 3),
            skipped=skipped,
        )

    async def publish_task(
        self,
        task_dir: Path,
        tags: set[str] | None = None,
        visibility: str = "public",
    ) -> PublishResult:
        paths = TaskPaths(task_dir)
        if not paths.config_path.exists():
            raise FileNotFoundError(f"task.toml not found in {task_dir}")

        config = TaskConfig.model_validate_toml(paths.config_path.read_text())
        if config.task is None:
            raise ValueError("task.toml must contain a [task] section with a name")

        if not paths.is_valid():
            raise ValueError(
                f"Task directory {task_dir} is not valid (missing required files: instruction.md, environment/, or tests/test.sh)"
            )

        build_start = time.monotonic()
        content_hash, files = Packager.compute_content_hash(task_dir)

        skipped = False
        with tempfile.TemporaryDirectory() as tmp:
            archive_path = Path(tmp) / ARCHIVE_FILENAME
            self._create_archive(task_dir, files, archive_path)
            archive_size = archive_path.stat().st_size
            build_time = time.monotonic() - build_start

            remote_path = (
                f"{PACKAGE_DIR}/{config.task.name}/{content_hash}/{ARCHIVE_FILENAME}"
            )
            upload_start = time.monotonic()
            try:
                await self.storage.upload_file(archive_path, remote_path)
            except StorageApiError as exc:
                if exc.status in (409, "409"):
                    skipped = True
                elif exc.status in (403, "403"):
                    raise PermissionError(
                        "You don't have permission to publish to this package. "
                        "Either the organization doesn't exist or you are not a member with publish rights."
                    ) from exc
                else:
                    raise
            upload_time = time.monotonic() - upload_start

        # DB registration via RPC
        db = self.registry_db

        instruction = ""
        if paths.instruction_path.exists():
            instruction = paths.instruction_path.read_text()

        readme = ""
        if paths.readme_path.exists():
            readme = paths.readme_path.read_text()

        applied_tags = {"latest"} | (tags or set())

        file_rows = []
        for f in files:
            rel = f.relative_to(task_dir).as_posix()
            file_data = f.read_bytes()
            fhash = f"{hashlib.sha256(file_data).hexdigest()}"
            file_rows.append(
                {
                    "path": rel,
                    "content_hash": fhash,
                    "size_bytes": len(file_data),
                }
            )

        rpc_start = time.monotonic()
        rpc_result = await db.publish_task_version(
            org=config.task.org,
            name=config.task.short_name,
            tags=sorted(applied_tags),
            content_hash=content_hash,
            archive_path=remote_path,
            description=config.task.description,
            authors=[a.model_dump(mode="json") for a in config.task.authors],
            keywords=config.task.keywords,
            metadata=config.metadata,
            verifier_config=config.verifier.model_dump(mode="json"),
            agent_config=config.agent.model_dump(mode="json"),
            environment_config=config.environment.model_dump(mode="json"),
            instruction=instruction,
            readme=readme,
            files=file_rows,
            visibility=visibility,
        )

        rpc_time = time.monotonic() - rpc_start
        created = rpc_result.get("created", True)
        revision = rpc_result.get("revision") if created else None

        return PublishResult(
            name=config.task.name,
            content_hash=content_hash,
            archive_path=remote_path,
            file_count=len(files),
            archive_size_bytes=archive_size,
            build_time_sec=round(build_time, 3),
            upload_time_sec=round(upload_time, 3),
            rpc_time_sec=round(rpc_time, 3),
            skipped=skipped,
            revision=revision,
            tags=sorted(applied_tags),
            db_skipped=not created,
        )

    async def publish_tasks(
        self,
        task_dirs: list[Path],
        *,
        max_concurrency: int = 100,
        tags: set[str] | None = None,
        visibility: str = "public",
        on_task_upload_start: Callable[[Path], Any] | None = None,
        on_task_upload_complete: Callable[[Path, PublishResult], Any] | None = None,
    ) -> BatchPublishResult:
        if not task_dirs:
            return BatchPublishResult(results=[], total_time_sec=0.0)

        start = time.monotonic()
        results: list[PublishResult | None] = [None] * len(task_dirs)
        semaphore = asyncio.Semaphore(max_concurrency)

        async def _publish(index: int, task_dir: Path) -> None:
            async with semaphore:
                if on_task_upload_start:
                    on_task_upload_start(task_dir)
                result = await self.publish_task(
                    task_dir, tags=tags, visibility=visibility
                )
                results[index] = result
                if on_task_upload_complete:
                    on_task_upload_complete(task_dir, result)

        async with asyncio.TaskGroup() as tg:
            for i, task_dir in enumerate(task_dirs):
                tg.create_task(_publish(i, task_dir))

        return BatchPublishResult(
            results=[r for r in results if r is not None],
            total_time_sec=round(time.monotonic() - start, 3),
        )

    async def publish_dataset(
        self,
        dataset_dir: Path,
        tags: set[str] | None = None,
        visibility: str = "public",
        promote_tasks: bool = False,
    ) -> DatasetPublishResult:
        paths = DatasetPaths(dataset_dir)
        if not paths.manifest_path.exists():
            raise FileNotFoundError(f"dataset.toml not found in {dataset_dir}")

        manifest = DatasetManifest.from_toml_file(paths.manifest_path)
        db = self.registry_db

        # Upload dataset-level files to storage before RPC call
        file_infos: list[dict[str, Any]] = []
        for file_ref in manifest.files:
            file_path = dataset_dir / file_ref.path
            if not file_path.exists():
                raise FileNotFoundError(
                    f"Dataset file '{file_ref.path}' not found in {dataset_dir}"
                )
            result = await self.publish_file(manifest.dataset.name, file_path)
            file_infos.append(
                {
                    "path": file_ref.path,
                    "content_hash": result.content_hash,
                    "size_bytes": result.file_size_bytes,
                    "storage_path": result.remote_path,
                }
            )

        applied_tags = {"latest"} | (tags or set())

        task_refs = [
            {"org": ref.org, "name": ref.short_name, "digest": ref.digest}
            for ref in manifest.tasks
        ]

        rpc_start = time.monotonic()
        rpc_result = await db.publish_dataset_version(
            org=manifest.dataset.org,
            name=manifest.dataset.short_name,
            tags=sorted(applied_tags),
            description=manifest.dataset.description,
            authors=[a.model_dump(mode="json") for a in manifest.dataset.authors],
            tasks=task_refs,
            files=file_infos,
            visibility=visibility,
            promote_tasks=promote_tasks,
        )
        rpc_time = time.monotonic() - rpc_start

        created = rpc_result.get("created", True)
        content_hash = rpc_result.get("content_hash", "")

        return DatasetPublishResult(
            name=manifest.dataset.name,
            content_hash=content_hash,
            revision=rpc_result.get("revision", 0),
            task_count=manifest.task_count,
            file_count=len(file_infos),
            skipped=not created,
            db_skipped=not created,
            rpc_time_sec=round(rpc_time, 3),
            tags=sorted(applied_tags),
        )
