#!/usr/bin/env python3
"""
Shared eval logic for the python-service-scan-authz-baseline skill.

Provider-specific entry points (copilot_runner.py, claude_runner.py) each
create a model instance and pass it to `main()` here.  All orchestration,
file tools, grading, and benchmarking live in this module.
"""

import argparse
import json
import logging
import math
import os
import sys
import time
from pathlib import Path
from typing import Callable

from langchain.agents import create_agent
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.tools import tool
from pydantic import BaseModel, model_validator

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------

SKILL_DIR = Path(__file__).resolve().parent.parent.parent  # .../python-service-scan-authz-baseline
EVALS_DIR = SKILL_DIR / "evals"
WORKSPACE_DIR = EVALS_DIR / "workspace"
EVALS_JSON = EVALS_DIR / "evals.json"
SKILL_MD = SKILL_DIR / "SKILL.md"

# ---------------------------------------------------------------------------
# Pricing  (USD per 1 000 000 tokens)
# GitHub Models prices from https://docs.github.com/en/billing/reference/costs-for-github-models
# Anthropic prices from https://www.anthropic.com/pricing
# ---------------------------------------------------------------------------

PRICING: dict[str, dict[str, float]] = {
    # GitHub Models (OpenAI)
    "gpt-4.1": {"input": 2.00, "output": 8.00},
    "gpt-4.1-mini": {"input": 0.40, "output": 1.60},
    "gpt-4o-mini": {"input": 0.15, "output": 0.60},
    "gpt-4o": {"input": 2.50, "output": 10.00},
    # Anthropic
    "claude-opus-4": {"input": 15.00, "output": 75.00},
    "claude-sonnet-4": {"input": 3.00, "output": 15.00},
    "claude-haiku-4": {"input": 0.80, "output": 4.00},
    "claude-haiku-3": {"input": 0.80, "output": 4.00},
}


def _find_pricing(model_name: str) -> dict[str, float] | None:
    """Return pricing entry using longest-substring match."""
    best_key, best_len = None, 0
    for key in PRICING:
        if key in model_name and len(key) > best_len:
            best_key, best_len = key, len(key)
    return PRICING[best_key] if best_key else None


def calculate_cost(model_name: str, input_tokens: int, output_tokens: int) -> float | None:
    """Return estimated USD cost, or None if the model is not in the pricing table."""
    pricing = _find_pricing(model_name)
    if pricing is None:
        return None
    return round(
        input_tokens * pricing["input"] / 1_000_000 + output_tokens * pricing["output"] / 1_000_000,
        6,
    )


# ---------------------------------------------------------------------------
# File expansion (resolves glob patterns from eval-case "files" entries)
# ---------------------------------------------------------------------------


def expand_files(patterns: list[str], base: Path) -> list[tuple[str, str]]:
    """
    Expand a list of path patterns (which may contain glob wildcards) relative
    to *base* and return ``(relative_path, content)`` pairs for every matched
    regular file.  Patterns without wildcards are treated as literal paths.
    All resolved paths must remain inside *base*; others are silently skipped.
    """
    base_resolved = base.resolve()
    results: list[tuple[str, str]] = []
    seen: set[Path] = set()
    for pattern in patterns:
        matches = (
            sorted(base.glob(pattern))
            if any(c in pattern for c in ("*", "?", "["))
            else [base / pattern]
        )
        for match in matches:
            resolved = match.resolve()
            if not str(resolved).startswith(str(base_resolved)):
                continue
            if not resolved.is_file() or resolved in seen:
                continue
            seen.add(resolved)
            rel = str(resolved.relative_to(base_resolved))
            logger.debug("expand_files: %s", rel)
            results.append((rel, resolved.read_text()))
    return results


# ---------------------------------------------------------------------------
# File-system tools (scoped to SKILL_DIR so the agent can read fixtures)
# ---------------------------------------------------------------------------


def make_tools(base: Path, written_files: dict) -> list:
    @tool
    def read_file(path: str) -> str:
        """Read a text file. Path is relative to the skill root directory."""
        target = (base / path).resolve()
        if not str(target).startswith(str(base.resolve())):
            return "Error: path is outside the allowed directory"
        if not target.exists():
            return f"Error: file not found: {path}"
        if target.is_dir():
            return f"Error: {path} is a directory — use list_directory to see its contents"
        return target.read_text()

    @tool
    def list_directory(path: str) -> str:
        """List contents of a directory. Path is relative to the skill root directory."""
        target = (base / path).resolve()
        if not str(target).startswith(str(base.resolve())):
            return "Error: path is outside the allowed directory"
        if not target.exists():
            return f"Error: path not found: {path}"
        lines = []
        for entry in sorted(target.iterdir()):
            kind = "dir " if entry.is_dir() else "file"
            lines.append(f"{kind}  {entry.relative_to(base)}")
        return "\n".join(lines)

    @tool
    def write_auth_report(content: str) -> str:
        """Write the completed AUTH.md report. Call this once with the full markdown content."""
        written_files["AUTH.md"] = content
        return f"AUTH.md written ({len(content)} characters)."

    return [read_file, list_directory, write_auth_report]


