"""Repo-approved G6 synthetic evaluator version comparison producer.

The repository benchmark owns gold labels and the baseline/candidate outputs.
Runtime code may bind those immutable observations to approved synthetic
measurement anchors, but it must never invent a candidate or read transcript
content.
"""

from __future__ import annotations

import hashlib
import json
from collections import defaultdict
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path
from typing import Any
from uuid import UUID, uuid5

import asyncpg

from ..contracts.supervision_research import (
    EvaluationVersionBatch,
    LedgerEvidencePointer,
    SupervisionResearchBenchmarkPack,
    VersionedEvaluationObservation,
)
from . import supervision_research_store
from .supervision_research import (
    compare_evaluation_versions,
    load_supervision_research_benchmark,
)


BENCHMARK_PATH = (
    Path(__file__).resolve().parents[1]
    / "data"
    / "supervision_research_benchmark_g6.v1.json"
)
_EVALUATOR_NAMESPACE = UUID("0ee4f677-979d-4635-9201-3aefc89ec71c")
_FORBIDDEN_SOURCE_KEYS = {"raw_transcript", "transcript", "utterance_text"}


@dataclass(frozen=True, slots=True)
class ValidatedRepositoryBenchmark:
    pack: SupervisionResearchBenchmarkPack
    content_sha256: str


@dataclass(frozen=True, slots=True)
class RepositoryComparisonInput:
    benchmark: ValidatedRepositoryBenchmark
    cohort_id: str
    baseline: EvaluationVersionBatch
    candidate: EvaluationVersionBatch
    pointers_by_event_id: Mapping[str, LedgerEvidencePointer]
    learner_ids_by_event_id: Mapping[str, UUID]


def _canonical_json(payload: Any) -> str:
    return json.dumps(
        payload,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
    )


def _forbidden_keys(payload: Any) -> set[str]:
    if isinstance(payload, Mapping):
        found = _FORBIDDEN_SOURCE_KEYS & {str(key) for key in payload}
        for value in payload.values():
            found.update(_forbidden_keys(value))
        return found
    if isinstance(payload, Sequence) and not isinstance(payload, (str, bytes)):
        found: set[str] = set()
        for value in payload:
            found.update(_forbidden_keys(value))
        return found
    return set()


@lru_cache(maxsize=4)
def _load_validated_repository_benchmark(
    resolved_path: str,
) -> ValidatedRepositoryBenchmark:
    path = Path(resolved_path)
    raw = json.loads(path.read_text(encoding="utf-8"))
    forbidden = _forbidden_keys(raw)
    if forbidden:
        raise ValueError(
            "G6 repository benchmark contains forbidden source text fields: "
            + ",".join(sorted(forbidden))
        )
    pack = load_supervision_research_benchmark(path)
    if pack.data_classification != "synthetic_educational":
        raise ValueError("G6 repository benchmark must remain synthetic educational")
    if pack.clinical_claim_allowed:
        raise ValueError("G6 repository benchmark cannot allow clinical claims")

    baseline_by_key = {
        (item.case_ref, item.competency_id): item
        for item in pack.baseline_batch.observations
    }
    candidate_by_key = {
        (item.case_ref, item.competency_id): item
        for item in pack.candidate_batch.observations
    }
    if set(baseline_by_key) != set(candidate_by_key):
        raise ValueError("G6 candidate must use the repository baseline gold case set")
    for key, baseline_item in baseline_by_key.items():
        candidate_item = candidate_by_key[key]
        if (
            baseline_item.gold_label != candidate_item.gold_label
            or baseline_item.synthetic_subgroup != candidate_item.synthetic_subgroup
        ):
            raise ValueError(
                "G6 candidate cannot replace repository gold labels or subgroups"
            )
    provenance = (
        pack.baseline_batch.model,
        pack.baseline_batch.prompt_version,
        pack.baseline_batch.instrument_id,
        pack.baseline_batch.instrument_version,
    )
    candidate_provenance = (
        pack.candidate_batch.model,
        pack.candidate_batch.prompt_version,
        pack.candidate_batch.instrument_id,
        pack.candidate_batch.instrument_version,
    )
    if provenance == candidate_provenance:
        raise ValueError("G6 candidate must identify a distinct evaluator version")

    source_ids = [
        item.evidence_event_id
        for batch in (pack.baseline_batch, pack.candidate_batch)
        for item in batch.observations
    ]
    if len(source_ids) != len(set(source_ids)):
        raise ValueError("G6 repository benchmark evidence ids must be unique")
    report = compare_evaluation_versions(pack.baseline_batch, pack.candidate_batch)
    if report.status != pack.expected_drift_status:
        raise ValueError("G6 repository benchmark expected drift status is stale")
    return ValidatedRepositoryBenchmark(
        pack=pack,
        content_sha256=hashlib.sha256(_canonical_json(raw).encode("utf-8")).hexdigest(),
    )


