"""B1-X focused live extraction; never changes the frozen corpus/gold."""

from __future__ import annotations

import argparse
import asyncio
import hashlib
import inspect
import itertools
import json
import sys
import time
from pathlib import Path
from types import SimpleNamespace
from uuid import UUID

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from sqlalchemy import text  # noqa: E402
from src.config import get_settings  # noqa: E402
from src.core.provider_capacity import watchdog_seconds  # noqa: E402
from src.domain.projects.discovery_b1 import capacity_trace_fields  # noqa: E402
from src.domain.projects.discovery_catalog import build_discovery_payload  # noqa: E402
from src.domain.projects.discovery_readiness import finalize_discovery_result  # noqa: E402
from src.domain.projects.discovery_schemas import (  # noqa: E402
    DiscoverySpec,
    DiscoveryTurnResult,
    Entity,
)
from src.integrations.project_intake_client import DISCOVERY_PROMPT, discover_project  # noqa: E402
from src.workers.tasks.project_discovery import reserve_discovery_call  # noqa: E402

EVALUATION_PROJECT_ID = "1f6556ff-4a16-4143-9b12-e9d9249d1d3d"
EVALUATION_WORKSPACE_ID = "5e890ff3-5c7b-43f5-a4ad-75c47f90bdba"
PROTECTED_OUTPUTS = {
    ROOT / "docs/audits/fixtures/13a_run_01",
    ROOT / "docs/audits/fixtures/13a_run_02",
    ROOT / "docs/audits/fixtures/13b_b1_completion/extraction_before",
    ROOT / "docs/audits/fixtures/13b_b1_admission_gated_01",
    ROOT / "docs/audits/fixtures/13b_b1_live_01",
    ROOT / "docs/audits/fixtures/13b_b1_live_02",
    ROOT / "docs/audits/fixtures/13b_b1_equal_load_01",
    ROOT / "docs/audits/fixtures/13b_b1_equal_load_02_with_rate_limit_fix",
    ROOT / "docs/audits/fixtures/13b_b1_equal_load_03_retry_only",
    ROOT / "docs/audits/fixtures/13b_b1_semantic_stability_live_01",
    ROOT / "docs/audits/fixtures/13b_b1_admission_prod_host_01.operator_summary.json",
}


def write_json(path: Path, value) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def witnesses():
    return [
        {
            "id": "13A-018",
            "class": "implicit_entity_explicit_workflow",
            "prompt": (
                "Create a contact page for our cleaning company, with a short form "
                "and a thank-you message. Staff must log in to see what was sent."
            ),
        },
        {
            "id": "13A-025",
            "class": "implicit_entity_explicit_workflow",
            "prompt": (
                "Our removal business needs to collect job details before giving a price. "
                "Please make a quote request page and a staff review area."
            ),
        },
        {
            "id": "13A-042",
            "class": "archetype_variance_booking",
            "prompt": (
                "A staff-only appointment booking website for a clinic. Receptionists "
                "reserve a room and start time for a customer; no payments or calendar integration."
            ),
        },
        {
            "id": "13A-051",
            "class": "explicit_entity_explicit_workflow",
            "prompt": (
                "Create internal documents for our staff: title, body content and draft or "
                "published status. Employees can create, view and edit their own documents."
            ),
        },
        {
            "id": "13A-060",
            "class": "multiple_entities_explicit_workflow",
            "prompt": (
                "Document records with separate file relationships: edit the document text "
                "and link an already admitted file to it. Signed-in employees only."
            ),
        },
        {
            "id": "13A-063",
            "class": "explicit_entity_explicit_workflow",
            "prompt": (
                "Our purchasing staff need an item directory with names and SKU codes, "
                "where they can add and edit business records. There is no selling online."
            ),
        },
        {
            "id": "13A-075",
            "class": "multiple_workflows",
            "prompt": (
                "An internal CRM with a sales pipeline and task tracking for staff. "
                "Stage changes and tasks are separate operations in our single-company workspace."
            ),
        },
        {
            "id": "13A-085",
            "class": "ambiguous",
            "prompt": "Help with bookings.",
        },
        {
            "id": "13A-103",
            "class": "explicit_workflow_informal",
            "prompt": (
                "hi! want ppl to msg us from the site no account. name email msg saved "
                "then our staff login to read. simple pls"
            ),
        },
        {
            "id": "X-nouns",
            "class": "explicit_distinct_entities",
            "prompt": (
                "Staff in one office save contact submissions and job requests as separate "
                "records. Staff add, view and update both lists. Keep those names."
            ),
            "expected_names": ["contact submissions", "job requests"],
        },
        {
            "id": "X-plural",
            "class": "explicit_plural_entity",
            "prompt": (
                "Keep inspection reports for signed-in staff in one workspace. Staff create, "
                "view and edit inspection reports. Use that name for the records."
            ),
            "expected_names": ["inspection reports"],
        },
        {
            "id": "X-singular",
            "class": "explicit_singular_entity",
            "prompt": (
                "Keep an inspection report for signed-in staff in one workspace. Staff create, "
                "view and edit each inspection report. Use that name for the records."
            ),
            "expected_names": ["inspection report"],
        },
        {
            "id": "X-none",
            "class": "no_meaningful_workflow",
            "prompt": "Hello",
            "expect_empty": True,
        },
    ]


