#!/usr/bin/env python3
"""Audit semantic duplicates and emit the frozen human-review sample."""

from __future__ import annotations

import argparse
import json
import sys
from collections import Counter
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Sequence

import numpy as np

from prepare_data import (
    DEFAULT_CURATION_REVIEW,
    DEFAULT_FIXTURES,
    DEFAULT_FROZEN_TEST,
    default_source_paths,
    load_reviewed_semantic_exclusions,
)
from purpose_data import (
    LABELS,
    DataError,
    SourceRecord,
    curate_records,
    exclude_reviewed_duplicates,
    load_classifiable_fixtures,
    load_sources,
    prompt_hash,
    write_json,
)
from review_contract import (
    DEFAULT_SAMPLE_FRACTION,
    DEFAULT_SAMPLE_SEED,
    stratified_review_sample,
    write_review_csv,
)
from train import (
    DEFAULT_MODEL,
    DEFAULT_MODEL_REVISION,
    HEAD_TOKENS,
    MAX_LENGTH,
    TAIL_TOKENS,
    encode_fixed_shape,
)


SCRIPT_DIR = Path(__file__).resolve().parent
REPOSITORY_ROOT = SCRIPT_DIR.parent.parent
DEFAULT_REPORT = SCRIPT_DIR / "data" / "semantic-audit-v1.json"
DEFAULT_REVIEW_CSV = SCRIPT_DIR / ".artifacts" / "human-review-v1.csv"
@dataclass(frozen=True)
class AuditRecord:
    prompt: str
    purpose: str
    prompt_hash: str
    origin: str
    source: str
    line: int | None


@dataclass(frozen=True)
class SemanticCandidate:
    left: int
    right: int
    similarity: float


def _relative(path: Path) -> str:
    try:
        return str(path.resolve().relative_to(REPOSITORY_ROOT))
    except ValueError:
        return str(path.resolve())


def semantic_candidates(
    embeddings: np.ndarray,
    purposes: Sequence[str],
    *,
    same_label_threshold: float,
    cross_label_threshold: float,
    neighbors: int,
    block_size: int,
) -> list[SemanticCandidate]:
    """Return de-duplicated high-cosine nearest-neighbor pairs."""

    if embeddings.ndim != 2 or embeddings.shape[0] != len(purposes):
        raise ValueError("embeddings and purposes must have matching record counts")
    if len(purposes) < 2:
        return []
    if not (
        0.0 < same_label_threshold <= 1.0
        and 0.0 < cross_label_threshold <= 1.0
    ):
        raise ValueError("semantic thresholds must be in (0, 1]")
    if neighbors <= 0 or block_size <= 0:
        raise ValueError("neighbors and block size must be positive")

    norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
    if np.any(norms == 0):
        raise ValueError("semantic embeddings must be non-zero")
    normalized = embeddings.astype(np.float32, copy=False) / norms
    top_count = min(neighbors, len(purposes) - 1)
    found: dict[tuple[int, int], float] = {}
    for start in range(0, len(purposes), block_size):
        stop = min(start + block_size, len(purposes))
        similarities = normalized[start:stop] @ normalized.T
        local_rows = np.arange(stop - start)
        similarities[local_rows, np.arange(start, stop)] = -1.0
        candidate_columns = np.argpartition(
            similarities, -top_count, axis=1
        )[:, -top_count:]
        for local_index, columns in enumerate(candidate_columns):
            left = start + local_index
            for right in columns:
                similarity = float(similarities[local_index, right])
                threshold = (
                    same_label_threshold
                    if purposes[left] == purposes[right]
                    else cross_label_threshold
                )
                if similarity < threshold:
                    continue
                pair = (min(left, int(right)), max(left, int(right)))
                found[pair] = max(found.get(pair, -1.0), similarity)
    return [
        SemanticCandidate(left, right, similarity)
        for (left, right), similarity in sorted(
            found.items(), key=lambda item: (-item[1], item[0])
        )
    ]