def load_validated_repository_benchmark(
    path: str | Path = BENCHMARK_PATH,
) -> ValidatedRepositoryBenchmark:
    return _load_validated_repository_benchmark(str(Path(path).resolve()))


def benchmark_anchor_metadata(
    benchmark: ValidatedRepositoryBenchmark,
    batch: EvaluationVersionBatch,
    observation: VersionedEvaluationObservation,
) -> dict[str, Any]:
    """Return the exact safe metadata contract required from a runtime anchor."""

    return {
        "benchmark_schema_version": benchmark.pack.schema_version,
        "benchmark_pack_version": benchmark.pack.version,
        "benchmark_content_sha256": benchmark.content_sha256,
        "data_classification": benchmark.pack.data_classification,
        "clinical_claim_allowed": False,
        "source_evidence_event_id": observation.evidence_event_id,
        "batch_id": batch.batch_id,
        "model": batch.model,
        "prompt_version": batch.prompt_version,
        "instrument_id": batch.instrument_id,
        "instrument_version": batch.instrument_version,
        "case_ref": observation.case_ref,
        "competency_id": observation.competency_id,
        "synthetic_subgroup": observation.synthetic_subgroup,
        "gold_label": observation.gold_label,
        "predicted_label": observation.predicted_label,
    }


def _row_metadata(row: Mapping[str, Any]) -> Mapping[str, Any]:
    metadata = row.get("metadata")
    if isinstance(metadata, str):
        metadata = json.loads(metadata)
    if not isinstance(metadata, Mapping):
        raise ValueError("G6 benchmark measurement anchor metadata is invalid")
    forbidden = _forbidden_keys(metadata)
    if forbidden:
        raise ValueError("G6 benchmark anchor contains forbidden source text fields")
    return metadata


def _bind_batch(
    batch: EvaluationVersionBatch,
    rows_by_source_id: Mapping[str, Mapping[str, Any]],
) -> EvaluationVersionBatch:
    payload = batch.model_dump(mode="json")
    for observation in payload["observations"]:
        source_id = str(observation["evidence_event_id"])
        observation["evidence_event_id"] = str(
            rows_by_source_id[source_id]["measurement_id"]
        )
    return EvaluationVersionBatch.model_validate(payload)


def build_repository_comparison_input(
    rows: Sequence[Mapping[str, Any]],
    *,
    benchmark: ValidatedRepositoryBenchmark | None = None,
) -> RepositoryComparisonInput | None:
    benchmark = benchmark or load_validated_repository_benchmark()
    expected: dict[
        str, tuple[EvaluationVersionBatch, VersionedEvaluationObservation]
    ] = {}
    for batch in (benchmark.pack.baseline_batch, benchmark.pack.candidate_batch):
        for observation in batch.observations:
            expected[observation.evidence_event_id] = (batch, observation)

    by_cohort: dict[str, dict[str, Mapping[str, Any]]] = defaultdict(dict)
    for row in rows:
        metadata = _row_metadata(row)
        source_id = str(metadata.get("source_evidence_event_id") or "")
        if source_id not in expected:
            raise ValueError("G6 benchmark anchor references an unapproved source id")
        cohort_id = str(row.get("cohort_id") or "")
        if not cohort_id:
            raise ValueError("G6 benchmark anchor cohort is missing")
        if source_id in by_cohort[cohort_id]:
            raise ValueError("G6 benchmark anchor source id is duplicated")
        batch, observation = expected[source_id]
        required = benchmark_anchor_metadata(benchmark, batch, observation)
        if any(metadata.get(key) != value for key, value in required.items()):
            raise ValueError(
                "G6 benchmark anchor provenance differs from repository gold"
            )
        by_cohort[cohort_id][source_id] = row

    complete = [
        (cohort_id, mapped)
        for cohort_id, mapped in by_cohort.items()
        if set(mapped) == set(expected)
    ]
    if not complete:
        return None
    if len(complete) > 1:
        raise ValueError(
            "G6 repository benchmark has multiple complete runtime cohorts"
        )
    cohort_id, rows_by_source_id = complete[0]
    baseline = _bind_batch(benchmark.pack.baseline_batch, rows_by_source_id)
    candidate = _bind_batch(benchmark.pack.candidate_batch, rows_by_source_id)
    pointers: dict[str, LedgerEvidencePointer] = {}
    learner_ids: dict[str, UUID] = {}
    for row in rows_by_source_id.values():
        event_id = str(row["measurement_id"])
        pointers[event_id] = LedgerEvidencePointer(
            ledger="measurement_event",
            event_id=event_id,
            session_id=str(row["session_id"]),
            route_hint=(
                f"/research/benchmarks/supervision-research/{benchmark.pack.version}"
            ),
        )
        learner_ids[event_id] = UUID(str(row["learner_id"]))
    return RepositoryComparisonInput(
        benchmark=benchmark,
        cohort_id=cohort_id,
        baseline=baseline,
        candidate=candidate,
        pointers_by_event_id=pointers,
        learner_ids_by_event_id=learner_ids,
    )