def extraction_projection(row):
    raw = row.get("discovery", {}).get("raw_turn")
    if raw is None:
        return None
    spec = raw["spec"]
    return {
        "entities": sorted(e["name"].strip().casefold() for e in spec.get("entities") or []),
        "workflow_count": len(spec.get("workflows") or []),
        "archetypes": sorted(spec.get("archetypes") or []),
        "actor_names": sorted(a["name"].strip().casefold() for a in spec.get("actors") or []),
        "requested_capabilities": sorted(raw.get("requested_capabilities") or []),
    }


def _stable(values):
    return len({json.dumps(value, sort_keys=True) for value in values}) <= 1


def diagnostics(runs):
    counters = {
        "pair_count": 0,
        "execution_error_pairs": 0,
        "entity_projection_disagreement": 0,
        "workflow_projection_disagreement": 0,
        "archetype_disagreement": 0,
        "actor_projection_disagreement": 0,
        "capability_projection_disagreement": 0,
        "explicit_workflow_omissions": 0,
        "source_name_failures": 0,
        "empty_input_invention": 0,
    }
    per_witness = {}
    for run in runs:
        for row in run:
            projection = extraction_projection(row)
            bucket = per_witness.setdefault(
                row["id"],
                {
                    "class": row["class"],
                    "completed_runs": 0,
                    "entities": [],
                    "workflow_counts": [],
                    "archetypes": [],
                    "actor_names": [],
                    "requested_capabilities": [],
                },
            )
            if projection is None:
                continue
            bucket["completed_runs"] += 1
            bucket["entities"].append(projection["entities"])
            bucket["workflow_counts"].append(projection["workflow_count"])
            bucket["archetypes"].append(projection["archetypes"])
            bucket["actor_names"].append(projection["actor_names"])
            bucket["requested_capabilities"].append(projection["requested_capabilities"])
            if "explicit_workflow" in row["class"] or row["class"] == "multiple_workflows":
                counters["explicit_workflow_omissions"] += projection["workflow_count"] == 0
            expected = row.get("expected_names", [])
            counters["source_name_failures"] += bool(
                expected and projection["entities"] != sorted(name.casefold() for name in expected)
            )
            counters["empty_input_invention"] += bool(
                row.get("expect_empty") and (projection["entities"] or projection["workflow_count"])
            )
    for witness, bucket in per_witness.items():
        completed = bucket["completed_runs"]
        per_witness[witness] = {
            **bucket,
            "entity_stable": completed > 1 and _stable(bucket["entities"]),
            "workflow_count_stable": completed > 1 and _stable(bucket["workflow_counts"]),
            "archetype_stable": completed > 1 and _stable(bucket["archetypes"]),
            "actor_stable": completed > 1 and _stable(bucket["actor_names"]),
            "capability_stable": completed > 1 and _stable(bucket["requested_capabilities"]),
        }
    for left, right in itertools.combinations(runs, 2):
        for a, b in zip(left, right, strict=True):
            assert a["id"] == b["id"] and a["prompt"] == b["prompt"]
            counters["pair_count"] += 1
            pa, pb = extraction_projection(a), extraction_projection(b)
            if pa is None or pb is None:
                counters["execution_error_pairs"] += 1
                continue
            for field, metric in (
                ("entities", "entity_projection_disagreement"),
                ("workflow_count", "workflow_projection_disagreement"),
                ("archetypes", "archetype_disagreement"),
                ("actor_names", "actor_projection_disagreement"),
                ("requested_capabilities", "capability_projection_disagreement"),
            ):
                counters[metric] += pa[field] != pb[field]
    counters["per_witness"] = per_witness
    return counters


def _refuse_protected(output: Path) -> None:
    resolved = output.resolve()
    for path in PROTECTED_OUTPUTS:
        if resolved == path.resolve() or path.resolve() in resolved.parents:
            raise ValueError("Refusing to write into a protected historical fixture")


