#!/usr/bin/env python
"""
Phase 2 Script: Run baseline extraction on test set

Usage:
  python scripts/run_baseline.py

This script:
1. Loads the prepared data splits
2. Runs zero-shot extraction on test set
3. Computes baseline metrics (F1, precision, recall)
4. Saves results and evaluation report
"""

import os
import sys
import json
import logging
import subprocess
from pathlib import Path
from datetime import datetime

# Add src to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.baseline import BaselineExtractor, extract_baseline
from src.evaluator import BaselineEvaluator
from src.schemas import CUAD_CATEGORIES
from src.model_providers import get_provider

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def _git_commit_hash() -> str:
    """Best-effort short commit hash so a report can be tied back to the
    exact code that produced it. Returns 'unknown' rather than raising if
    git isn't available or the tree has no commits yet -- this is metadata
    for reproducibility, not something worth failing a run over."""
    try:
        return subprocess.check_output(
            ["git", "rev-parse", "--short", "HEAD"],
            cwd=Path(__file__).resolve().parent.parent,
            stderr=subprocess.DEVNULL,
        ).decode().strip()
    except Exception:
        return "unknown"


def _split_provenance() -> dict:
    """Read the actual split file's own recorded seed/fractions rather than
    hardcoding them in the report -- if prepare_datasets.py's defaults ever
    change, the report must not silently keep reporting stale values."""
    splits_file = Path("data/processed/cuad_splits.json")
    if not splits_file.exists():
        return {}
    with open(splits_file, encoding="utf-8") as f:
        splits = json.load(f)
    return {
        "random_seed": splits.get("random_seed"),
        "split_fractions": splits.get("split_fractions"),
        "n_test_ids": len(splits.get("test_ids", [])),
    }

def load_test_contracts():
    """
    Load test contracts (full text + gold clauses) prepared by
    scripts/prepare_datasets.py.

    NOTE: prepare_datasets.py must be run first -- it persists
    data/processed/cuad_test_contracts.json (full text/clauses), not just
    the ID list in cuad_splits.json. An earlier version of this function
    only read the ID list and always returned [], making Phase 2
    unrunnable even when data existed.
    """
    contracts_file = Path("data/processed/cuad_test_contracts.json")

    if not contracts_file.exists():
        logger.warning(f"Test contracts file not found: {contracts_file}")
        logger.info("Run Phase 1 first: python scripts/prepare_datasets.py")
        return []

    with open(contracts_file, encoding="utf-8") as f:
        contracts = json.load(f)

    logger.info(f"Loaded {len(contracts)} test contracts")
    return contracts

def _not_measured_report(reason: str) -> str:
    split_info = _split_provenance()
    from src.model_providers.local_ollama import DEFAULT_MODEL
    return f"""# Baseline Extraction Evaluation Report

**Phase**: 2 - Baseline Extraction
**Date**: {datetime.utcnow().isoformat()}
**Model**: {DEFAULT_MODEL} (local, $0 API cost) -- not actually invoked, no run occurred
**Code version**: git {_git_commit_hash()}
**Split provenance**: {split_info or 'no split file found'}

## Status: NOT MEASURED

Reason: {reason}

## Next Steps

```bash
python scripts/prepare_datasets.py   # Phase 1: download CUAD, build splits
ollama serve                         # start the local model server (if not already running)
ollama pull {DEFAULT_MODEL}          # one-time model download (~4-5GB, 4-bit quantized)
python scripts/run_baseline.py       # Phase 2: extract + score, no API key needed
```
"""


