"""Bernstein cost — spend visibility across all recorded metrics."""

from __future__ import annotations

import contextlib
import json
from collections import defaultdict
from pathlib import Path
from typing import Any

import click
from rich.console import Console
from rich.panel import Panel

console = Console()

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _ascii_bar(value: float, max_value: float, width: int = 30) -> str:
    """Return a block-character bar proportional to value/max_value."""
    if max_value <= 0 or value <= 0:
        return "░" * width
    filled = max(1, round((value / max_value) * width))
    filled = min(filled, width)
    return "█" * filled + "░" * (width - filled)


def _count_task_status(task_records: list[dict[str, Any]]) -> tuple[int, int]:
    """Return (done_count, failed_count) deduped by task_id."""
    seen: dict[str, dict[str, Any]] = {}
    for rec in task_records:
        tid = rec.get("task_id", "")
        if tid:
            seen[tid] = rec
    done = sum(1 for r in seen.values() if r.get("status") == "done")
    failed = sum(1 for r in seen.values() if r.get("status") == "failed")
    # If status not recorded, fall back to presence of cost as "done"
    if done == 0 and failed == 0 and seen:
        done = sum(1 for r in seen.values() if float(r.get("cost_usd", 0) or 0) > 0)
    return done, failed


def _render_savings_comparison(
    cons: Console,
    actual_cost: float,
    savings_vs_opus: float,
) -> None:
    """Print an ASCII bar chart comparing Bernstein vs all-Opus baseline."""
    single_agent_cost = actual_cost + savings_vs_opus
    if single_agent_cost <= 0:
        return

    savings_pct = (savings_vs_opus / single_agent_cost) * 100

    bar_width = 34
    single_bar = _ascii_bar(single_agent_cost, single_agent_cost, bar_width)
    bernstein_bar = _ascii_bar(actual_cost, single_agent_cost, bar_width)

    cons.print()
    cons.print("[bold]Cost Comparison[/bold]  (Bernstein vs all-Opus baseline)")
    cons.print(f"  Single agent  [red]{single_bar}[/red]  [dim]${single_agent_cost:.4f}[/dim]")
    cons.print(f"  Bernstein     [green]{bernstein_bar}[/green]  [bold green]${actual_cost:.4f}[/bold green]")
    if savings_pct > 0:
        cons.print(
            f"\n  [bold green]You saved ${savings_vs_opus:.4f} "
            f"({savings_pct:.0f}%) by using Bernstein's model cascade[/bold green]"
        )