# ---------------------------------------------------------------------------
# Agent runner
# ---------------------------------------------------------------------------


def run_agent(eval_case: dict, with_skill: bool, model: BaseChatModel) -> tuple[str, int, int, int]:
    """Run one eval case. Returns (response_text, input_tokens, output_tokens, duration_ms)."""
    skill_md = SKILL_MD.read_text()

    if with_skill:
        system = (
            "You are a software engineering assistant. "
            "Follow the skill instructions below exactly before doing anything else.\n\n"
            f"<skill>\n{skill_md}\n</skill>\n\n"
            "The fixture files are under evals/files/. "
            "Use read_file and list_directory to access them. "
            "When you have finished the report, call write_auth_report with the full "
            "markdown content. "
            "Do not print the report as your text response — write it via the tool."
        )
    else:
        system = (
            "You are a software engineering assistant helping with "
            "Python microservice authorization. "
            "The fixture files are under evals/files/. "
            "Use read_file and list_directory to access them. "
            "When you have finished the report, call write_auth_report with the full "
            "markdown content. "
            "Do not print the report as your text response — write it via the tool."
        )

    written_files: dict = {}
    tools = make_tools(SKILL_DIR, written_files)
    agent = create_agent(model, tools, system_prompt=system)

    user_content = eval_case["prompt"]
    file_patterns = eval_case.get("files", [])
    if file_patterns:
        injected = expand_files(file_patterns, SKILL_DIR)
        if injected:
            sections = "\n\n".join(f"### {path}\n```\n{content}\n```" for path, content in injected)
            user_content = f"{user_content}\n\n<context>\n{sections}\n</context>"

    start = time.time()
    result = agent.invoke(
        {"messages": [HumanMessage(content=user_content)]},
        config={"recursion_limit": 50},
    )
    duration_ms = int((time.time() - start) * 1000)

    # Prefer content written via write_auth_report tool; fall back to last AIMessage
    response_text = written_files.get("AUTH.md", "")
    if not response_text:
        for msg in reversed(result["messages"]):
            if isinstance(msg, AIMessage) and not msg.tool_calls:
                response_text = msg.content if isinstance(msg.content, str) else str(msg.content)
                break

    # Sum token usage across all messages that carry it
    input_tokens = 0
    output_tokens = 0
    for msg in result["messages"]:
        um = getattr(msg, "usage_metadata", None)
        if isinstance(um, dict):
            input_tokens += um.get("input_tokens", 0)
            output_tokens += um.get("output_tokens", 0)
        else:
            rm = getattr(msg, "response_metadata", {}) or {}
            usage = rm.get("usage", {}) or {}
            input_tokens += usage.get("input_tokens", 0) or usage.get("prompt_tokens", 0)
            output_tokens += usage.get("output_tokens", 0) or usage.get("completion_tokens", 0)

    return response_text, input_tokens, output_tokens, duration_ms


# ---------------------------------------------------------------------------
# Grading
# ---------------------------------------------------------------------------


class _AssertionResult(BaseModel):
    text: str
    passed: bool
    evidence: str


class _GradingResponse(BaseModel):
    assertion_results: list[_AssertionResult]

    @model_validator(mode="before")
    @classmethod
    def _coerce_stringified_list(cls, v):
        # Some models return assertion_results as a JSON-encoded string instead
        # of an inline array. Parse it if needed.
        if isinstance(v, dict) and isinstance(v.get("assertion_results"), str):
            v["assertion_results"] = json.loads(v["assertion_results"])
        return v


