from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING, Annotated

from postgrest.exceptions import APIError
from typer import Argument, Option, echo

if TYPE_CHECKING:
    from rich.console import Console

    from harbor.publisher.publisher import Publisher


def _humanize_bytes(n: int) -> str:
    for unit in ("B", "KB", "MB", "GB"):
        if n < 1024:
            return f"{n:.1f} {unit}" if unit != "B" else f"{n} {unit}"
        n /= 1024  # type: ignore[assignment]
    return f"{n:.1f} TB"


def _resolve_paths(
    paths: list[Path], *, no_tasks: bool = False
) -> tuple[list[Path], list[Path], set[Path], set[Path]]:
    """Resolve paths to task and dataset directories by auto-detecting type.

    Returns (task_dirs, dataset_dirs, explicit_task_dirs, explicit_dataset_dirs).
    ``explicit_task_dirs`` contains only those task dirs that the user directly
    specified (not auto-discovered from dataset or parent directories).
    ``explicit_dataset_dirs`` contains only those dataset dirs that the user
    directly specified.
    """
    from harbor.models.dataset.paths import DatasetPaths
    from harbor.models.task.paths import TaskPaths

    task_dirs: list[Path] = []
    dataset_dirs: list[Path] = []
    explicit_task_dirs: set[Path] = set()
    explicit_dataset_dirs: set[Path] = set()

    for p in paths:
        resolved = p.resolve()
        # dataset.toml file passed directly
        if resolved.is_file() and resolved.name == DatasetPaths.MANIFEST_FILENAME:
            dataset_dirs.append(resolved.parent)
            explicit_dataset_dirs.add(resolved.parent)
            continue
        if not resolved.is_dir():
            echo(f"Warning: {p} is not a directory, skipping.")
            continue
        has_task = (resolved / TaskPaths.CONFIG_FILENAME).exists()
        has_dataset = (resolved / DatasetPaths.MANIFEST_FILENAME).exists()
        if has_task and has_dataset:
            echo(
                f"Error: {p} contains both {TaskPaths.CONFIG_FILENAME} and "
                f"{DatasetPaths.MANIFEST_FILENAME}. Cannot determine type."
            )
            raise SystemExit(1)
        if has_task:
            task_dirs.append(resolved)
            explicit_task_dirs.add(resolved)
        elif has_dataset:
            dataset_dirs.append(resolved)
            explicit_dataset_dirs.add(resolved)
            # Also collect task subdirs so they are published before the dataset
            if not no_tasks:
                for child in sorted(resolved.iterdir()):
                    if child.is_dir() and (child / TaskPaths.CONFIG_FILENAME).exists():
                        task_dirs.append(child)
        else:
            # Scan immediate subdirectories
            found_any = False
            for child in sorted(resolved.iterdir()):
                if not child.is_dir():
                    continue
                if (child / TaskPaths.CONFIG_FILENAME).exists():
                    task_dirs.append(child)
                    explicit_task_dirs.add(child)
                    found_any = True
            if not found_any:
                echo(f"Warning: {p} contains no tasks, skipping.")

    return task_dirs, dataset_dirs, explicit_task_dirs, explicit_dataset_dirs


