"""
cli.py — comfy-resolve command-line interface

Entry points:
  resolve       Parse a workflow, resolve assets, prompt for review, download.
  cache-list    Show all cached asset mappings.
  cache-clear   Remove one or all cached mappings.
  config-init   Run (or re-run) first-time setup wizard.
  config-show   Print active config with API keys redacted.
"""

from __future__ import annotations

import json
import logging
import traceback
from pathlib import Path
from typing import Any, Optional

import toml  # type: ignore
import typer
from rich.console import Console
from rich.panel import Panel
from rich.prompt import Prompt
from rich.rule import Rule
from rich.table import Table
from rich import box

from .config import config_exists, get_model_path, load_config, write_default_config
from .downloader import DownloadTask, Downloader
from .parser import parse_workflow
from .resolver import Resolver
from .sources.base import ResolutionResult
from .storage import delete_mapping, list_mappings, save_mapping, save_mappings
from .ui import UISession

# ---------------------------------------------------------------------------
# App & module-level singletons
# ---------------------------------------------------------------------------

app = typer.Typer(
    name="comfy-resolve",
    help="Resolve and download model dependencies for ComfyUI workflows.",
    add_completion=False,
    rich_markup_mode="rich",
)
console = Console()
logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# resolve
# ---------------------------------------------------------------------------

@app.command()
def resolve(
    workflow: Optional[Path] = typer.Argument(
        None,
        help=(
            "Path to the ComfyUI workflow JSON file. "
            "If omitted, workflows are discovered from your ComfyUI install directory."
        ),
        exists=False,           # checked manually for nicer error messages
        file_okay=True,
        dir_okay=False,
        readable=True,
        resolve_path=True,
    ),
    auto: bool = typer.Option(
        False, "--auto",
        help="Skip prompts; use best guess for HIGH-confidence assets only.",
    ),
    auto_all: bool = typer.Option(
        False, "--auto-all",
        help="Skip prompts; use best guess for ALL assets regardless of confidence.",
    ),
    dry_run: bool = typer.Option(
        False, "--dry-run",
        help="Resolve and show what would be downloaded — no files are written.",
    ),
    no_cache: bool = typer.Option(
        False, "--no-cache",
        help="Ignore cached mappings and re-resolve every asset from scratch.",
    ),
    download_dir: Optional[Path] = typer.Option(
        None, "--download-dir",
        help="Override the destination directory for every asset type.",
    ),
    concurrent: Optional[int] = typer.Option(
        None, "--concurrent",
        min=1, max=16,
        help="Number of parallel downloads (overrides config value).",
    ),
    verbose: bool = typer.Option(
        False, "--verbose", "-v",
        help="Emit DEBUG-level logs.",
    ),
) -> None:
    """
    [bold cyan]Resolve and download[/bold cyan] all model assets referenced in a ComfyUI
    workflow JSON file.

    Resolution order per asset:
      1. Local filesystem  2. Cache  3. Config mappings
      4. HuggingFace       5. Civitai (token required)  6. Heuristic

    If no workflow path is supplied, workflows are discovered automatically
    from your ComfyUI install's [dim]user/default/workflows[/dim] folder and
    presented as a numbered pick-list.

    After resolution you are shown a full pre-download review where you can
    edit sources, skip individual assets, or abort entirely before any file
    is written to disk.
    """
    _setup_logging(verbose)

    try:
        _resolve_workflow(
            workflow=workflow,
            auto=auto,
            auto_all=auto_all,
            dry_run=dry_run,
            no_cache=no_cache,
            download_dir=download_dir,
            concurrent=concurrent,
            verbose=verbose,
        )
    except json.JSONDecodeError as exc:
        console.print(f"[red]✗ Invalid JSON in {workflow}[/red]")
        console.print(f"[red]  Line {exc.lineno}, col {exc.colno}: {exc.msg}[/red]")
        if verbose:
            console.print(f"[dim]{exc}[/dim]")
        raise typer.Exit(code=1)
    except KeyboardInterrupt:
        console.print("\n[yellow]Interrupted.[/yellow]")
        raise typer.Exit(code=130)
    except typer.Exit:
        raise
    except Exception as exc:
        console.print(f"[red]✗ Unexpected error: {exc}[/red]")
        if verbose:
            console.print(f"[dim]{traceback.format_exc()}[/dim]")
        else:
            console.print("[dim]Re-run with --verbose for a full traceback.[/dim]")
        raise typer.Exit(code=1)


