"""alice compare — Send the same prompt to multiple models and compare responses."""

from __future__ import annotations

import difflib
import json
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from typing import Any

import click
from botocore.exceptions import ClientError
from rich.panel import Panel
from rich.table import Table

from alice_cli import personality
from alice_cli.aws.bedrock_context import create_bedrock_context
from alice_cli.bedrock_utils import build_converse_kwargs, parse_converse_response
from alice_cli.console import (
    err_console,
    print_banner,
    print_rule,
    print_success,
    spinner,
)
from alice_cli.constants import get_memory_bucket
from alice_cli.errors import AliceCLIError
from alice_cli.locker import CloudLocker
from alice_cli.models import (
    TokenUsage,
    resolve_model_id,
)
from alice_cli.session_record import build_session_record


@dataclass
class _ModelResult:
    """Result from a single model invocation."""

    alias: str
    resolved_id: str
    text: str
    input_tokens: int
    output_tokens: int
    error: str | None = None


def _invoke_model(
    client: Any,
    resolved_id: str,
    alias: str,
    prompt_text: str,
    system_prompt: str | None,
) -> _ModelResult:
    """Invoke a single model via the Converse API. Returns a _ModelResult."""
    try:
        response = client.converse(
            **build_converse_kwargs(
                resolved_id,
                [{"role": "user", "content": [{"text": prompt_text}]}],
                system_prompt=system_prompt,
            )
        )
        text, tokens = parse_converse_response(response)
        return _ModelResult(
            alias=alias,
            resolved_id=resolved_id,
            text=text,
            input_tokens=tokens.input,
            output_tokens=tokens.output,
        )
    except ClientError as exc:
        return _ModelResult(
            alias=alias,
            resolved_id=resolved_id,
            text="",
            input_tokens=0,
            output_tokens=0,
            error=str(exc),
        )


def _display_table(results: list[_ModelResult]) -> None:
    """Display model responses in a Rich table with one column per model."""
    table = Table(title="Model Comparison", show_lines=True)
    for r in results:
        table.add_column(r.alias, style="bold")

    cells: list[str] = []
    for r in results:
        if r.error:
            cells.append(f"[red]Error: {r.error}[/red]")
        else:
            cells.append(r.text or "[dim](empty response)[/dim]")

    table.add_row(*cells)
    err_console.print(table)


def _display_diff(results: list[_ModelResult]) -> None:
    """Display a line-by-line diff between the first two model responses."""
    if len(results) < 2:
        err_console.print("[yellow]Need at least 2 models for diff mode.[/yellow]")
        return

    a_lines = (results[0].text or "").splitlines(keepends=True)
    b_lines = (results[1].text or "").splitlines(keepends=True)

    diff = difflib.unified_diff(
        a_lines,
        b_lines,
        fromfile=results[0].alias,
        tofile=results[1].alias,
    )
    diff_text = "".join(diff)
    if diff_text:
        err_console.print(Panel(diff_text, title="Diff", border_style="cyan"))
    else:
        err_console.print("[dim]No differences found between the two responses.[/dim]")


