"""Multi-model comparison plot: funds curves + cost vs budget.

Run from the repo root:
    uv run python scripts/plot_multi_model.py [--seed 1] [--config hard] [--budget 30]
"""
import argparse
import json
import sqlite3
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from datetime import datetime

ROOT = Path(__file__).parent.parent

# ---------------------------------------------------------------------------
# Default runs — edit or pass via CLI
# ---------------------------------------------------------------------------

DEFAULT_RUNS = [
    {"label": "gemini-flash",  "model_slug": "openrouter_google_gemini-3-flash-preview", "color": "#4fc3f7"},
    {"label": "minimax-m2.5",  "model_slug": "openrouter_minimax_minimax-m2.5",          "color": "#f39c12"},
    {"label": "kimi-k2.5",    "model_slug": "openrouter_moonshotai_kimi-k2.5",           "color": "#2ecc71"},
]

INITIAL_FUNDS_CENTS = 15_000_000  # $150K (default; presets may override)


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--seed", type=int, default=1)
    p.add_argument("--config", default="hard")
    p.add_argument("--budget", type=float, default=30.0)
    p.add_argument("--out", default=None, help="Output PNG path (default: plots/funds_curves.png)")
    return p.parse_args()


def load_funds_curve(db_path: Path):
    """Reconstruct running balance from ledger entries."""
    con = sqlite3.connect(str(db_path))
    rows = con.execute(
        "SELECT occurred_at, amount_cents FROM ledger_entries ORDER BY occurred_at ASC"
    ).fetchall()
    con.close()

    if not rows:
        return [], []

    times, balances = [], []
    running = INITIAL_FUNDS_CENTS
    # Prepend the sim start (day before first ledger event, pegged to Jan 1)
    start = datetime.fromisoformat(rows[0][0]).replace(
        month=1, day=1, hour=9, minute=0, second=0, microsecond=0
    )
    times.append(start)
    balances.append(INITIAL_FUNDS_CENTS / 100)

    for occurred_at, amount_cents in rows:
        running += int(amount_cents)
        times.append(datetime.fromisoformat(occurred_at))
        balances.append(running / 100)

    return times, balances


def load_meta(result_path: Path):
    with open(result_path) as f:
        d = json.load(f)
    return {
        "turns": d.get("turns_completed", 0),
        "terminal_reason": d.get("terminal_reason", "unknown"),
        "cost_usd": d.get("total_cost_usd", 0.0),
        "horizon_years": d.get("horizon_years", 3),
    }


def load_run_data(runs, seed):
    run_data = []
    for run in runs:
        slug = run["model_slug"]
        db_path = ROOT / "db" / f"{seed}_{slug}.db"
        result_path = ROOT / "results" / f"yc_bench_result_{seed}_{slug}.json"

        if not db_path.exists() or not result_path.exists():
            print(f"  Skipping {run['label']}: missing {db_path.name} or {result_path.name}")
            continue

        times, balances = load_funds_curve(db_path)
        meta = load_meta(result_path)
        run_data.append({**run, "times": times, "balances": balances, **meta})
        print(f"  Loaded {run['label']}: {meta['turns']} turns, {meta['terminal_reason']}, ${meta['cost_usd']:.4f}")

    return run_data


# ---------------------------------------------------------------------------
# Plot
# ---------------------------------------------------------------------------