def _resolve_workflow(
    workflow: Optional[Path],
    auto: bool,
    auto_all: bool,
    dry_run: bool,
    no_cache: bool,
    download_dir: Optional[Path],
    concurrent: Optional[int],
    verbose: bool,
) -> None:
    """
    Core resolve pipeline — separated from the @app.command so that the
    exception handlers in the command wrapper stay clean and focused.
    """

    # ── 1. Config ──────────────────────────────────────────────────────────
    ensure_config()
    config = load_config()

    if concurrent is not None:
        config.setdefault("download", {})["concurrent"] = concurrent

    # ── 2. Validate / discover workflow path ────────────────────────────────
    if workflow is None:
        workflow = _pick_workflow(config)
        if workflow is None:
            console.print("[yellow]No workflow selected — exiting.[/yellow]")
            raise typer.Exit(code=0)

    if not workflow.exists():
        console.print(f"[red]✗ Workflow file not found: {workflow}[/red]")
        raise typer.Exit(code=1)
    if not workflow.is_file():
        console.print(f"[red]✗ Path is not a file: {workflow}[/red]")
        raise typer.Exit(code=1)

    # ── 3. Parse ────────────────────────────────────────────────────────────
    console.print(Rule(f"[bold cyan]comfy-resolve[/bold cyan]  [dim]{workflow.name}[/dim]"))

    try:
        assets = parse_workflow(workflow)
    except json.JSONDecodeError:
        raise  # re-raise so the outer handler formats it nicely
    except Exception as exc:
        console.print(f"[red]✗ Error parsing workflow: {exc}[/red]")
        raise typer.Exit(code=1)

    if not assets:
        console.print("[yellow]No model assets found in this workflow.[/yellow]")
        raise typer.Exit(code=0)

    console.print(f"[cyan]Found {len(assets)} asset(s) in workflow[/cyan]\n")

    # ── 4. Resolve ──────────────────────────────────────────────────────────
    resolver = Resolver(config)

    # --no-cache: monkey-patch the cache check so it always returns None
    if no_cache:
        resolver._check_cache = lambda asset: None  # type: ignore[method-assign]

    with console.status("[cyan]Resolving…[/cyan]", spinner="dots") as status:
        def _progress(name: str) -> None:
            short = name if len(name) <= 50 else name[:47] + "…"
            status.update(f"[cyan]Resolving:[/cyan] [dim]{short}[/dim]")
        results = resolver.resolve_all(assets, on_progress=_progress)

    # ── 5. Auto / interactive resolution modes ──────────────────────────────
    if auto_all:
        confirmed_results = results
        unresolved = sum(1 for r in results if r.source is None)
        console.print(
            f"[yellow]--auto-all: accepting all {len(results)} resolution(s)"
            + (f" ({unresolved} unresolved)" if unresolved else "")
            + "[/yellow]"
        )
    elif auto:
        confirmed_results = [r for r in results if r.confidence == "high"]
        skipped = len(results) - len(confirmed_results)
        console.print(
            f"[yellow]--auto: {len(confirmed_results)} high-confidence, "
            f"{skipped} skipped[/yellow]"
        )
    else:
        session = UISession(results, config)
        confirmed_results = session.run()

    if not confirmed_results:
        console.print("[yellow]Nothing to do after review.[/yellow]")
        raise typer.Exit(code=0)

    # ── 6. Build download task list ─────────────────────────────────────────
    downloader = Downloader(config)
    tasks: list[DownloadTask] = []

    for result in confirmed_results:
        if not result.source:
            continue
        asset_type = result.asset_type
        dest = download_dir if download_dir else _safe_get_model_path(asset_type, config)
        tasks.append(DownloadTask(result=result, destination=dest, asset_type=asset_type))

    if not tasks:
        console.print("[yellow]No resolvable assets — nothing to download.[/yellow]")
        raise typer.Exit(code=0)

    # ── 7. Dry-run ──────────────────────────────────────────────────────────
    if dry_run:
        _ui(config).display_dry_run_table(tasks)
        raise typer.Exit(code=0)

    # ── 8. Pre-download review (interactive only) ───────────────────────────
    if not auto and not auto_all:
        tasks = _ui(config).prompt_pre_download_review(tasks)
        if not tasks:
            console.print("[yellow]No assets selected — nothing to download.[/yellow]")
            raise typer.Exit(code=0)

    # ── 9. Disk-space check ─────────────────────────────────────────────────
    download_tasks = [
        t for t in tasks
        if t.result.source and t.result.source.source_type != "local"
    ]

    if download_tasks:
        try:
            enough, required, available = downloader.check_disk_space(download_tasks)
            if not enough:
                console.print(
                    Panel(
                        f"[bold]Required :[/bold]  {_fmt(required)}\n"
                        f"[bold]Available:[/bold]  {_fmt(available)}",
                        title="[red]⚠ Insufficient disk space[/red]",
                        border_style="red",
                    )
                )
                if not typer.confirm("Continue anyway?", default=False):
                    console.print("[yellow]Aborted.[/yellow]")
                    raise typer.Exit(code=0)
            else:
                console.print(
                    f"[dim]Disk space: {_fmt(required)} required, "
                    f"{_fmt(available)} available[/dim]"
                )
        except typer.Exit:
            raise
        except Exception as exc:
            # Non-fatal — warn and proceed
            console.print(f"[yellow]⚠ Could not check disk space: {exc}[/yellow]")

    # ── 10. Download ────────────────────────────────────────────────────────
    actual_count = len(download_tasks)
    local_count  = len(tasks) - actual_count

    if actual_count:
        msg = f"Downloading [bold]{actual_count}[/bold] asset(s)"
        if local_count:
            msg += f"  [dim]({local_count} already local, skipped)[/dim]"
        console.print(f"\n[cyan]{msg}[/cyan]")
        download_results = downloader.download_all(tasks)
    else:
        console.print("[green]✓ All assets are already present locally.[/green]")
        raise typer.Exit(code=0)

    # ── 11. Persist non-local mappings to cache ─────────────────────────────
    for result in confirmed_results:
        if (
            result.source
            and result.source.source_type != "local"
            and result.confidence in ("high", "medium")
        ):
            try:
                save_mapping(result.input_name, result)
            except Exception as exc:
                logger.warning(
                    "Failed to cache mapping for %s: %s", result.input_name, exc
                )

    # ── 12. Summary ─────────────────────────────────────────────────────────
    _ui(config).display_download_summary(download_results)