def _select_device(torch: Any, requested: str) -> Any:
    if requested != "auto":
        return torch.device(requested)
    if torch.cuda.is_available():
        return torch.device("cuda")
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def embed_prompts(
    records: Sequence[AuditRecord],
    *,
    model_name: str,
    model_revision: str,
    device_name: str,
    batch_size: int,
) -> tuple[np.ndarray, str]:
    try:
        import torch
        from transformers import AutoModel, AutoTokenizer
    except ImportError as exc:
        raise DataError(
            "semantic-audit dependencies are missing; install requirements.txt"
        ) from exc

    device = _select_device(torch, device_name)
    tokenizer = AutoTokenizer.from_pretrained(
        model_name, revision=model_revision, use_fast=True
    )
    model = AutoModel.from_pretrained(model_name, revision=model_revision).to(device)
    model.eval()
    chunks: list[np.ndarray] = []
    with torch.inference_mode():
        for start in range(0, len(records), batch_size):
            batch = records[start : start + batch_size]
            encoded = encode_fixed_shape(
                tokenizer,
                [record.prompt for record in batch],
                torch,
            )
            encoded = {key: value.to(device) for key, value in encoded.items()}
            hidden = model(**encoded).last_hidden_state
            mask = encoded["attention_mask"].unsqueeze(-1).expand(hidden.size()).float()
            pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
            pooled = torch.nn.functional.normalize(pooled, p=2, dim=1)
            chunks.append(pooled.cpu().numpy())
    return np.concatenate(chunks), str(device)


def _audit_records(
    curated: Sequence[SourceRecord],
    fixtures: Sequence[dict[str, str]],
    fixtures_path: Path,
) -> list[AuditRecord]:
    records = [
        AuditRecord(
            prompt=record.value["prompt"],
            purpose=record.value["purpose"],
            prompt_hash=prompt_hash(record.value["prompt"]),
            origin="synthetic",
            source=_relative(record.source),
            line=record.line,
        )
        for record in curated
    ]
    records.extend(
        AuditRecord(
            prompt=fixture["prompt"],
            purpose=fixture["purpose"],
            prompt_hash=prompt_hash(fixture["prompt"]),
            origin="shipped-fixture",
            source=_relative(fixtures_path),
            line=None,
        )
        for fixture in fixtures
    )
    return records


def _split_membership() -> dict[str, str]:
    membership: dict[str, str] = {}
    artifact_dir = SCRIPT_DIR / ".artifacts" / "dataset-v1"
    for split in ("train", "validation"):
        path = artifact_dir / f"{split}.jsonl"
        if not path.exists():
            continue
        for line in path.read_text(encoding="utf-8").splitlines():
            value = json.loads(line)
            membership[prompt_hash(value["prompt"])] = split
    for line in DEFAULT_FROZEN_TEST.read_text(encoding="utf-8").splitlines():
        value = json.loads(line)
        membership[prompt_hash(value["prompt"])] = "test"
    return membership