def _render_shareable_summary(
    cons: Console,
    actual_cost: float,
    savings_vs_opus: float,
    tasks_done: int,
    tasks_failed: int,
    total_duration_s: float,
) -> None:
    """Print a copy-pasteable markdown run summary."""
    single_agent_cost = actual_cost + savings_vs_opus
    savings_pct = (savings_vs_opus / single_agent_cost) * 100 if single_agent_cost > 0 else 0

    mins = int(total_duration_s // 60)
    secs = int(total_duration_s % 60)
    time_str = f"{mins}m {secs}s" if mins > 0 else f"{secs}s"

    lines: list[str] = [
        "🎼 Bernstein run summary",
        f"   Tasks: {tasks_done} completed" + (f", {tasks_failed} failed" if tasks_failed else ""),
    ]
    if total_duration_s > 0:
        lines.append(f"   Time:  {time_str}")
    if single_agent_cost > actual_cost:
        lines.append(f"   Cost:  ${actual_cost:.2f} (vs ~${single_agent_cost:.2f} single agent)")
        lines.append(f"   Saved: ${savings_vs_opus:.2f} ({savings_pct:.0f}%)")
    else:
        lines.append(f"   Cost:  ${actual_cost:.2f}")

    cons.print()
    cons.print(
        Panel(
            "\n".join(lines),
            title="[bold]Shareable summary[/bold]",
            border_style="dim",
            expand=False,
        )
    )


# ---------------------------------------------------------------------------
# Data loading
# ---------------------------------------------------------------------------


def _load_tasks_jsonl(metrics_dir: Path) -> list[dict[str, Any]]:
    p = metrics_dir / "tasks.jsonl"
    if not p.exists():
        return []
    records: list[dict[str, Any]] = []
    for line in p.read_text().splitlines():
        line = line.strip()
        if line:
            with contextlib.suppress(json.JSONDecodeError):
                records.append(json.loads(line))
    return records


def _load_api_usage_jsonl(metrics_dir: Path) -> list[dict[str, Any]]:
    records: list[dict[str, Any]] = []
    for p in sorted(metrics_dir.glob("api_usage_*.jsonl")):
        for line in p.read_text().splitlines():
            line = line.strip()
            if line:
                with contextlib.suppress(json.JSONDecodeError):
                    records.append(json.loads(line))
    return records


# ---------------------------------------------------------------------------
# Aggregation
# ---------------------------------------------------------------------------


def _aggregate_fast_path_savings(task_records: list[dict[str, Any]]) -> dict[str, Any]:
    """Aggregate fast-path savings from task records.

    Returns dict with tasks_bypassed, estimated_savings_usd, and action breakdown.
    """
    bypassed = 0
    savings = 0.0
    actions: dict[str, int] = defaultdict(int)
    for rec in task_records:
        if rec.get("model") == "fast-path":
            bypassed += 1
            savings += float(rec.get("estimated_savings_usd", 0.0) or 0.0)
            action = rec.get("fast_path_action", "unknown")
            actions[action] += 1
    return {
        "tasks_bypassed": bypassed,
        "estimated_savings_usd": savings,
        "actions": dict(actions),
    }


def _aggregate(
    task_records: list[dict[str, Any]],
    api_records: list[dict[str, Any]],
) -> dict[str, dict[str, Any]]:
    """Return per-model aggregated stats.

    Keys: model name (or "unknown")
    Values: dict with tasks, tokens_in, tokens_out, cost_usd, duration_total, duration_count
    """
    rows: dict[str, dict[str, Any]] = defaultdict(
        lambda: {
            "tasks": 0,
            "tokens_in": 0,
            "tokens_out": 0,
            "cost_usd": 0.0,
            "duration_total": 0.0,
            "duration_count": 0,
        }
    )

    # Deduplicate task records by task_id, keeping the last entry
    seen: dict[str, dict[str, Any]] = {}
    for rec in task_records:
        tid = rec.get("task_id", "")
        if tid:
            seen[tid] = rec
        else:
            # No task_id — treat as anonymous
            model = rec.get("model") or "unknown"
            row = rows[model]
            row["tasks"] += 1
            row["tokens_in"] += int(rec.get("tokens_prompt", 0) or 0)
            row["tokens_out"] += int(rec.get("tokens_completion", 0) or 0)
            row["cost_usd"] += float(rec.get("cost_usd", 0.0) or 0.0)
            dur = float(rec.get("duration_seconds", 0.0) or 0.0)
            if dur > 0:
                row["duration_total"] += dur
                row["duration_count"] += 1

    for rec in seen.values():
        model = rec.get("model") or "unknown"
        row = rows[model]
        row["tasks"] += 1
        row["tokens_in"] += int(rec.get("tokens_prompt", 0) or 0)
        row["tokens_out"] += int(rec.get("tokens_completion", 0) or 0)
        row["cost_usd"] += float(rec.get("cost_usd", 0.0) or 0.0)
        dur = float(rec.get("duration_seconds", 0.0) or 0.0)
        if dur > 0:
            row["duration_total"] += dur
            row["duration_count"] += 1

    # api_usage records have labels with provider/model but no token breakdown
    for rec in api_records:
        labels = rec.get("labels", {})
        model = labels.get("model") or "unknown"
        if model not in rows:
            rows[model]  # ensure key exists (defaultdict)

    return dict(rows)


@click.command("estimate")
@click.argument("goal")
@click.option("--role", default="backend", help="Agent role for the task.")
@click.option("--scope", type=click.Choice(["small", "medium", "large"]), default="medium", help="Task scope.")
@click.option("--complexity", type=click.Choice(["low", "medium", "high"]), default="medium", help="Task complexity.")
@click.option(
    "--metrics-dir",
    default=".sdd/metrics",
    show_default=True,
    help="Directory containing historical metrics.",
)
def estimate_cmd(goal: str, role: str, scope: str, complexity: str, metrics_dir: str) -> None:
    """Predict the cost of a task before running it.

    \b
      bernstein estimate "Fix all typos in src/" --scope small
      bernstein estimate "Implement RAG system" --scope large --complexity high
    """
    from bernstein.core.cost import predict_task_cost
    from bernstein.core.models import Complexity, Scope, Task

    # Create a dummy task for prediction
    task = Task(
        id="estimate",
        title=goal[:100],
        description=goal,
        role=role,
        scope=Scope(scope),
        complexity=Complexity(complexity),
    )

    est_cost = predict_task_cost(task, metrics_dir=Path(metrics_dir))

    console.print(
        Panel(
            f"[bold]Cost Prediction[/bold]\n\n"
            f"Goal:       [cyan]{goal}[/cyan]\n"
            f"Role:       {role}\n"
            f"Scope:      {scope}\n"
            f"Complexity: {complexity}\n\n"
            f"Estimated total: [bold green]${est_cost:.4f}[/bold green] (±20%)\n\n"
            f"[dim]Note: Predictions use historical data when available and assume\n"
            f"average token consumption for the given scope/complexity.[/dim]",
            border_style="green",
            expand=False,
        )
    )


# ---------------------------------------------------------------------------
# Command
# ---------------------------------------------------------------------------


@click.command("cost")
@click.option(
    "--metrics-dir",
    default=".sdd/metrics",
    show_default=True,
    help="Directory containing metrics JSONL files.",
)
@click.option("--json", "as_json", is_flag=True, default=False, help="Output raw JSON.")
@click.option("--share", is_flag=True, default=False, help="Print only the shareable summary snippet.")
def cost_cmd(metrics_dir: str, as_json: bool, share: bool) -> None:
    """Show agent spend: cost, tokens, and duration per model."""
    mdir = Path(metrics_dir)
    if not mdir.exists():
        if as_json:
            click.echo(json.dumps({"error": f"Metrics directory not found: {mdir}"}))
        else:
            console.print(f"[red]Metrics directory not found:[/red] {mdir}")
        raise SystemExit(1)

    task_records = _load_tasks_jsonl(mdir)
    api_records = _load_api_usage_jsonl(mdir)

    rows = _aggregate(task_records, api_records)

    if not rows:
        if as_json:
            click.echo(json.dumps({"rows": [], "totals": {}}))
        else:
            console.print("[dim]No metrics data found.[/dim]")
        return

    # Sort by cost descending, then by task count
    sorted_models = sorted(rows.items(), key=lambda kv: (-kv[1]["cost_usd"], -kv[1]["tasks"]))

    totals: dict[str, Any] = {
        "tasks": 0,
        "tokens_in": 0,
        "tokens_out": 0,
        "cost_usd": 0.0,
        "avg_duration_s": None,
    }
    total_dur = 0.0
    total_dur_count = 0
    for _, v in sorted_models:
        totals["tasks"] += v["tasks"]
        totals["tokens_in"] += v["tokens_in"]
        totals["tokens_out"] += v["tokens_out"]
        totals["cost_usd"] += v["cost_usd"]
        total_dur += v["duration_total"]
        total_dur_count += v["duration_count"]
    if total_dur_count > 0:
        totals["avg_duration_s"] = round(total_dur / total_dur_count, 1)

    fast_path_savings = _aggregate_fast_path_savings(task_records)

    from bernstein.core.cost import (
        compute_daily_cost,
        compute_savings_vs_opus,
        project_monthly_cost,
    )

    savings_vs_opus = compute_savings_vs_opus(task_records)
    daily_costs = compute_daily_cost(task_records, days=7)
    projected_monthly = project_monthly_cost(task_records, window_days=7)

    tasks_done, tasks_failed = _count_task_status(task_records)

    if as_json:
        output = {
            "rows": [
                {
                    "model": model,
                    "tasks": v["tasks"],
                    "tokens_in": v["tokens_in"],
                    "tokens_out": v["tokens_out"],
                    "cost_usd": round(v["cost_usd"], 6),
                    "cost_per_task": round(v["cost_usd"] / v["tasks"], 6) if v["tasks"] > 0 else 0,
                    "avg_duration_s": (
                        round(v["duration_total"] / v["duration_count"], 1) if v["duration_count"] > 0 else None
                    ),
                }
                for model, v in sorted_models
            ],
            "totals": totals,
            "fast_path": fast_path_savings,
            "savings_vs_opus_usd": round(savings_vs_opus, 6),
            "daily_costs": daily_costs,
            "projected_monthly_usd": round(projected_monthly, 4),
            "tasks_done": tasks_done,
            "tasks_failed": tasks_failed,
        }
        click.echo(json.dumps(output, indent=2))
        return

    # --share: print only the shareable snippet and exit
    if share:
        _render_shareable_summary(
            console,
            actual_cost=totals["cost_usd"],
            savings_vs_opus=savings_vs_opus,
            tasks_done=tasks_done,
            tasks_failed=tasks_failed,
            total_duration_s=total_dur,
        )
        return

    from rich.table import Table

    table = Table(title="Bernstein Cost Report", header_style="bold cyan", show_lines=False)
    table.add_column("Model", min_width=20)
    table.add_column("Tasks", justify="right", min_width=6)
    table.add_column("Tokens In", justify="right", min_width=10)
    table.add_column("Tokens Out", justify="right", min_width=10)
    table.add_column("Cost USD", justify="right", min_width=10)
    table.add_column("Cost/Task", justify="right", min_width=10)
    table.add_column("Avg Duration", justify="right", min_width=12)

    for model, v in sorted_models:
        avg_dur = f"{v['duration_total'] / v['duration_count']:.1f}s" if v["duration_count"] > 0 else "—"
        cost_str = f"${v['cost_usd']:.4f}" if v["cost_usd"] > 0 else "$0.0000"
        cost_per_task = f"${v['cost_usd'] / v['tasks']:.4f}" if v["tasks"] > 0 else "—"
        table.add_row(
            model,
            str(v["tasks"]),
            f"{v['tokens_in']:,}",
            f"{v['tokens_out']:,}",
            cost_str,
            cost_per_task,
            avg_dur,
        )

    # Totals row
    avg_total = f"{total_dur / total_dur_count:.1f}s" if total_dur_count > 0 else "—"
    total_cost_per_task = f"${totals['cost_usd'] / totals['tasks']:.4f}" if totals["tasks"] > 0 else "—"
    table.add_section()
    table.add_row(
        "[bold]TOTAL[/bold]",
        f"[bold]{totals['tasks']}[/bold]",
        f"[bold]{totals['tokens_in']:,}[/bold]",
        f"[bold]{totals['tokens_out']:,}[/bold]",
        f"[bold green]${totals['cost_usd']:.4f}[/bold green]",
        f"[bold]{total_cost_per_task}[/bold]",
        f"[bold]{avg_total}[/bold]",
    )

    console.print(table)

    # Fast-path savings summary
    if fast_path_savings["tasks_bypassed"] > 0:
        bp = fast_path_savings["tasks_bypassed"]
        sv = fast_path_savings["estimated_savings_usd"]
        actions = fast_path_savings["actions"]
        action_parts = [f"{v} {k}" for k, v in sorted(actions.items(), key=lambda x: -x[1])]
        console.print(
            f"\n[bold green]Fast-path:[/bold green] Saved ~${sv:.2f} by "
            f"bypassing LLM for {bp} task(s) ({', '.join(action_parts)})"
        )

    # ASCII bar chart: Bernstein vs single-agent baseline
    _render_savings_comparison(console, totals["cost_usd"], savings_vs_opus)

    # Projected monthly cost
    if projected_monthly > 0:
        console.print(f"\n[dim]Projected monthly cost (30d):[/dim] ${projected_monthly:.2f}")

    # Shareable run summary
    _render_shareable_summary(
        console,
        actual_cost=totals["cost_usd"],
        savings_vs_opus=savings_vs_opus,
        tasks_done=tasks_done,
        tasks_failed=tasks_failed,
        total_duration_s=total_dur,
    )