def generate_eval_report(
    extraction_results: dict,
    metrics: dict = None,
    output_file: str = "docs/EVAL_REPORT_PHASE2.md",
    not_measured_reason: str = None,
):
    """
    Generate the Phase 2 evaluation report.

    If `metrics` is provided (from BaselineEvaluator.compute_f1 +
    compute_confidence_calibration), the report contains real, measured
    numbers. Otherwise it is written as an explicit NOT MEASURED report --
    it must never silently emit the same template regardless of whether
    a run actually happened.
    """
    logger.info(f"Generating evaluation report: {output_file}")
    Path(output_file).parent.mkdir(parents=True, exist_ok=True)

    if metrics is None:
        report = _not_measured_report(not_measured_reason or "Run did not complete")
        with open(output_file, "w") as f:
            f.write(report)
        logger.info(f"Report saved to {output_file} (NOT MEASURED)")
        return

    em = extraction_results["metrics"]
    per_cat = metrics["per_category"]
    # Categories with zero gold support AND zero predictions have f1=None
    # (excluded from the macro average -- see src/evaluator.py); they can't
    # be sorted or formatted as floats alongside scored categories.
    scored_cats = {c: m for c, m in per_cat.items() if m["f1"] is not None}
    sorted_cats = sorted(scored_cats.items(), key=lambda x: x[1]["f1"], reverse=True)
    excluded_cats = metrics.get("excluded_categories", [])

    total_input_tokens = em.get("total_input_tokens", 0)
    total_output_tokens = em.get("total_output_tokens", 0)
    n_contracts = em["total_contracts"]

    per_cat_lines = "\n".join(
        f"| {cat} | {m['f1']:.3f} | {m['precision']:.3f} | {m['recall']:.3f} | {m['support']} |"
        for cat, m in sorted_cats
    )
    excluded_note = (
        f"\n**Excluded from macro F1** (zero gold examples in this split AND zero predictions -- "
        f"no evidence either way): {', '.join(excluded_cats)}\n"
        if excluded_cats else ""
    )

    calibration = metrics.get("calibration") or {}

    def _fmt(value, spec=".4f"):
        return format(value, spec) if value is not None else "N/A"

    brier_str = _fmt(calibration.get("brier_score"))
    spearman_str = _fmt(calibration.get("spearman_rho"))
    mean_conf_str = _fmt(calibration.get("mean_confidence"))
    calibration_n = calibration.get("n", 0)

    split_info = _split_provenance()
    seed = split_info.get("random_seed", "unknown")
    fractions = split_info.get("split_fractions", {})

    report = f"""# Baseline Extraction Evaluation Report

**Phase**: 2 - Baseline Extraction
**Date**: {datetime.utcnow().isoformat()}
**Code version**: git {_git_commit_hash()}
**Model**: {extraction_results.get('model', 'unknown')} (zero-shot, temperature={extraction_results.get('temperature', 'unknown')})
**Prompt version**: {extraction_results.get('prompt_version', 'unknown')}
**Dataset split**: CUAD test ({n_contracts} contracts, seed={seed}, fractions={fractions})
**Metric granularity**: (contract, category) presence -- see src/evaluator.py and docs/DECISIONS.md ADR-018 for why this, not per-span matching, is the correct unit for this task.

## Key Metrics (MEASURED)

| Metric | Value |
|--------|-------|
| Macro F1 (primary, over {metrics.get('macro_f1_n_categories', 'N/A')} of 41 categories) | {metrics['macro_f1']:.4f} |
| Micro F1 | {metrics['micro_f1']:.4f} |
| Micro precision / recall | {metrics.get('micro_precision', 0.0):.4f} / {metrics.get('micro_recall', 0.0):.4f} |
| Total TP / FP / FN (contract,category pairs) | {metrics['global_stats']['tp']} / {metrics['global_stats']['fp']} / {metrics['global_stats']['fn']} |
| Mean confidence | {em.get('avg_confidence', 0.0):.4f} |
| Total time (s) | {em.get('total_time_seconds', 0.0):.1f} |
| Avg time/contract (s) | {(em.get('total_time_seconds', 0.0) / n_contracts) if n_contracts else 0.0:.2f} |
| Contracts with extraction errors | {em.get('failed', 0)}/{n_contracts} ({(em.get('failed', 0) / n_contracts * 100) if n_contracts else 0:.1f}%) |
| Total input tokens | {total_input_tokens} |
| Total output tokens | {total_output_tokens} |
| API cost (USD) | $0 -- local inference, no paid API used |
| GPU | {em.get('gpu_name', 'not recorded')} |
{excluded_note}
## Confidence Calibration

| Metric | Value |
|--------|-------|
| Brier score (lower is better) | {brier_str} |
| Spearman rho (confidence vs. correctness) | {spearman_str} |
| Mean confidence | {mean_conf_str} |
| N (distinct contract,category predictions scored) | {calibration_n} |

## Per-Category F1

| Category | F1 | Precision | Recall | Support (gold) |
|----------|----|-----------|--------|-----------------|
{per_cat_lines}

## Failure Modes

NOT YET ANALYZED -- requires manual review of `data/processed/baseline_results.json`
against gold clauses to categorize false positives/negatives by type.
"""
    with open(output_file, "w") as f:
        f.write(report)
    logger.info(f"Report saved to {output_file} (measured)")

