#!/usr/bin/env python3
"""Populate a quality ledger from a small adaptive-routing fixture batch."""

from __future__ import annotations

import argparse
import sys
from dataclasses import dataclass
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from llm_connect.adapter import LLMAdapter
from llm_connect.grading import ExactMatchJudge, PairedGrader
from llm_connect.models import LLMResponse, RunConfig
from llm_connect.quality import QualityLedger
from llm_connect.routing import AdaptiveRoutingPolicy, RoutingRule
from llm_connect.shadowing import ShadowingAdapter


@dataclass
class FixtureAdapter(LLMAdapter):
    adapter_id: str
    response_text: str
    cost_usd: float

    def execute_prompt(self, prompt: str, config: RunConfig) -> LLMResponse:
        prompt_tokens = len(prompt.split())
        completion_tokens = len(self.response_text.split())
        return LLMResponse(
            content=self.response_text,
            model=self.adapter_id,
            usage={
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "total_tokens": prompt_tokens + completion_tokens,
            },
            metadata={"cost_usd": self.cost_usd, "latency_ms": 25.0},
        )

    def validate_config(self, config: RunConfig) -> bool:
        return True


def build_candidates() -> dict[str, FixtureAdapter]:
    return {
        "openrouter-cheap-fixture": FixtureAdapter(
            "openrouter-cheap-fixture",
            "summary",
            0.001,
        ),
        "openrouter-mid-fixture": FixtureAdapter(
            "openrouter-mid-fixture",
            "summary with entities and relations",
            0.004,
        ),
        "openrouter-premium-fixture": FixtureAdapter(
            "openrouter-premium-fixture",
            "summary with entities and relations",
            0.012,
        ),
        "claude-code-baseline-fixture": FixtureAdapter(
            "claude-code-baseline-fixture",
            "summary with entities and relations",
            0.0,
        ),
    }


def populate_ledger(ledger: QualityLedger) -> dict[str, FixtureAdapter]:
    candidates = build_candidates()
    baseline = candidates["claude-code-baseline-fixture"]
    grader = PairedGrader(ExactMatchJudge())
    prompts = [
        "Summarize chapter one and keep entity names.",
        "Extract relations from chapter two.",
        "Evaluate whether the entity graph is coherent.",
    ]
    config = RunConfig(model_name="fixture")

    for task_type, prompt in zip(
        ["summarize-source", "extract-relations", "evaluate-entity"],
        prompts,
    ):
        for adapter_id, candidate in candidates.items():
            if candidate is baseline:
                continue
            ShadowingAdapter(
                candidate_adapter=candidate,
                baseline_adapter=baseline,
                grader=grader,
                ledger=ledger,
                task_type=task_type,
                adapter_id=adapter_id,
                baseline_adapter_id=baseline.adapter_id,
                shadow_rate=1.0,
                tags={"fixture": "adaptive-routing"},
            ).execute_prompt(prompt, config)

    return candidates


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--ledger",
        default="quality-ledger.jsonl",
        help="Path to the JSONL ledger to populate.",
    )
    args = parser.parse_args()

    ledger = QualityLedger(Path(args.ledger))
    candidates = populate_ledger(ledger)
    policy = AdaptiveRoutingPolicy(
        rules=[
            RoutingRule(
                "summarize-source",
                prefer=candidates["claude-code-baseline-fixture"],
                fallback=candidates["openrouter-mid-fixture"],
            )
        ],
        ledger=ledger,
        adapters_by_id=candidates,
    )

    selected = policy.resolve("summarize-source", quality_floor=0.8)
    print(f"ledger={ledger.path}")
    print(f"observations={len(ledger.read_all())}")
    print(f"selected={selected.adapter_id}")


if __name__ == "__main__":
    main()