def publish_command(
    paths: Annotated[
        list[Path] | None,
        Argument(help="Task or dataset directories to publish."),
    ] = None,
    tag: Annotated[
        list[str] | None,
        Option(
            "--tag",
            "-t",
            help="Tag(s) to apply (repeatable). 'latest' is always added.",
        ),
    ] = None,
    concurrency: Annotated[
        int, Option("--concurrency", "-c", help="Max concurrent uploads.")
    ] = 50,
    no_tasks: Annotated[
        bool, Option("--no-tasks", help="Skip publishing tasks for datasets.")
    ] = False,
    public: Annotated[
        bool, Option("--public/--private", help="Set visibility (default: private).")
    ] = False,
    debug: Annotated[
        bool,
        Option(
            "--debug", help="Show extra timing details (e.g. RPC time).", hidden=True
        ),
    ] = False,
) -> None:
    """Publish tasks and datasets to the Harbor registry."""
    if paths is None:
        paths = [Path(".")]

    from rich.console import Console

    from harbor.cli.utils import run_async
    from harbor.publisher.publisher import Publisher

    console = Console()

    async def _publish() -> None:
        publisher = Publisher()

        # Auth check
        try:
            await publisher.registry_db.get_user_id()
        except RuntimeError as exc:
            echo(str(exc))
            raise SystemExit(1)

        task_dirs, dataset_dirs, explicit_task_dirs, explicit_dataset_dirs = (
            _resolve_paths(paths, no_tasks=no_tasks)
        )
        if not task_dirs and not dataset_dirs:
            echo("No tasks or datasets found.")
            raise SystemExit(1)

        visibility = "public" if public else "private"
        tags = set(tag) if tag else None
        if task_dirs:
            await _publish_tasks(
                publisher,
                console,
                task_dirs,
                tags,
                concurrency,
                visibility,
                debug,
                explicit_task_dirs=explicit_task_dirs,
            )
        if dataset_dirs:
            await _publish_datasets(
                publisher,
                console,
                dataset_dirs,
                tags,
                visibility,
                explicit_dataset_dirs=explicit_dataset_dirs,
            )

    run_async(_publish())


async def _publish_tasks(
    publisher: Publisher,
    console: Console,
    task_dirs: list[Path],
    tags: set[str] | None,
    concurrency: int,
    visibility: str,
    debug: bool = False,
    *,
    explicit_task_dirs: set[Path] | None = None,
) -> None:
    from rich.console import Group
    from rich.live import Live
    from rich.progress import (
        BarColumn,
        MofNCompleteColumn,
        Progress,
        SpinnerColumn,
        TextColumn,
        TimeElapsedColumn,
        TimeRemainingColumn,
    )
    from rich.progress import TaskID as ProgressTaskID
    from rich.table import Table

    from harbor.publisher.publisher import PublishResult

    overall_progress = Progress(
        SpinnerColumn(),
        MofNCompleteColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TimeElapsedColumn(),
        TimeRemainingColumn(),
    )
    running_progress = Progress(
        SpinnerColumn(),
        TimeElapsedColumn(),
        TextColumn("[progress.description]{task.description}"),
    )

    running_tasks: dict[str, ProgressTaskID] = {}
    explicit_result_names: list[str] = []

    def on_start(task_dir: Path) -> None:
        running_tasks[str(task_dir)] = running_progress.add_task(
            f"{task_dir.name}: publishing...", total=None
        )

    def on_complete(task_dir: Path, result: PublishResult) -> None:
        key = str(task_dir)
        if key in running_tasks:
            running_progress.remove_task(running_tasks.pop(key))
        overall_progress.advance(overall_task)
        if explicit_task_dirs and task_dir in explicit_task_dirs:
            explicit_result_names.append(result.name)

    with Live(
        Group(overall_progress, running_progress),
        console=console,
        refresh_per_second=10,
    ):
        overall_task = overall_progress.add_task(
            "Publishing tasks...", total=len(task_dirs)
        )
        batch = await publisher.publish_tasks(
            task_dirs,
            max_concurrency=concurrency,
            tags=tags,
            visibility=visibility,
            on_task_upload_start=on_start,
            on_task_upload_complete=on_complete,
        )

    table = Table()
    table.add_column("Task")
    table.add_column("Hash", max_width=20)
    table.add_column("Rev", justify="right")
    table.add_column("Files", justify="right")
    table.add_column("Size", justify="right")
    table.add_column("Build", justify="right")
    table.add_column("Upload", justify="right")
    if debug:
        table.add_column("RPC", justify="right")

    for r in batch.results:
        short_hash = r.content_hash.split(":")[-1][:12]
        rev_str = (
            str(r.revision)
            if r.revision is not None
            else ("exists" if r.db_skipped else "-")
        )
        row = [
            r.name,
            short_hash,
            rev_str,
            str(r.file_count),
            _humanize_bytes(r.archive_size_bytes),
            f"{r.build_time_sec:.2f}s",
            "skipped" if r.skipped else f"{r.upload_time_sec:.2f}s",
        ]
        if debug:
            row.append(f"{r.rpc_time_sec:.2f}s")
        table.add_row(*row)

    console.print(table)
    published = sum(1 for r in batch.results if not r.skipped or not r.db_skipped)
    skipped = sum(1 for r in batch.results if r.skipped and r.db_skipped)
    parts = [f"Published {published}"]
    if skipped:
        parts.append(f"skipped {skipped}")
        echo(f"\n{', '.join(parts)} task(s) in {batch.total_time_sec:.2f}s")

    if explicit_result_names:
        from harbor.constants import HARBOR_REGISTRY_TASKS_URL

        for name in explicit_result_names:
            echo(f"{HARBOR_REGISTRY_TASKS_URL}/{name}")


