from pathlib import Path
from typing import Annotated

from rich.console import Console
from rich.table import Table
from typer import Argument, Option, Typer

datasets_app = Typer(no_args_is_help=True)
console = Console()


@datasets_app.command()
def list(
    registry_url: Annotated[
        str | None,
        Option(
            "--registry-url",
            help="Registry URL for remote dataset listing",
            show_default="The default harbor registry.",
        ),
    ] = None,
    registry_path: Annotated[
        Path | None,
        Option(
            "--registry-path",
            help="Path to local registry for dataset listing",
            show_default=False,
        ),
    ] = None,
    verbose: Annotated[
        bool,
        Option(
            "-v",
            "--verbose",
            help="Show detailed information including task counts",
        ),
    ] = False,
):
    """List all datasets available in a registry.

    By default, uses the Harbor default registry. You can specify a custom
    registry using --registry-url or --registry-path.
    """
    from harbor.models.registry import LocalRegistryInfo, RemoteRegistryInfo
    from harbor.registry.client import RegistryClient

    try:
        # Create registry client based on provided options
        if registry_url is not None and registry_path is not None:
            console.print(
                "[red]Error: Cannot specify both --registry-url and --registry-path[/red]"
            )
            return

        if registry_path is not None:
            registry_info = LocalRegistryInfo(path=registry_path)
            client = RegistryClient.from_info(registry_info)
            console.print(f"[blue]Using local registry: {registry_path}[/blue]\n")
        elif registry_url is not None:
            registry_info = RemoteRegistryInfo(url=registry_url)
            client = RegistryClient.from_info(registry_info)
            console.print(f"[blue]Using remote registry: {registry_url}[/blue]\n")
        else:
            client = RegistryClient()
            console.print(
                f"[blue]Using default registry: {RegistryClient.DEFAULT_REGISTRY_URL}[/blue]\n"
            )

        # Get datasets from the registry
        datasets = client.datasets

        if not datasets:
            console.print("[yellow]No datasets found in registry[/yellow]")
            return

        # Create table
        if verbose:
            table = Table(title="Available Datasets", show_lines=True)
            table.add_column("Name", style="cyan", no_wrap=True)
            table.add_column("Version", style="magenta")
            table.add_column("Tasks", style="green", justify="right")
            table.add_column("Description", style="white")
        else:
            table = Table(title="Available Datasets")
            table.add_column("Name", style="cyan", no_wrap=True)
            table.add_column("Version", style="magenta")
            table.add_column("Description", style="white")

        # Add rows for each dataset
        total_datasets = 0
        total_tasks = 0

        for dataset_name in sorted(datasets.keys()):
            versions = datasets[dataset_name]
            for version in sorted(versions.keys()):
                dataset = versions[version]
                task_count = len(dataset.tasks)
                total_datasets += 1
                total_tasks += task_count

                if verbose:
                    table.add_row(
                        dataset_name,
                        version,
                        str(task_count),
                        dataset.description,
                    )
                else:
                    table.add_row(
                        dataset_name,
                        version,
                        dataset.description,
                    )

        console.print(table)
        console.print(
            f"\n[green]Total: {total_datasets} dataset(s)"
            + (f" with {total_tasks} task(s)" if verbose else "")
            + "[/green]"
        )

        if not verbose:
            console.print(
                "[dim]Tip: Use --verbose to see task counts for each dataset[/dim]"
            )

    except Exception as e:
        console.print(f"[red]Error: {e}[/red]")
        raise


@datasets_app.command()
def download(
    dataset: Annotated[
        str,
        Argument(
            help="Dataset to download in format 'name@version' or 'name' (defaults to @head)"
        ),
    ],
    registry_url: Annotated[
        str | None,
        Option(
            "--registry-url",
            help="Registry URL for remote dataset",
            show_default="The default harbor registry.",
        ),
    ] = None,
    registry_path: Annotated[
        Path | None,
        Option(
            "--registry-path",
            help="Path to local registry",
            show_default=False,
        ),
    ] = None,
    output_dir: Annotated[
        Path | None,
        Option(
            "--output-dir",
            "-o",
            help="Directory to download tasks to",
            show_default="~/.cache/harbor/tasks",
        ),
    ] = None,
    overwrite: Annotated[
        bool,
        Option(
            "--overwrite",
            help="Overwrite cached tasks",
        ),
    ] = False,
):
    """Download a dataset from a registry.

    This command downloads all tasks in a dataset to the local cache.
    Tasks are fetched using shallow clones with sparse checkout for efficiency.

    Examples:
        harbor datasets download my-dataset
        harbor datasets download my-dataset@v1.0
        harbor datasets download my-dataset@head --overwrite
    """
    from harbor.dataset.client import DatasetClient
    from harbor.models.job.config import RegistryDatasetConfig
    from harbor.models.registry import LocalRegistryInfo, RemoteRegistryInfo
    from harbor.registry.client import RegistryClient

    if "@" in dataset:
        name, version = dataset.split("@", 1)
    else:
        name = dataset
        version = "head"

    if registry_url is not None and registry_path is not None:
        console.print(
            "[red]Error: Cannot specify both --registry-url and --registry-path[/red]"
        )
        return

    if registry_path is not None:
        registry_info = LocalRegistryInfo(path=registry_path)
        console.print(f"[blue]Using local registry: {registry_path}[/blue]")
    elif registry_url is not None:
        registry_info = RemoteRegistryInfo(url=registry_url)
        console.print(f"[blue]Using remote registry: {registry_url}[/blue]")
    else:
        registry_info = RemoteRegistryInfo(url=RegistryClient.DEFAULT_REGISTRY_URL)
        console.print(
            f"[blue]Using default registry: {RegistryClient.DEFAULT_REGISTRY_URL}[/blue]"
        )

    config = RegistryDatasetConfig(
        registry=registry_info,
        name=name,
        version=version,
        overwrite=overwrite,
        download_dir=output_dir,
    )

    console.print(f"[cyan]Downloading dataset: {name} (version: {version})[/cyan]")

    client = DatasetClient()

    try:
        with console.status("[bold green]Downloading tasks..."):
            downloaded_tasks = client.download_dataset_from_config(config)
    except KeyError:
        console.print(
            f"[red]Error: Dataset '{name}' (version: '{version}') not found in registry[/red]"
        )
        return
    except Exception as e:
        console.print(f"[red]Error: {e}[/red]")
        raise

    console.print(
        f"\n[green]Successfully downloaded {len(downloaded_tasks)} task(s)[/green]"
    )