def preflight():
    settings = get_settings()
    if "EXTRACTION DISCIPLINE" not in DISCOVERY_PROMPT:
        raise ValueError("Extraction-discipline prompt is not loaded")
    if "INTENT FIDELITY" not in DISCOVERY_PROMPT:
        raise ValueError("Intent-fidelity prompt is not loaded")
    entity_name = Entity.model_fields["name"].description or ""
    workflows = DiscoverySpec.model_fields["workflows"].description or ""
    if "grammatical number" not in entity_name:
        raise ValueError("Entity.name extraction description is not loaded")
    if "including when decision is clarify" not in workflows:
        raise ValueError("Workflow extraction description is not loaded")
    if settings.project_intake_timeout_seconds != 8:
        raise ValueError("Live extraction requires the frozen 8s timeout")
    if settings.project_intake_model != "gpt-4o-mini":
        raise ValueError("Model differs from the frozen baseline")
    if not settings.openai_api_key:
        raise ValueError("Live extraction requires the configured provider credential")
    if not settings.redis_url:
        raise ValueError("Live extraction requires configured Redis")
    if settings.project_intake_calls_per_minute != 60:
        raise ValueError("Expected unchanged 60/min admission policy")
    source = inspect.getsource(reserve_discovery_call)
    if "project-intake:budget:" not in source:
        raise ValueError("Admission implementation is not the production reserve_discovery_call")
    return {
        "preflight": "passed",
        "measurement_kind": "focused_live_extraction",
        "working_directory": str(ROOT),
        "witness_count": len(witnesses()),
        "repeats_default": 4,
        "expected_provider_calls": len(witnesses()) * 4,
        "timeout_seconds": 8,
        "model": settings.project_intake_model,
        "admission_policy_enabled": True,
        "extraction_discipline": True,
        "provider_calls": 0,
        "required_environment_names": ["REDIS_URL", "DATABASE_URL", "OPENAI_API_KEY"],
        "script_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
    }


def verify(output: Path):
    scorecard = json.loads((output / "scorecard.json").read_text(encoding="utf-8"))
    contract = json.loads((output / "request_contract.json").read_text(encoding="utf-8"))
    metadata = json.loads((output / "metadata.json").read_text(encoding="utf-8"))
    assert scorecard["measurement_kind"] == "focused_live_extraction"
    assert contract["measurement_kind"] == "focused_live_extraction"
    assert metadata["environment"] == "production"
    assert metadata["admission_policy_enabled"] is True
    assert "EXTRACTION DISCIPLINE" in contract["prompt"]
    assert contract["repeats"] == metadata["repeats"]
    diagnostics = scorecard["diagnostics"]
    for key in (
        "entity_projection_disagreement",
        "workflow_projection_disagreement",
        "archetype_disagreement",
        "actor_projection_disagreement",
        "capability_projection_disagreement",
        "explicit_workflow_omissions",
        "empty_input_invention",
        "per_witness",
    ):
        assert key in diagnostics
    cases = json.loads((output / "witnesses.json").read_text(encoding="utf-8"))
    for repeat in range(1, metadata["repeats"] + 1):
        for case in cases:
            path = output / f"run_{repeat:02}" / f"{case['id']}.json"
            assert path.is_file(), f"missing {path.name}"
    return {
        "verification": "passed",
        "environment": metadata["environment"],
        "measurement_kind": scorecard["measurement_kind"],
        "diagnostics": {
            key: diagnostics[key]
            for key in diagnostics
            if key != "per_witness"
        },
        "per_witness_ids": sorted(diagnostics["per_witness"]),
    }


async def live_context():
    from src.core.database import async_session_factory, engine
    from src.domain.projects.discovery_service import available_templates

    engine.echo = False
    try:
        async with async_session_factory() as db:
            await db.execute(text("SET TRANSACTION READ ONLY"))
            project = (
                (
                    await db.execute(
                        text(
                            "SELECT type, source_kind FROM projects "
                            "WHERE id=:pid AND workspace_id=:wid"
                        ),
                        {
                            "pid": UUID(EVALUATION_PROJECT_ID),
                            "wid": UUID(EVALUATION_WORKSPACE_ID),
                        },
                    )
                )
                .mappings()
                .first()
            )
            if not project:
                raise ValueError("Project does not belong to the evaluation workspace")
            templates = [
                SimpleNamespace(
                    **{
                        key: getattr(item, key)
                        for key in ("ref", "name", "description", "framework", "tags")
                    }
                )
                for item in await available_templates(db, UUID(EVALUATION_WORKSPACE_ID))
            ]
            revisions = list(
                (await db.execute(text("SELECT version_num FROM alembic_version"))).scalars()
            )
            await db.rollback()
            return dict(project), templates, revisions
    finally:
        await engine.dispose()