def grade(response: str, eval_case: dict, model: BaseChatModel) -> dict:
    """Grade each assertion in eval_case against the response text."""
    assertions = eval_case.get("assertions", [])
    if not assertions:
        return {
            "assertion_results": [],
            "summary": {"passed": 0, "failed": 0, "total": 0, "pass_rate": 0.0},
        }

    numbered = "\n".join(f"{i + 1}. {a}" for i, a in enumerate(assertions))
    prompt = f"""Grade the AI output below against each assertion.

Rules:
- PASS requires concrete evidence quoted or referenced from the output.
- Do not give benefit of the doubt — if not clearly present, it FAILS.
- Evidence must be specific (quote a phrase, cite a section heading, etc.).

## Output
{response}

## Assertions
{numbered}

For each assertion, set passed=true only if concrete evidence is present in the output."""

    grader = model.with_structured_output(_GradingResponse, method="json_schema")
    grading: _GradingResponse = grader.invoke([HumanMessage(content=prompt)])

    passed = sum(1 for r in grading.assertion_results if r.passed)
    total = len(grading.assertion_results)
    return {
        "assertion_results": [r.model_dump() for r in grading.assertion_results],
        "summary": {
            "passed": passed,
            "failed": total - passed,
            "total": total,
            "pass_rate": round(passed / total, 3) if total else 0.0,
        },
    }


# ---------------------------------------------------------------------------
# Benchmark aggregation
# ---------------------------------------------------------------------------


def _stats(values: list[float]) -> dict:
    if not values:
        return {"mean": 0.0, "stddev": 0.0}
    mean = sum(values) / len(values)
    variance = sum((v - mean) ** 2 for v in values) / len(values)
    return {"mean": round(mean, 3), "stddev": round(math.sqrt(variance), 3)}


def compute_benchmark(iteration_dir: Path, evals: list[dict]) -> dict:
    summary = {}
    for config in ("with_skill", "without_skill"):
        pass_rates, times, tokens, costs = [], [], [], []
        for ev in evals:
            grading_path = iteration_dir / _eval_dir_name(ev) / config / "grading.json"
            timing_path = iteration_dir / _eval_dir_name(ev) / config / "timing.json"
            if grading_path.exists():
                g = json.loads(grading_path.read_text())
                pass_rates.append(g["summary"]["pass_rate"])
            if timing_path.exists():
                t = json.loads(timing_path.read_text())
                times.append(t["duration_ms"] / 1000)
                tokens.append(t["total_tokens"])
                if t.get("cost_usd") is not None:
                    costs.append(t["cost_usd"])
        summary[config] = {
            "pass_rate": _stats(pass_rates),
            "time_seconds": _stats(times),
            "tokens": _stats(tokens),
            "cost_usd": _stats(costs) if costs else None,
        }

    ws = summary.get("with_skill", {})
    wos = summary.get("without_skill", {})
    summary["delta"] = {
        "pass_rate": round(
            ws.get("pass_rate", {}).get("mean", 0) - wos.get("pass_rate", {}).get("mean", 0), 3
        ),
        "time_seconds": round(
            ws.get("time_seconds", {}).get("mean", 0) - wos.get("time_seconds", {}).get("mean", 0),
            3,
        ),
        "tokens": round(
            ws.get("tokens", {}).get("mean", 0) - wos.get("tokens", {}).get("mean", 0), 3
        ),
        "cost_usd": round(
            (ws.get("cost_usd") or {}).get("mean", 0) - (wos.get("cost_usd") or {}).get("mean", 0),
            6,
        ),
    }
    return {"run_summary": summary}


# ---------------------------------------------------------------------------
# Eval directory naming
# ---------------------------------------------------------------------------


def _eval_dir_name(ev: dict) -> str:
    case_id = ev.get("id") or ev["name"]
    service = ev.get("service", "")
    return f"eval-{case_id}-{service}" if service else f"eval-{case_id}"


# ---------------------------------------------------------------------------
# Orchestration helpers
# ---------------------------------------------------------------------------


def resolve_iteration(force: int | None) -> Path:
    WORKSPACE_DIR.mkdir(parents=True, exist_ok=True)
    if force is not None:
        d = WORKSPACE_DIR / f"iteration-{force}"
        d.mkdir(exist_ok=True)
        return d
    existing = sorted(WORKSPACE_DIR.glob("iteration-*"), key=lambda p: int(p.name.split("-")[1]))
    n = (int(existing[-1].name.split("-")[1]) + 1) if existing else 1
    d = WORKSPACE_DIR / f"iteration-{n}"
    d.mkdir()
    return d