# ---------------------------------------------------------------------------
# cache-list
# ---------------------------------------------------------------------------

@app.command(name="cache-list")
def cache_list() -> None:
    """List all cached asset→source mappings."""
    mappings = list_mappings()

    if not mappings:
        console.print("[yellow]No cached mappings.[/yellow]")
        return

    table = Table(title="Cached Mappings", box=box.ROUNDED, show_lines=False)
    table.add_column("Asset",            style="cyan",   no_wrap=True)
    table.add_column("Source Type",      style="blue")
    table.add_column("Repository / URL", style="green")
    table.add_column("Confidence",       style="yellow")

    for m in mappings:
        table.add_row(
            m.get("filename", "—"),
            m.get("source", {}).get("type", "—"),
            m.get("source", {}).get("repo", "—"),
            m.get("confidence", "—"),
        )

    console.print(table)
    console.print(f"\n[dim]{len(mappings)} mapping(s) cached[/dim]")


# ---------------------------------------------------------------------------
# cache-clear
# ---------------------------------------------------------------------------

@app.command(name="cache-clear")
def cache_clear(
    name: Optional[str] = typer.Argument(
        None,
        help="Filename of a specific mapping to remove. Omit to clear everything.",
    ),
) -> None:
    """
    Remove cached asset mappings.

    Pass a filename to remove just that entry, or omit it to wipe the entire
    cache (with confirmation).
    """
    if name:
        delete_mapping(name)
        console.print(f"[green]✓ Cleared mapping for: {name}[/green]")
    else:
        count = len(list_mappings())
        if count == 0:
            console.print("[yellow]Cache is already empty.[/yellow]")
            return
        if typer.confirm(f"Clear all {count} cached mapping(s)?", default=False):
            save_mappings({})
            console.print("[green]✓ Cache cleared.[/green]")
        else:
            console.print("[yellow]Cancelled.[/yellow]")