def make_plot(run_data, seed, config_name, budget_usd, out_path: Path):
    total_cost = sum(r["cost_usd"] for r in run_data)
    budget_pct = (total_cost / budget_usd) * 100 if budget_usd else 0

    fig, (ax_funds, ax_cost) = plt.subplots(
        1, 2,
        figsize=(17, 6.5),
        facecolor="#0f1117",
        gridspec_kw={"width_ratios": [3, 1]},
    )

    for ax in [ax_funds, ax_cost]:
        ax.set_facecolor("#1a1d27")
        ax.tick_params(colors="#aaaaaa", labelsize=9)
        for spine in ax.spines.values():
            spine.set_edgecolor("#333344")

    # Horizon annotation (approx end date)
    horizon_years = run_data[0]["horizon_years"] if run_data else 3
    horizon_label = f"{horizon_years}-year horizon"

    fig.suptitle(
        f"YC-Bench · {len(run_data)}-Model Comparison · seed={seed} · {config_name}  ({horizon_label}, 5 employees)\n"
        f"Total API spend: ${total_cost:.2f} / ${budget_usd:.0f} budget  ({budget_pct:.1f}%)",
        color="white", fontsize=13, y=1.02,
    )

    # ── Funds curves ─────────────────────────────────────────────────────────
    ax_funds.axhline(0, color="#e74c3c", linewidth=0.9, linestyle="--", alpha=0.4, zorder=1)
    ax_funds.axhline(INITIAL_FUNDS_CENTS / 100, color="#555577", linewidth=0.7, linestyle=":", alpha=0.6, zorder=1)

    for r in run_data:
        if not r["times"]:
            continue
        reason = r["terminal_reason"]
        turns = r["turns"]

        reason_short = {"bankruptcy": "bankrupt", "horizon_end": "survived!", "error": "error"}.get(reason, reason)
        label = f"{r['label']}  ({turns}t · {reason_short})"

        lw = 2.2 if reason == "horizon_end" else 1.8
        alpha = 1.0 if reason == "horizon_end" else 0.85
        ax_funds.plot(r["times"], r["balances"], color=r["color"],
                      linewidth=lw, alpha=alpha, label=label, zorder=3)

        # Mark terminal point
        marker = "★" if reason == "horizon_end" else "x"
        msize = 100 if reason == "horizon_end" else 70
        ax_funds.scatter([r["times"][-1]], [r["balances"][-1]],
                         color=r["color"], s=msize, zorder=5,
                         marker="*" if reason == "horizon_end" else "x", linewidths=2)

    ax_funds.yaxis.set_major_formatter(
        plt.FuncFormatter(lambda x, _: f"${x/1000:.0f}K" if abs(x) < 1_000_000 else f"${x/1_000_000:.1f}M")
    )
    ax_funds.xaxis.set_major_formatter(mdates.DateFormatter("%b '%y"))
    ax_funds.xaxis.set_major_locator(mdates.MonthLocator(interval=3))
    plt.setp(ax_funds.xaxis.get_majorticklabels(), rotation=30, ha="right")
    ax_funds.set_title("Company Funds Over Sim Time", color="white", fontsize=12, pad=8)
    ax_funds.set_ylabel("Balance (USD)", color="#aaaaaa", fontsize=9)
    ax_funds.legend(fontsize=9, facecolor="#1a1d27", edgecolor="#333344",
                    labelcolor="white", loc="upper right")
    ax_funds.grid(axis="y", color="#333344", linewidth=0.5, linestyle="--")
    ax_funds.text(0.005, 0.03, "← bankruptcy", transform=ax_funds.transAxes,
                  color="#e74c3c", fontsize=7.5, alpha=0.6)

    # ── Cost vs budget bars ──────────────────────────────────────────────────
    labels = [r["label"] for r in run_data]
    costs  = [r["cost_usd"] for r in run_data]
    colors = [r["color"] for r in run_data]
    y_pos  = list(range(len(labels)))

    bars = ax_cost.barh(y_pos, costs, color=colors, alpha=0.85, height=0.55, zorder=3)

    # Total bar
    total_y = len(labels) + 0.6
    ax_cost.barh(total_y, total_cost, color="#f1c40f", alpha=0.45, height=0.45, zorder=3)

    # Budget line
    ax_cost.axvline(budget_usd, color="#f1c40f", linewidth=1.6, linestyle="--",
                    zorder=4, label=f"${budget_usd:.0f} budget")

    # Value labels
    for i, cost in enumerate(costs):
        ax_cost.text(cost + budget_usd * 0.01, i, f"${cost:.3f}",
                     va="center", color="#dddddd", fontsize=8)
    ax_cost.text(total_cost + budget_usd * 0.01, total_y, f"${total_cost:.2f}",
                 va="center", color="#f1c40f", fontsize=8.5, fontweight="bold")

    ax_cost.set_yticks(y_pos + [total_y])
    ax_cost.set_yticklabels(labels + ["TOTAL"], color="#aaaaaa", fontsize=8)
    ax_cost.set_xlabel("API Cost (USD)", color="#aaaaaa", fontsize=9)
    ax_cost.set_title(f"Cost vs ${budget_usd:.0f} Budget", color="white", fontsize=12, pad=8)
    ax_cost.legend(fontsize=8, facecolor="#1a1d27", edgecolor="#333344", labelcolor="white")
    ax_cost.set_xlim(0, max(budget_usd * 1.15, max(costs) * 1.4 if costs else 1))
    ax_cost.grid(axis="x", color="#333344", linewidth=0.5, linestyle="--")

    plt.tight_layout()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    plt.savefig(out_path, dpi=150, bbox_inches="tight", facecolor=fig.get_facecolor())
    print(f"Saved: {out_path}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    args = parse_args()
    out = Path(args.out) if args.out else ROOT / "plots" / "funds_curves.png"

    print(f"Loading runs: seed={args.seed}, config={args.config}")
    run_data = load_run_data(DEFAULT_RUNS, args.seed)

    if not run_data:
        print("No data found. Run benchmarks first.")
    else:
        make_plot(run_data, args.seed, args.config, args.budget, out)
