"""
Tool discovery and shared execution utilities.

Tools are organized in category subfolders (folding/, utilities/, etc.).
Each tool file exports METADATA dict and register() function.
"""

import importlib
from pathlib import Path
from typing import Iterator, Optional
import typer
from rich.console import Console
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn

from amina_cli.commands.tools.display import render_tool_output

console = Console()


def discover_tools() -> Iterator[tuple[str, dict, callable]]:
    """
    Discover all tools in category subfolders.

    Yields:
        (module_name, METADATA dict, register function)
    """
    tools_dir = Path(__file__).parent

    for category_dir in sorted(tools_dir.iterdir()):
        if not category_dir.is_dir() or category_dir.name.startswith("_"):
            continue

        for tool_file in sorted(category_dir.glob("*.py")):
            if tool_file.name.startswith("_"):
                continue

            module_name = f"amina_cli.commands.tools.{category_dir.name}.{tool_file.stem}"
            try:
                module = importlib.import_module(module_name)
                if hasattr(module, "METADATA") and hasattr(module, "register"):
                    yield module_name, module.METADATA, module.register
            except ImportError as e:
                console.print(f"[dim]Warning: Could not load {module_name}: {e}[/dim]")


def get_all_tools() -> list[dict]:
    """Get METADATA for all discovered tools."""
    return [metadata for _, metadata, _ in discover_tools()]


def get_tool(name: str) -> Optional[dict]:
    """Get METADATA for a specific tool by name."""
    for _, metadata, _ in discover_tools():
        if metadata["name"] == name:
            return metadata
    return None


def register_all_tools(app: typer.Typer):
    """Register all discovered tool commands with the app."""
    for _, metadata, register_fn in discover_tools():
        register_fn(app)


def run_tool_with_progress(
    tool_name: str,
    params: dict,
    output_dir: Optional[Path] = None,
    timeout: float = 36000.0,  # 10 hours default
    background: bool = False,
):
    """
    Execute a tool with progress display and result downloading.

    This is the shared execution function used by all tool commands.

    Args:
        tool_name: Name of the tool to run
        params: Tool-specific parameters
        output_dir: Directory to save output files (required unless background=True)
        timeout: Request timeout in seconds (default: 10 hours)
        background: If True, submit job and return immediately without waiting
    """
    from amina_cli.client import (
        run_tool_sync,
        submit_tool_sync,
        AuthenticationError,
        ToolExecutionError,
        InsufficientCreditsError,
        JobLimitError,
    )
    from amina_cli.registry import ToolNotFoundError
    from amina_cli.storage import download_results, StorageError

    # Handle background submission
    if background:
        try:
            from amina_cli.auth import save_job

            tool_info = get_tool(tool_name)
            display_name = tool_info.get("display_name", tool_name) if tool_info else tool_name

            with Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                console=console,
            ) as progress:
                progress.add_task(f"Submitting {display_name}...", total=None)
                result = submit_tool_sync(tool_name, params)

            # Save to local job history
            save_job(result)

            # Handle queued vs submitted status
            if result.get("status") == "queued":
                queue_position = result.get("queue_position", 0)
                console.print(f"\n[yellow]⏳[/yellow] Job queued: [bold]{result['job_id']}[/bold]")
                console.print(f"  Tool: {display_name}")
                console.print(f"  Queue position: {queue_position}")
                console.print("[dim]Will start automatically when capacity available[/dim]")
            else:
                console.print(f"\n[green]\u2713[/green] Job submitted: [bold]{result['job_id']}[/bold]")
                console.print(f"  Tool: {display_name}")
                console.print(f"  Call ID: {result.get('call_id', 'N/A')}")
            console.print(f"\n[dim]Check status:[/dim]  amina jobs status {result['job_id'][:8]}")
            console.print(f"[dim]Wait for completion:[/dim]  amina jobs wait {result['job_id'][:8]}")
            console.print(f"[dim]Download results:[/dim]  amina jobs download {result['job_id'][:8]} -o <output_dir>")
            return

        except AuthenticationError as e:
            console.print(
                Panel(
                    f"[red]Authentication failed[/red]\n\n{e}",
                    title="Error",
                    border_style="red",
                )
            )
            raise typer.Exit(1)

        except InsufficientCreditsError as e:
            console.print(
                Panel(
                    f"[red]Insufficient Credits[/red]\n\n"
                    f"Required: {e.required} credits\n"
                    f"Remaining: {e.remaining:.1f} credits\n\n"
                    f"Upgrade at: {e.upgrade_url}",
                    title="Payment Required",
                    border_style="yellow",
                )
            )
            raise typer.Exit(1)

        except JobLimitError as e:
            console.print(
                Panel(
                    f"[red]Job Limit Reached[/red]\n\n"
                    f"{e}\n\n"
                    f"Current: {e.current} / Limit: {e.limit}\n\n"
                    f"Use [cyan]amina jobs list[/cyan] to see running jobs.",
                    title="Rate Limited",
                    border_style="yellow",
                )
            )
            raise typer.Exit(1)

        except ToolExecutionError as e:
            console.print(
                Panel(
                    f"[red]Job submission failed[/red]\n\n{e}",
                    title="Error",
                    border_style="red",
                )
            )
            raise typer.Exit(1)

        except ToolNotFoundError as e:
            console.print(
                Panel(
                    f"[red]Tool not found[/red]\n\n{e}",
                    title="Error",
                    border_style="red",
                )
            )
            raise typer.Exit(1)

    # Normal blocking execution
    try:
        tool_info = get_tool(tool_name)
        display_name = tool_info.get("display_name", tool_name) if tool_info else tool_name

        # Track queue position for progress updates
        queue_position_holder = {"position": 0, "is_queued": False}

        def on_queue_position_update(position: int):
            queue_position_holder["position"] = position
            queue_position_holder["is_queued"] = True

        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            TimeElapsedColumn(),
            console=console,
        ) as progress:
            task = progress.add_task(f"Running {display_name}...", total=None)

            # Queue position callback for progress updates
            last_position = -1

            def position_callback(position: int):
                nonlocal last_position
                queue_position_holder["position"] = position
                queue_position_holder["is_queued"] = True
                if position != last_position:
                    progress.update(
                        task,
                        description=f"[yellow]Queued[/yellow] (position {position}) - {display_name}",
                    )
                    last_position = position

            # Execute tool (blocking)
            result = run_tool_sync(
                tool_name,
                params,
                timeout=timeout,
                on_queue_position_update=position_callback,
            )

            # Update progress when job completes
            if queue_position_holder["is_queued"]:
                progress.update(task, description=f"Running {display_name}...")

            progress.update(task, description="Downloading results...")

            # Download output files
            downloaded = []
            try:
                downloaded = download_results(result, output_dir)
            except StorageError as e:
                # If signed_urls are available, show them as fallback
                signed_urls = result.get("signed_urls", {})
                if signed_urls:
                    console.print("\n[yellow]Warning:[/yellow] Could not download files automatically.")
                    console.print("\n[bold]Files available at:[/bold]")
                    for file_type, url in signed_urls.items():
                        if url:
                            console.print(f"  {file_type}: {url}")
                else:
                    console.print(f"\n[yellow]Warning:[/yellow] Download failed: {e}")

        _display_result(result, downloaded, output_dir, params, tool_info)

    except AuthenticationError as e:
        console.print(
            Panel(
                f"[red]Authentication failed[/red]\n\n{e}",
                title="Error",
                border_style="red",
            )
        )
        raise typer.Exit(1)

    except InsufficientCreditsError as e:
        console.print(
            Panel(
                f"[red]Insufficient Credits[/red]\n\n"
                f"Required: {e.required} credits\n"
                f"Remaining: {e.remaining:.1f} credits\n\n"
                f"Upgrade at: {e.upgrade_url}",
                title="Payment Required",
                border_style="yellow",
            )
        )
        raise typer.Exit(1)

    except JobLimitError as e:
        console.print(
            Panel(
                f"[red]Job Limit Reached[/red]\n\n"
                f"{e}\n\n"
                f"Current: {e.current} / Limit: {e.limit}\n\n"
                f"Use [cyan]amina jobs list[/cyan] to see running jobs.",
                title="Rate Limited",
                border_style="yellow",
            )
        )
        raise typer.Exit(1)

    except ToolExecutionError as e:
        console.print(
            Panel(
                f"[red]Tool execution failed[/red]\n\n{e}",
                title="Error",
                border_style="red",
            )
        )
        raise typer.Exit(1)

    except ToolNotFoundError as e:
        console.print(
            Panel(
                f"[red]Tool not found[/red]\n\n{e}",
                title="Error",
                border_style="red",
            )
        )
        raise typer.Exit(1)