@click.command("compare")
@click.option(
    "--models",
    required=True,
    help="Comma-separated list of model aliases or IDs to compare (at least 2).",
)
@click.option(
    "--diff",
    "diff_mode",
    is_flag=True,
    help="Show a line-by-line diff between the first two models instead of a table",
)
@click.option(
    "--system", "system_prompt", default=None, help="System prompt for all model invocations"
)
@click.option(
    "-o",
    "--output",
    "output_path",
    default=None,
    type=click.Path(),
    help="Write comparison results to a JSON file",
)
@click.argument("prompt", nargs=-1, type=click.UNPROCESSED, required=True)
@click.pass_context
def compare_cmd(
    ctx: click.Context,
    models: str,
    diff_mode: bool,
    system_prompt: str | None,
    output_path: str | None,
    prompt: tuple[str, ...],
) -> None:
    """Compare responses from multiple models side-by-side.

    \b
    Examples:
        alice compare --models sonnet,haiku "Explain quantum computing"
        alice compare --models sonnet,opus --diff "What is recursion?"
        alice compare --models sonnet,haiku,nova-pro --output results.json "Hello"
        alice compare --models sonnet,opus --system "Be concise" "Explain AI"
    """
    config = ctx.obj["config"]

    prompt_text = " ".join(prompt).strip()
    if not prompt_text:
        raise AliceCLIError("Please provide a prompt after the options.")

    model_list = [m.strip() for m in models.split(",") if m.strip()]
    if len(model_list) < 2:
        raise AliceCLIError("Please provide at least 2 models separated by commas.")

    print_banner(personality.BANNER, config.quiet)

    # ── Authenticate and build Bedrock client ─────────────────────────
    context = create_bedrock_context(config)
    username = context.username
    client = context.client

    # ── Resolve model aliases ─────────────────────────────────────────
    resolved_models: list[tuple[str, str]] = []  # (alias, resolved_id)
    for alias in model_list:
        resolved_id = resolve_model_id(alias, context.inference_profile_arns)
        resolved_models.append((alias, resolved_id))

    # ── Invoke models concurrently ────────────────────────────────────
    results: list[_ModelResult] = []

    with (
        spinner("Comparing models…", config.quiet),
        ThreadPoolExecutor(max_workers=len(resolved_models)) as executor,
    ):
        futures = {
            executor.submit(
                _invoke_model, client, resolved_id, alias, prompt_text, system_prompt
            ): alias
            for alias, resolved_id in resolved_models
        }
        for future in as_completed(futures):
            results.append(future.result())

    # Preserve the original model order
    order = {alias: i for i, (alias, _) in enumerate(resolved_models)}
    results.sort(key=lambda r: order.get(r.alias, 0))

    # ── Report per-model errors ───────────────────────────────────────
    for r in results:
        if r.error:
            err_console.print(f"  [red]✗[/red] {r.alias}: {r.error}")

    successful = [r for r in results if r.error is None]
    if not successful:
        raise AliceCLIError("All model invocations failed.")

    # ── Display results ───────────────────────────────────────────────
    print_rule(config.quiet)
    if diff_mode:
        _display_diff(results)
    else:
        _display_table(results)
    print_rule(config.quiet)

    # ── Write JSON output ─────────────────────────────────────────────
    if output_path:
        output_data = {
            "prompt": prompt_text,
            "system": system_prompt,
            "models": [
                {
                    "alias": r.alias,
                    "model_id": r.resolved_id,
                    "response": r.text if not r.error else None,
                    "error": r.error,
                    "tokens": {
                        "input": r.input_tokens,
                        "output": r.output_tokens,
                        "total": r.input_tokens + r.output_tokens,
                    },
                }
                for r in results
            ],
        }
        try:
            with open(output_path, "w", encoding="utf-8") as f:
                json.dump(output_data, f, indent=2, ensure_ascii=False)
            print_success(f"Results written to {output_path}", config.quiet)
        except Exception as exc:
            raise AliceCLIError(f"Could not write output file: {exc}") from exc

    # ── Persist to Cloud Locker ───────────────────────────────────────
    total_input = sum(r.input_tokens for r in results)
    total_output = sum(r.output_tokens for r in results)

    metadata: dict[str, Any] = {
        "command": "compare",
        "models": model_list,
        "responses": [
            {
                "model": r.alias,
                "text": r.text if not r.error else None,
                "error": r.error,
                "tokens": {
                    "input": r.input_tokens,
                    "output": r.output_tokens,
                    "total": r.input_tokens + r.output_tokens,
                },
            }
            for r in results
        ],
    }
    if system_prompt:
        metadata["system"] = system_prompt

    locker = CloudLocker(
        config=config,
        username=username,
        bucket=get_memory_bucket(config.namespace, config.environment),
    )
    record = build_session_record(
        record_type="compare",
        model=",".join(alias for alias, _ in resolved_models),
        prompt=prompt_text,
        response=None,
        tokens=TokenUsage(
            input=total_input,
            output=total_output,
            total=total_input + total_output,
        ),
        metadata=metadata,
    )
    try:
        locker.persist(record, "compare")
        print_success("Session saved to Cloud Locker", config.quiet)
    except Exception as exc:
        err_console.print(f"[yellow]⚠ Could not save session: {exc}[/yellow]")