def run_and_save(
    eval_case: dict, config: str, iteration_dir: Path, model: BaseChatModel, model_name: str
) -> float | None:
    """Run one eval case and save outputs. Returns cost_usd (or None if unknown)."""
    with_skill = config == "with_skill"
    out_dir = iteration_dir / _eval_dir_name(eval_case) / config
    (out_dir / "outputs").mkdir(parents=True, exist_ok=True)

    print(f"  [{config}] running agent...", flush=True)
    response, input_tokens, output_tokens, duration_ms = run_agent(eval_case, with_skill, model)
    total_tokens = input_tokens + output_tokens
    cost_usd = calculate_cost(model_name, input_tokens, output_tokens)

    cost_str = f"${cost_usd:.4f}" if cost_usd is not None else "$unknown"
    print(
        f"  [{config}] done — {duration_ms / 1000:.1f}s, {total_tokens:,} tokens ({cost_str})",
        flush=True,
    )

    (out_dir / "outputs" / "response.md").write_text(response)
    (out_dir / "timing.json").write_text(
        json.dumps(
            {
                "model": model_name,
                "input_tokens": input_tokens,
                "output_tokens": output_tokens,
                "total_tokens": total_tokens,
                "duration_ms": duration_ms,
                "cost_usd": cost_usd,
            },
            indent=2,
        )
    )

    print(f"  [{config}] grading...", flush=True)
    grading = grade(response, eval_case, model)
    (out_dir / "grading.json").write_text(json.dumps(grading, indent=2))
    s = grading["summary"]
    passed_str = f"{s['passed']}/{s['total']} assertions passed ({s['pass_rate']:.0%})"
    print(f"  [{config}] {passed_str}", flush=True)

    return cost_usd


def grade_only(eval_case: dict, config: str, iteration_dir: Path, model: BaseChatModel):
    out_dir = iteration_dir / _eval_dir_name(eval_case) / config
    response_path = out_dir / "outputs" / "response.md"
    if not response_path.exists():
        print(f"  [{config}] no output found, skipping")
        return
    grading = grade(response_path.read_text(), eval_case, model)
    (out_dir / "grading.json").write_text(json.dumps(grading, indent=2))
    s = grading["summary"]
    passed_str = f"{s['passed']}/{s['total']} assertions passed ({s['pass_rate']:.0%})"
    print(f"  [{config}] {passed_str}", flush=True)


# ---------------------------------------------------------------------------
# Main entry point — called by provider runners
# ---------------------------------------------------------------------------


def main(make_model: Callable[[str], BaseChatModel], default_model: str):
    """
    Orchestrate the full eval loop.

    Args:
        make_model:     Factory that accepts a model name string and returns a
                        configured BaseChatModel instance.
        default_model:  Default model name for this provider.
    """
    parser = argparse.ArgumentParser(description="Skill eval runner")
    parser.add_argument("--case", help="Run only this eval case by id")
    parser.add_argument(
        "--model",
        default=os.getenv("EVAL_MODEL", default_model),
        help="Model name (overrides EVAL_MODEL env var)",
    )
    parser.add_argument(
        "--grade-only", action="store_true", help="Re-grade existing outputs, skip agent runs"
    )
    parser.add_argument("--no-baseline", action="store_true", help="Skip without_skill runs")
    parser.add_argument("--iter", type=int, help="Force iteration number (default: auto-increment)")
    parser.add_argument(
        "-d", "--debug", action="store_true", help="Enable debug logging (e.g. expand_files output)"
    )
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.DEBUG if args.debug else logging.WARNING,
        format="%(levelname)s %(name)s: %(message)s",
        stream=sys.stderr,
    )

    all_evals = json.loads(EVALS_JSON.read_text())["evals"]
    evals = all_evals
    if args.case:
        evals = [e for e in all_evals if e.get("id", e["name"]) == args.case]
        if not evals:
            raise SystemExit(f"No eval with id '{args.case}'")

    iteration_dir = resolve_iteration(args.iter)
    print(f"Iteration directory: {iteration_dir.relative_to(SKILL_DIR)}")
    print(f"Model: {args.model}\n")

    model = make_model(args.model)
    configs = ["with_skill"] + ([] if args.no_baseline else ["without_skill"])

    total_cost: float | None = None

    for ev in evals:
        print(f"=== {ev.get('id', ev['name'])} — {ev['name']} ===")
        for config in configs:
            if args.grade_only:
                grade_only(ev, config, iteration_dir, model)
            else:
                cost = run_and_save(ev, config, iteration_dir, model, args.model)
                if cost is not None:
                    total_cost = (total_cost or 0.0) + cost
        print()

    benchmark = compute_benchmark(iteration_dir, all_evals)
    benchmark_path = iteration_dir / "benchmark.json"
    benchmark_path.write_text(json.dumps(benchmark, indent=2))
    print(f"Benchmark saved to {benchmark_path.relative_to(SKILL_DIR)}")

    delta = benchmark["run_summary"].get("delta", {})
    pr = delta.get("pass_rate")
    if isinstance(pr, float):
        print(f"Pass rate delta (with - without skill): {pr:+.1%}")

    if total_cost is not None:
        print(f"Total estimated cost (this run): ${total_cost:.4f}")
    else:
        print("Total estimated cost: $unknown (model not in pricing table)")