def _display_result(
    result: dict,
    downloaded: list[Path],
    output_dir: Path,
    input_params: dict | None = None,
    metadata: dict | None = None,
):
    """Display tool execution result using the generic renderer."""
    status = result.get("status", "unknown")

    if status == "success":
        # Build success message
        exec_time = result.get("execution_time_seconds", 0)
        time_str = f" ({exec_time:.1f}s)" if exec_time else ""

        if downloaded:
            console.print(f"\n[green]\u2713[/green] Completed{time_str}")
            console.print(f"[green]\u2713[/green] Downloaded {len(downloaded)} file(s) to {output_dir}/")
            for path in downloaded:
                console.print(f"  - {path.name}")
        else:
            console.print(f"\n[green]\u2713[/green] Completed{time_str}")

        # Show credits charged (from gateway)
        credits_charged = result.get("credits_charged", 0)
        if credits_charged:
            console.print(f"[dim]Credits charged: {credits_charged:.2f}[/dim]")

        # Show warnings if present (check both top-level and in data)
        data = result.get("data", {}) or {}
        warnings = result.get("warnings") or data.get("warnings", [])
        if warnings:
            console.print(f"\n[yellow]⚠ {len(warnings)} warning(s):[/yellow]")
            for warning in warnings:
                console.print(f"  [yellow]•[/yellow] {warning}")

        # Render tool output using declarative display config or auto-discovery
        render_tool_output(result, metadata, input_params)

        # Show input parameters (exclude large content fields)
        if input_params:
            # Keys to exclude from display (large or internal)
            exclude_keys = {
                "pdb_content",
                "sequence",
                "fasta_content",
                "model_pdb_content",
                "reference_pdb_content",  # DockQ large fields
                "mobile_pdb_content",  # Simple RMSD / US-align large fields
                "target_pdb_content",  # US-align large fields
                "yaml_content",  # Boltz-2 large field
            }
            display_params = {k: v for k, v in input_params.items() if k not in exclude_keys and v is not None}
            if display_params:
                console.print("\n[bold]Input parameters:[/bold]")
                for key, value in display_params.items():
                    console.print(f"  {key}: {value}")

    else:
        # Check error_details first (from gateway), then error, then message
        error_msg = result.get("error_details") or result.get("error") or result.get("message", "Unknown error")

        # Show credits charged even for failures (20% partial charge)
        credits_charged = result.get("credits_charged", 0)
        credits_info = f"\nCredits charged: {credits_charged:.2f}" if credits_charged else ""

        console.print(
            Panel(
                f"[red]Tool execution failed[/red]\n\nError: {error_msg}{credits_info}",
                title="Error",
                border_style="red",
            )
        )