async def _load_repository_benchmark_anchors(
    conn: asyncpg.Connection,
    benchmark: ValidatedRepositoryBenchmark,
) -> Sequence[Mapping[str, Any]]:
    return await conn.fetch(
        """
        SELECT m.measurement_id, m.session_id, s.learner_id, u.cohort AS cohort_id,
               m.metadata
        FROM app.measurement_event m
        JOIN app.sessions s ON s.id = m.session_id
        JOIN app.app_user u ON u.user_id = s.learner_id
        WHERE m.status = 'ready'
          AND m.source_kind = 'observed_runtime'
          AND m.perspective = 'runtime_observation'
          AND m.metadata->>'benchmark_schema_version' = $1
          AND m.metadata->>'benchmark_pack_version' = $2
          AND m.metadata->>'benchmark_content_sha256' = $3
          AND m.metadata->>'data_classification' = 'synthetic_educational'
          AND m.metadata->>'clinical_claim_allowed' = 'false'
        ORDER BY u.cohort, m.measurement_id
        """,
        benchmark.pack.schema_version,
        benchmark.pack.version,
        benchmark.content_sha256,
    )


def _stable_ids(comparison: RepositoryComparisonInput) -> dict[str, UUID]:
    key = f"{comparison.benchmark.content_sha256}:{comparison.cohort_id}"
    return {
        name: uuid5(_EVALUATOR_NAMESPACE, f"{name}:{key}")
        for name in (
            "comparison-submission",
            "drift-report",
            "baseline-submission",
            "baseline-batch",
            "candidate-submission",
            "candidate-batch",
        )
    }


async def produce_repository_version_comparison(
    conn: asyncpg.Connection,
    *,
    benchmark_path: str | Path = BENCHMARK_PATH,
) -> dict[str, Any]:
    benchmark = load_validated_repository_benchmark(benchmark_path)
    comparison = build_repository_comparison_input(
        await _load_repository_benchmark_anchors(conn, benchmark),
        benchmark=benchmark,
    )
    if comparison is None:
        return {
            "status": "skipped",
            "reason": "repo_approved_synthetic_evidence_incomplete",
            "benchmark_schema_version": benchmark.pack.schema_version,
            "benchmark_version": benchmark.pack.version,
            "benchmark_content_sha256": benchmark.content_sha256,
            "data_classification": benchmark.pack.data_classification,
            "raw_transcript_included": False,
            "clinical_claim_allowed": False,
        }
    ids = _stable_ids(comparison)
    result = await supervision_research_store.append_evaluation_comparison(
        conn,
        submission_id=ids["comparison-submission"],
        drift_report_id=ids["drift-report"],
        baseline_submission_id=ids["baseline-submission"],
        baseline_batch_record_id=ids["baseline-batch"],
        candidate_submission_id=ids["candidate-submission"],
        candidate_batch_record_id=ids["candidate-batch"],
        cohort_id=comparison.cohort_id,
        baseline=comparison.baseline,
        candidate=comparison.candidate,
        pointers_by_event_id=comparison.pointers_by_event_id,
        learner_ids_by_event_id=comparison.learner_ids_by_event_id,
    )
    return {
        **result,
        "benchmark_schema_version": benchmark.pack.schema_version,
        "benchmark_version": benchmark.pack.version,
        "benchmark_content_sha256": benchmark.content_sha256,
        "data_classification": benchmark.pack.data_classification,
        "baseline_provenance": {
            "model": comparison.baseline.model,
            "prompt_version": comparison.baseline.prompt_version,
            "instrument_id": comparison.baseline.instrument_id,
            "instrument_version": comparison.baseline.instrument_version,
        },
        "candidate_provenance": {
            "model": comparison.candidate.model,
            "prompt_version": comparison.candidate.prompt_version,
            "instrument_id": comparison.candidate.instrument_id,
            "instrument_version": comparison.candidate.instrument_version,
        },
        "raw_transcript_included": False,
        "clinical_claim_allowed": False,
    }


__all__ = [
    "BENCHMARK_PATH",
    "RepositoryComparisonInput",
    "ValidatedRepositoryBenchmark",
    "benchmark_anchor_metadata",
    "build_repository_comparison_input",
    "load_validated_repository_benchmark",
    "produce_repository_version_comparison",
]