# ---------------------------------------------------------------------------
# config-init
# ---------------------------------------------------------------------------

@app.command(name="config-init")
def config_init() -> None:
    """Run the first-time setup wizard (or reinitialise an existing config)."""
    if config_exists():
        console.print("[yellow]A config file already exists.[/yellow]")
        if not typer.confirm("Overwrite it?", default=False):
            console.print("[yellow]Cancelled.[/yellow]")
            raise typer.Exit(code=0)
    run_first_time_setup()


# ---------------------------------------------------------------------------
# config-show
# ---------------------------------------------------------------------------

@app.command(name="config-show")
def config_show() -> None:
    """Print the active configuration (API keys are redacted)."""
    ensure_config()
    config = load_config()

    # Deep-redact every value under [api_keys]
    if "api_keys" in config:
        config["api_keys"] = {
            k: "***" if v else "" for k, v in config["api_keys"].items()
        }

    console.print(Rule("[bold cyan]Active Configuration[/bold cyan]"))
    console.print(toml.dumps(config))


# ---------------------------------------------------------------------------
# Workflow picker (shown when no path is given on the CLI)
# ---------------------------------------------------------------------------

def _pick_workflow(config: dict[str, Any]) -> Optional[Path]:
    """
    Scan the ComfyUI ``user/default/workflows`` directory for ``.json`` files
    and present them as a numbered pick-list.

    The user may:
    - Enter a number to select a discovered workflow.
    - Type any file path to use a workflow outside the default directory.
    - Press Enter (blank) to abort.

    Returns the selected :class:`Path`, or ``None`` if the user aborts.
    """
    comfyui_dir = config.get("comfyui_dir", "")
    workflows_dir = Path(comfyui_dir) / "user" / "default" / "workflows"

    discovered: list[Path] = []
    if workflows_dir.is_dir():
        discovered = sorted(workflows_dir.glob("*.json"))

    console.print(Rule("[bold cyan]Select a Workflow[/bold cyan]"))

    if discovered:
        table = Table(box=box.SIMPLE, show_header=True, header_style="bold cyan")
        table.add_column("#",        style="dim",    width=4,  justify="right")
        table.add_column("Filename", style="cyan",   no_wrap=True)
        table.add_column("Size",     style="yellow", justify="right")
        table.add_column("Path",     style="dim")

        for i, p in enumerate(discovered, 1):
            try:
                size = _fmt(p.stat().st_size)
            except OSError:
                size = "—"
            table.add_row(str(i), p.name, size, str(p.parent))

        console.print(table)
        console.print(
            f"[dim]Found {len(discovered)} workflow(s) in {workflows_dir}[/dim]\n"
        )
        console.print(
            "[dim]Enter a [bold]number[/bold] to select  •  "
            "a [bold]file path[/bold] to use a different workflow  •  "
            "[bold]blank[/bold] to abort[/dim]"
        )
    else:
        if comfyui_dir:
            console.print(
                f"[yellow]No workflows found in {workflows_dir}[/yellow]\n"
            )
        else:
            console.print("[yellow]ComfyUI directory not configured.[/yellow]\n")
        console.print(
            "[dim]Enter a [bold]file path[/bold] to a workflow JSON, "
            "or leave blank to abort.[/dim]"
        )

    raw = Prompt.ask("\nWorkflow", default="").strip()

    if not raw:
        return None

    # ── Numeric selection ───────────────────────────────────────────────────
    if raw.isdigit():
        idx = int(raw) - 1
        if 0 <= idx < len(discovered):
            chosen = discovered[idx]
            console.print(f"[cyan]Selected: {chosen.name}[/cyan]\n")
            return chosen
        console.print(
            f"[red]✗ Invalid selection '{raw}' — valid range is 1–{len(discovered)}.[/red]"
        )
        return None

    # ── Treat as an explicit file path ──────────────────────────────────────
    path = Path(raw)
    if not path.exists():
        console.print(f"[red]✗ File not found: {path}[/red]")
        return None
    if not path.is_file():
        console.print(f"[red]✗ Not a file: {path}[/red]")
        return None
    return path


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------