def audit(args: argparse.Namespace) -> dict[str, Any]:
    sources = (
        [path.resolve() for path in args.source]
        if args.source
        else default_source_paths()
    )
    fixtures = load_classifiable_fixtures(args.fixtures.resolve())
    source_records = load_sources(sources)
    lexical_curation = curate_records(
        source_records,
        fixtures,
        near_duplicate_threshold=args.word_duplicate_threshold,
    )
    curated = exclude_reviewed_duplicates(
        lexical_curation.records,
        load_reviewed_semantic_exclusions(args.curation_review.resolve()),
    )
    sample = stratified_review_sample(
        curated.records,
        fraction=args.review_fraction,
        seed=args.review_seed,
    )
    write_review_csv(
        args.review_csv.resolve(),
        sample,
        source_formatter=_relative,
    )

    audit_records = _audit_records(
        curated.records, fixtures, args.fixtures.resolve()
    )
    embeddings, device = embed_prompts(
        audit_records,
        model_name=args.model,
        model_revision=args.model_revision,
        device_name=args.device,
        batch_size=args.batch_size,
    )
    candidates = semantic_candidates(
        embeddings,
        [record.purpose for record in audit_records],
        same_label_threshold=args.same_label_threshold,
        cross_label_threshold=args.cross_label_threshold,
        neighbors=args.neighbors,
        block_size=args.block_size,
    )
    membership = _split_membership()

    candidate_values = []
    leakage_candidates = 0
    for candidate in candidates:
        left = audit_records[candidate.left]
        right = audit_records[candidate.right]
        left_split = (
            "shipped-fixture"
            if left.origin == "shipped-fixture"
            else membership.get(left.prompt_hash, "unknown")
        )
        right_split = (
            "shipped-fixture"
            if right.origin == "shipped-fixture"
            else membership.get(right.prompt_hash, "unknown")
        )
        crosses_frozen_boundary = (
            left_split == "train"
            and right_split in {"validation", "test", "shipped-fixture"}
        ) or (
            right_split == "train"
            and left_split in {"validation", "test", "shipped-fixture"}
        )
        leakage_candidates += int(crosses_frozen_boundary)
        candidate_values.append(
            {
                "similarity": round(candidate.similarity, 6),
                "sameLabel": left.purpose == right.purpose,
                "crossesFrozenBoundary": crosses_frozen_boundary,
                "left": {
                    "promptHash": left.prompt_hash,
                    "purpose": left.purpose,
                    "origin": left.origin,
                    "split": left_split,
                    "source": left.source,
                    "line": left.line,
                },
                "right": {
                    "promptHash": right.prompt_hash,
                    "purpose": right.purpose,
                    "origin": right.origin,
                    "split": right_split,
                    "source": right.source,
                    "line": right.line,
                },
            }
        )

    report = {
        "schemaVersion": 1,
        "auditVersion": "purpose-semantic-audit-v1",
        "embedding": {
            "model": args.model,
            "revision": args.model_revision,
            "fixedInputShape": [1, MAX_LENGTH],
            "truncation": {
                "strategy": "head-tail-pair",
                "headTokens": HEAD_TOKENS,
                "tailTokens": TAIL_TOKENS,
            },
            "pooling": "attention-mask mean pooling followed by L2 normalization",
            "device": device,
        },
        "population": {
            "canonicalSourceRecords": len(source_records),
            "curatedSyntheticRecords": len(curated.records),
            "classifiableShippedFixtures": len(fixtures),
            "auditedRecords": len(audit_records),
        },
        "candidatePolicy": {
            "sameLabelCosineThreshold": args.same_label_threshold,
            "crossLabelCosineThreshold": args.cross_label_threshold,
            "nearestNeighborsPerRecord": args.neighbors,
            "automaticExclusion": False,
            "note": (
                "Embedding similarity proposes review candidates only. It cannot safely "
                "distinguish deliberate boundary pairs from duplicated labels."
            ),
        },
        "summary": {
            "candidates": len(candidate_values),
            "sameLabelCandidates": sum(item["sameLabel"] for item in candidate_values),
            "crossLabelCandidates": sum(
                not item["sameLabel"] for item in candidate_values
            ),
            "frozenBoundaryCandidates": leakage_candidates,
        },
        "humanReviewSample": {
            "status": "planned",
            "fraction": args.review_fraction,
            "seed": args.review_seed,
            "populationRecords": len(curated.records),
            "sampleRecords": len(sample),
            "strata": ["purpose", "slice", "primary language"],
            "generatedCSV": _relative(args.review_csv),
            "samplePurposeCounts": dict(
                sorted(Counter(record.value["purpose"] for record in sample).items())
            ),
            "sampleSliceCounts": dict(
                sorted(Counter(record.value["slice"] for record in sample).items())
            ),
            "requiredFields": [
                "reviewedPurpose",
                "reviewedSecondary",
                "reviewedDifficulty",
                "reviewedSlice",
                "reviewStatus",
            ],
        },
        "candidates": candidate_values,
    }
    write_json(args.report.resolve(), report)
    return report


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", action="append", type=Path)
    parser.add_argument("--fixtures", type=Path, default=DEFAULT_FIXTURES)
    parser.add_argument("--report", type=Path, default=DEFAULT_REPORT)
    parser.add_argument("--review-csv", type=Path, default=DEFAULT_REVIEW_CSV)
    parser.add_argument(
        "--curation-review", type=Path, default=DEFAULT_CURATION_REVIEW
    )
    parser.add_argument("--review-fraction", type=float, default=DEFAULT_SAMPLE_FRACTION)
    parser.add_argument("--review-seed", type=int, default=DEFAULT_SAMPLE_SEED)
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--model-revision", default=DEFAULT_MODEL_REVISION)
    parser.add_argument("--device", default="auto")
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--block-size", type=int, default=256)
    parser.add_argument("--neighbors", type=int, default=8)
    parser.add_argument("--same-label-threshold", type=float, default=0.97)
    parser.add_argument("--cross-label-threshold", type=float, default=0.985)
    parser.add_argument("--word-duplicate-threshold", type=float, default=0.92)
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    if args.batch_size <= 0 or args.block_size <= 0 or args.neighbors <= 0:
        parser.error("batch size, block size, and neighbors must be positive")
    try:
        report = audit(args)
    except (DataError, OSError, UnicodeError, ValueError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1
    print(
        f"Semantic audit: {report['population']['auditedRecords']} records, "
        f"{report['summary']['candidates']} candidates "
        f"({report['summary']['frozenBoundaryCandidates']} cross a frozen boundary); "
        f"{report['humanReviewSample']['sampleRecords']} prompts sampled for review."
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