def _gpu_name() -> str:
    """Best-effort GPU identification for the report -- not required for
    the pipeline to run (CPU inference works, just slower), only for
    honest hardware provenance in the generated report."""
    try:
        out = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
            stderr=subprocess.DEVNULL,
        ).decode().strip()
        return out or "no GPU detected"
    except Exception:
        return "nvidia-smi not available"


def main():
    """Main execution."""

    logger.info("=" * 70)
    logger.info("PHASE 2: BASELINE EXTRACTION (local, open-weight model, $0 API cost)")
    logger.info("=" * 70)

    # Check the local model provider is actually reachable and the model is
    # pulled -- this replaces the old ANTHROPIC_API_KEY gate. No cloud
    # credential is required or checked anywhere in this script.
    provider = get_provider()
    if not provider.is_available():
        logger.error(f"Local model provider not ready: {type(provider).__name__} (model={provider.model_id})")
        logger.info("Ensure Ollama is running ('ollama serve') and the model is pulled ('ollama pull %s')" % provider.model_id)
        generate_eval_report(
            {}, not_measured_reason=f"Local provider not ready: {type(provider).__name__} could not reach model {provider.model_id!r}"
        )
        return

    logger.info(f"Local provider ready: {type(provider).__name__} (model={provider.model_id})")

    # Load test contracts
    test_contracts = load_test_contracts()

    if not test_contracts:
        logger.warning("No test contracts loaded")
        logger.info("\nTo proceed:")
        logger.info("1. Run Phase 1: python scripts/prepare_datasets.py")
        logger.info("2. Verify data/processed/cuad_test_contracts.json exists")
        logger.info("3. Run: python scripts/run_baseline.py")
        generate_eval_report(
            {}, not_measured_reason="No test contracts found; run scripts/prepare_datasets.py first"
        )
        return

    # Run baseline extraction
    logger.info(f"\nRunning baseline extraction on {len(test_contracts)} test contracts...")

    results = extract_baseline(
        test_contracts,
        output_file="data/processed/baseline_results.json",
        provider=provider,
    )

    # Aggregate token usage (real, from local inference responses -- not estimated)
    total_input_tokens = sum(
        (r.get("input_tokens") or 0) for r in results["extractions"]
    )
    total_output_tokens = sum(
        (r.get("output_tokens") or 0) for r in results["extractions"]
    )
    results["metrics"]["total_input_tokens"] = total_input_tokens
    results["metrics"]["total_output_tokens"] = total_output_tokens
    results["metrics"]["gpu_name"] = _gpu_name()

    # Flatten predictions and gold clauses (tagging each with contract_id so
    # cross-contract matches can't accidentally score as true positives) and
    # compute real F1 against the gold clauses persisted by prepare_datasets.py.
    predictions = [
        {**clause, "contract_id": r["contract_id"]}
        for r in results["extractions"]
        for clause in r.get("clauses", [])
    ]
    gold = [
        {**clause, "contract_id": contract["id"]}
        for contract in test_contracts
        for clause in contract.get("clauses", [])
    ]

    evaluator = BaselineEvaluator()
    metrics = evaluator.compute_f1(predictions, gold)
    calibration = evaluator.compute_confidence_calibration(predictions, gold)
    metrics["calibration"] = calibration
    evaluator.log_results(metrics)

    # Generate report
    logger.info("\nGenerating evaluation report...")
    generate_eval_report(results, metrics=metrics)

    logger.info("\n" + "=" * 70)
    logger.info("PHASE 2 COMPLETE")
    logger.info("=" * 70)
    logger.info("Results saved to:")
    logger.info("  - data/processed/baseline_results.json")
    logger.info("  - docs/EVAL_REPORT_PHASE2.md")
    logger.info("\nNext: Phase 3 - Retrieval experiments")
    logger.info("=" * 70)

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        logger.error(f"Phase 2 failed: {e}", exc_info=True)
        sys.exit(1)