def _setup_logging(verbose: bool) -> None:
    """Configure root logger verbosity."""
    level = logging.DEBUG if verbose else logging.WARNING
    fmt   = "%(asctime)s  %(name)s  %(levelname)s  %(message)s" if verbose else "%(message)s"
    logging.basicConfig(level=level, format=fmt, force=True)
    if verbose:
        logger.debug("Verbose logging enabled.")


def _safe_get_model_path(asset_type: str, config: dict[str, Any]) -> Path:
    """
    Return the configured model path for *asset_type*.

    Falls back to ``<comfyui_dir>/models/<asset_type>`` if the type is
    unknown or the config lookup fails, so a bad type never crashes downstream.
    """
    try:
        return get_model_path(asset_type, config)
    except Exception:
        base = Path(config.get("comfyui_dir", ".")) / "models"
        fallback = base / asset_type
        logger.warning(
            "Could not resolve model path for type '%s'; falling back to %s",
            asset_type,
            fallback,
        )
        return fallback


def _ui(config: dict[str, Any]) -> UISession:
    """Return a throw-away UISession (used to call display helpers)."""
    return UISession([], config)


def _fmt(size_bytes: int) -> str:
    """Return a human-readable byte size string, e.g. ``1.45 GB``."""
    size = float(size_bytes)
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if size < 1024.0:
            return f"{size:.2f} {unit}"
        size /= 1024.0
    return f"{size:.2f} PB"


def run_first_time_setup() -> dict[str, Any]:
    """
    Prompt for the ComfyUI installation directory, write the default config,
    and return the freshly-loaded config dict.
    """
    console.print(
        Panel(
            "[bold]Welcome to [cyan]comfy-resolve[/cyan]![/bold]\n\n"
            "Let's configure your ComfyUI installation path.\n"
            "All model sub-directories will be inferred automatically.\n"
            "You can override individual paths in the generated config file.",
            title="First-time setup",
            border_style="cyan",
        )
    )

    comfyui_dir = typer.prompt("ComfyUI installation directory")

    if not comfyui_dir.strip():
        console.print("[red]✗ Path cannot be empty.[/red]")
        raise typer.Exit(code=1)

    dir_path = Path(comfyui_dir)
    if not dir_path.exists():
        console.print(f"[yellow]⚠ Directory does not exist: {dir_path}[/yellow]")
        if not typer.confirm("Create it and continue anyway?", default=False):
            raise typer.Exit(code=1)

    config_path = write_default_config(str(dir_path))
    console.print(f"\n[green]✓ Config written to {config_path}[/green]")
    console.print(
        "[dim]Edit the file to set API keys or override individual model paths.[/dim]"
    )
    return load_config()


def ensure_config() -> None:
    """
    Ensure a config file exists.  If not, run the first-time setup wizard
    interactively before proceeding.
    """
    if not config_exists():
        run_first_time_setup()


# ---------------------------------------------------------------------------
# App callback — global help text
# ---------------------------------------------------------------------------

@app.callback()
def _callback(ctx: typer.Context) -> None:
    """
    [bold cyan]comfy-resolve[/bold cyan] — ComfyUI Workflow Dependency Resolver

    Parses a workflow JSON file, locates every referenced model on HuggingFace,
    Civitai, or your local filesystem, then downloads anything that is missing.
    """
    pass


# ---------------------------------------------------------------------------
# Direct invocation
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    app()