"""Phase 4: Final test-set evaluation + LLM zero-shot comparison.

Run ONCE after Phase 3 is complete. Touches test.jsonl for the first time.

Steps:
  1. Eval checkpoints/best/ on test split  → results/test_final.json
  2. Run zero-shot LLM baseline on a sample → results/llm_comparison.json
  3. Print the comparison table

Usage:
    python scripts/run_phase4_eval.py
    python scripts/run_phase4_eval.py --skip-llm   # skip LLM call (no API key)
    python scripts/run_phase4_eval.py --llm-sample 100
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from rich.console import Console
from rich.table import Table

from finner.config import RESULTS_DIR, CHECKPOINTS_DIR, SPLITS_DIR, settings
from finner.data.split import load_jsonl
from finner.eval.metrics import compute_seqeval
from finner.eval.run_eval import evaluate_checkpoint, _to_serializable

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
console = Console()


def run_test_eval() -> dict:
    console.print("\n[bold cyan]Step 1: Final test-set evaluation (touching test split for the first time)[/bold cyan]")
    metrics = evaluate_checkpoint(
        checkpoint_path=CHECKPOINTS_DIR / "best",
        split="test",
        save_name="test_final",
    )
    console.print(f"\n[bold green]Final test entity-F1: {metrics['overall_f1']:.4f}[/bold green]")
    return metrics


def run_llm_baseline(sample_size: int) -> dict:
    console.print(f"\n[bold cyan]Step 2: Zero-shot LLM baseline ({sample_size} examples)[/bold cyan]")
    console.print("[yellow]Note: This uses the LLM API for comparison ONLY. Not in the core inference path.[/yellow]")

    from finner.eval.llm_baseline import run_llm_baseline as _run
    test_examples = load_jsonl(SPLITS_DIR / "test.jsonl")

    t0 = time.perf_counter()
    baseline_result = _run(test_examples, sample_size=sample_size)
    elapsed = time.perf_counter() - t0

    llm_metrics = compute_seqeval(baseline_result["flat_labels"], baseline_result["flat_preds"])

    result = {
        "llm_model": settings.llm_model,
        "llm_provider": settings.llm_provider,
        "llm_f1": llm_metrics["overall_f1"],
        "llm_precision": llm_metrics["overall_precision"],
        "llm_recall": llm_metrics["overall_recall"],
        "llm_per_class": llm_metrics["per_class"],
        "llm_avg_latency_ms": baseline_result["avg_latency_ms"],
        "llm_total_examples": baseline_result["total_examples"],
        "llm_total_elapsed_s": elapsed,
    }

    out = RESULTS_DIR / "llm_comparison.json"
    with open(out, "w") as f:
        json.dump(_to_serializable(result), f, indent=2)
    console.print(f"Saved → {out}")
    return result


def merge_comparison(test_metrics: dict, llm_result: dict | None) -> None:
    """Load best-val F1 from runs.jsonl and merge into llm_comparison.json."""
    runs_path = RESULTS_DIR / "runs.jsonl"
    best_val_f1 = 0.0
    if runs_path.exists():
        runs = [json.loads(l) for l in open(runs_path) if l.strip()]
        if runs:
            best_val_f1 = max(r["best_val_entity_f1"] for r in runs)

    comparison = {
        "model_name": settings.model_name,
        "model_params_M": 108,
        "model_f1": test_metrics["overall_f1"],
        "model_per_class": test_metrics["per_class"],
        "best_val_f1": best_val_f1,
    }
    if llm_result:
        comparison.update(llm_result)

    out = RESULTS_DIR / "llm_comparison.json"
    with open(out, "w") as f:
        json.dump(_to_serializable(comparison), f, indent=2)

    # Print comparison table
    table = Table(title="Fine-tuned Encoder vs Zero-shot LLM", show_header=True)
    table.add_column("System", style="bold")
    table.add_column("entity-F1", justify="right")
    table.add_column("Latency/sent", justify="right")
    table.add_column("Cost/1k sent", justify="right")
    table.add_column("Params", justify="right")

    table.add_row(
        "Fine-tuned BERT",
        f"[bold green]{test_metrics['overall_f1']:.4f}[/bold green]",
        "~50ms (CPU)",
        "~free (local)",
        "108M",
    )
    if llm_result:
        table.add_row(
            f"Zero-shot {llm_result.get('llm_model', 'LLM')}",
            f"{llm_result.get('llm_f1', 0):.4f}",
            f"{llm_result.get('llm_avg_latency_ms', 0):.0f}ms",
            "~$0.25 (API)",
            ">> 1B",
        )
    console.print(table)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--skip-llm", action="store_true", help="Skip LLM API call")
    parser.add_argument("--llm-sample", type=int, default=100, help="Number of test examples for LLM eval")
    args = parser.parse_args()

    test_metrics = run_test_eval()

    llm_result = None
    if not args.skip_llm:
        try:
            llm_result = run_llm_baseline(args.llm_sample)
        except Exception as exc:
            console.print(f"[red]LLM baseline failed: {exc}[/red]")
            console.print("[yellow]Run with --skip-llm to skip this step.[/yellow]")

    merge_comparison(test_metrics, llm_result)
    console.print("\n[bold green]✓ Phase 4 complete.[/bold green]")


if __name__ == "__main__":
    main()