async def measure(output: Path, repeats: int, environment: str):
    _refuse_protected(output)
    if output.exists():
        raise ValueError("Select a new experiment directory")
    if repeats < 2:
        raise ValueError("Focused extraction requires at least two repeats")
    settings = get_settings()
    preflight()
    if settings.project_intake_timeout_seconds != 8:
        raise ValueError("Live extraction requires the frozen 8s timeout")
    project, templates, revisions = await live_context()
    cases = witnesses()
    write_json(output / "witnesses.json", cases)
    write_json(
        output / "request_contract.json",
        {
            "prompt": DISCOVERY_PROMPT,
            "schema": DiscoveryTurnResult.model_json_schema(),
            "model": settings.project_intake_model,
            "timeout_seconds": 8,
            "repeats": repeats,
            "temperature": 0,
            "catalog_revisions": revisions,
            "measurement_kind": "focused_live_extraction",
            "admission_policy_enabled": True,
            "environment": environment,
            "witnesses_sha256": hashlib.sha256(
                json.dumps(cases, sort_keys=True).encode()
            ).hexdigest(),
            "acceptance_rule": (
                "Reduce extraction disagreements/omissions without new source-name "
                "or negative-case regressions; execution failures remain separate."
            ),
        },
    )
    write_json(
        output / "metadata.json",
        {
            "phase": "13B",
            "gate": "B1-X",
            "measurement_kind": "focused_live_extraction",
            "environment": environment,
            "admission_policy_enabled": True,
            "repeats": repeats,
            "witness_count": len(cases),
            "expected_provider_calls": len(cases) * repeats,
            "timeout_policy_changed": False,
            "admission_policy_changed": False,
            "core_entity_canonicalization": False,
            "comparison_note": (
                "Codex-host extraction_before is a different environment and is not "
                "a clean controlled before/after pair with this production-host run."
            ),
        },
    )
    semaphore, lock = asyncio.Semaphore(4), asyncio.Lock()
    last_start = 0.0

    async def one(case, repeat):
        nonlocal last_start
        async with semaphore:
            async with lock:
                await asyncio.sleep(max(0, last_start + 2.2 - time.monotonic()))
                last_start = time.monotonic()
            row = {**case, "oracle_kind": "focused", "stratum": case["class"]}
            payload = build_discovery_payload(
                project["type"],
                project["source_kind"],
                {"messages": [{"role": "user", "content": case["prompt"]}]},
                templates,
            )
            started = time.monotonic()
            trace: dict = {}
            try:
                await reserve_discovery_call(settings)
            except Exception as exc:
                row["discovery"] = {
                    "status": "error",
                    "error_type": f"Admission:{type(exc).__name__}",
                }
            else:
                try:
                    turn = await asyncio.wait_for(
                        discover_project(payload, settings, trace=trace),
                        watchdog_seconds(settings, 8),
                    )
                    result = finalize_discovery_result(turn, payload)
                    row["discovery"] = {
                        "status": "success",
                        "raw_turn": turn.model_dump(mode="json"),
                        "result": result.model_dump(mode="json"),
                        **capacity_trace_fields(trace),
                    }
                except Exception as exc:
                    row["discovery"] = {
                        "status": "error",
                        "error_type": type(exc).__name__,
                        **capacity_trace_fields(trace),
                    }
            row["discovery"]["elapsed_seconds"] = time.monotonic() - started
            write_json(output / f"run_{repeat:02}" / f"{case['id']}.json", row)
            print(
                f"repeat={repeat} witness={case['id']} status={row['discovery']['status']}",
                flush=True,
            )
            return row

    runs = []
    for repeat in range(1, repeats + 1):
        runs.append(await asyncio.gather(*(one(case, repeat) for case in cases)))
    result = {
        "measurement_kind": "focused_live_extraction",
        "admission_policy_enabled": True,
        "environment": environment,
        "diagnostics": diagnostics(runs),
    }
    write_json(output / "scorecard.json", result)
    print(
        json.dumps(
            {k: v for k, v in result["diagnostics"].items() if k != "per_witness"},
            indent=2,
        )
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out", type=Path)
    parser.add_argument("--repeats", type=int, default=4)
    parser.add_argument("--environment", choices=("local", "production"), default="local")
    parser.add_argument("--preflight", action="store_true")
    parser.add_argument("--verify", type=Path)
    args = parser.parse_args()
    try:
        if args.preflight:
            print(json.dumps(preflight(), indent=2))
        elif args.verify:
            print(json.dumps(verify(args.verify), indent=2))
        elif args.out:
            asyncio.run(measure(args.out, args.repeats, args.environment))
        else:
            parser.error("Supply --preflight, --verify RECEIPT_DIR, or --out NEW_DIRECTORY")
    except Exception as exc:
        print(json.dumps({"diagnostic_failed": type(exc).__name__}), file=sys.stderr)
        raise SystemExit(1) from None