async def _publish_datasets(
    publisher: Publisher,
    console: Console,
    dataset_dirs: list[Path],
    tags: set[str] | None,
    visibility: str,
    *,
    explicit_dataset_dirs: set[Path] | None = None,
) -> None:
    from rich.table import Table

    echo(f"Publishing {len(dataset_dirs)} dataset(s)...")
    table = Table()
    table.add_column("Dataset")
    table.add_column("Hash", max_width=20)
    table.add_column("Rev", justify="right")
    table.add_column("Tasks", justify="right")
    table.add_column("Files", justify="right")
    table.add_column("Status")

    from harbor.cli.sync import sync_dataset
    from harbor.publisher.publisher import DatasetPublishResult

    published = 0
    skipped = 0
    results: list[tuple[Path, DatasetPublishResult]] = []
    for dataset_dir in dataset_dirs:
        # Auto-sync local digests before publishing
        sync_changes = sync_dataset(dataset_dir)
        updated = sum(1 for c in sync_changes if c.old != c.new)
        if updated:
            echo(f"Synced {updated} digest(s) in dataset.toml")

        promote_tasks = False
        if visibility == "public":
            answer = console.input(
                f'Publishing dataset "{dataset_dir.name}" as public will make all its tasks public. Proceed? (y/N): '
            )
            if answer.strip().lower() != "y":
                echo(f"Skipping {dataset_dir.name}.")
                skipped += 1
                continue
            promote_tasks = True
        try:
            result = await publisher.publish_dataset(
                dataset_dir,
                tags=tags,
                visibility=visibility,
                promote_tasks=promote_tasks,
            )
        except APIError as exc:
            from rich.console import Console

            error_message = exc.message or str(exc)
            Console().print(
                f"[red]Failed to publish dataset at {dataset_dir}: {error_message}[/red]"
            )
            raise SystemExit(1)
        short_hash = result.content_hash[:12]
        status = "skipped (exists)" if result.skipped else "published"
        table.add_row(
            result.name,
            short_hash,
            str(result.revision) if not result.skipped else "-",
            str(result.task_count),
            str(result.file_count),
            status,
        )
        results.append((dataset_dir, result))
        if result.skipped:
            skipped += 1
        else:
            published += 1

    console.print(table)
    parts = [f"Published {published}"]
    if skipped:
        parts.append(f"skipped {skipped}")
    echo(f"\n{', '.join(parts)} dataset(s)")

    if explicit_dataset_dirs:
        from harbor.constants import HARBOR_REGISTRY_DATASETS_URL

        for dataset_dir, result in results:
            if dataset_dir in explicit_dataset_dirs:
                echo(f"{HARBOR_REGISTRY_DATASETS_URL}/{result.name}")
