#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""G4.3 canary — STRICT reproduction surface composite-ID counter (Task 4.3).

Detects composite IDs (`C##` OR `C##O##`) within ±50 chars window of any of 9
reproduction-surface keywords. Surface keywords are imported from
`render_prompt_card._ID_REPRODUCTION_SURFACES`.

Override R1R2-B3 carry: char-level ±50 window — NOT sentence boundary regex
(LLM punctuation is variable, sentence split would yield false negatives).
`feedback_no_regex_postprocessing.md` carry.

Override R4-M5 carry: dedup — if a single id occurrence (id_token + id_position
start) is within ±50 chars of multiple surface keywords, count violation
1× per (id_token, id_position[0]) pair.

Scope: ALL shots of pinned (si, shi) tuples (R3-I1).

Exit criteria (analyzer, OUT-OF-BAND):
  candidate.metrics.reproduction_surface_composite_id_count == 0  (STRICT)
  baseline value is informational only.

Usage:
    python scripts/canary/g4_3_reproduction_surface.py \\
        --config canary_config.json \\
        --cp-root data/projects/<pid>/checkpoints/scene_detail \\
        --prompt-version 17.<timestamp> \\
        --role candidate \\
        --output results/reproduction_surface_candidate.json
"""
from __future__ import annotations

import argparse
import json
import re
import sys
from datetime import datetime
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from _g4_3_common import (  # noqa: E402
    _load_scenes_or_fail,
    _validate_scene_preflight,
    build_pinning_scene_set,
    is_close_framing,
    load_cp_manifest,
    load_pinning_from_args,
    render_prompt_card_imports,
    validate_pinning,
)

_CONSTS = render_prompt_card_imports()
_ID_REPRODUCTION_SURFACES = _CONSTS["_ID_REPRODUCTION_SURFACES"]

# Match bare C## OR composite C##O## — spec §5.1 covers both because either
# form near a reproduction surface is a violation (R1R2-B3 / R3-I1).
_ID_REGEX = re.compile(r"\bC\d{2}(?:O\d{2})?\b")
WINDOW_CHARS = 50


def iso8601_now() -> str:
    return datetime.utcnow().isoformat() + "Z"


def find_reproduction_violations(prompt: str) -> list[dict]:
    """Return list of violation records.

    R4-M5 dedup: a single id occurrence (keyed by `(id_token, id_position[0])`)
    is counted at most once even if it falls within ±50 of multiple surface
    keywords. The chosen surface is the FIRST one whose window overlaps the
    id; remaining matches are recorded inside the violation under
    `additional_surfaces` (advisory) but contribute 0 to the counted total.

    char-level ±50 window proximity check (NO sentence regex — R1R2-B3 carry):
      id and surface are considered "near" iff the id span [id_start, id_end]
      and the surface span [s_start, s_end] are within `WINDOW_CHARS` of each
      other on the character axis, i.e.
        id_start - s_end <= WINDOW_CHARS  AND  s_start - id_end <= WINDOW_CHARS
    """
    surface_positions: list[tuple[int, int, str]] = []
    pl = prompt.lower()
    for surface in _ID_REPRODUCTION_SURFACES:
        s_lower = surface.lower()
        slen = len(s_lower)
        start = 0
        while True:
            idx = pl.find(s_lower, start)
            if idx == -1:
                break
            surface_positions.append((idx, idx + slen, surface))
            start = idx + 1

    violations: list[dict] = []
    seen_keys: set[tuple[str, int]] = set()
    for match in _ID_REGEX.finditer(prompt):
        id_start = match.start()
        id_end = match.end()
        id_token = match.group()
        key = (id_token, id_start)

        # Find ALL surfaces within ±50 chars window.
        nearby: list[tuple[int, int, str]] = []
        for s_start, s_end, surface in surface_positions:
            # char-level proximity: id and surface spans within WINDOW_CHARS on
            # both sides. Equivalent to "the gap between the two spans is
            # <= WINDOW_CHARS, regardless of which precedes which".
            gap = max(id_start - s_end, s_start - id_end)
            if gap <= WINDOW_CHARS:
                nearby.append((s_start, s_end, surface))

        if not nearby:
            continue
        if key in seen_keys:
            # R4-M5 dedup: this (id_token, id_position[0]) already counted.
            continue
        seen_keys.add(key)

        primary_s_start, primary_s_end, primary_surface = nearby[0]
        violations.append({
            "id_token": id_token,
            "surface": primary_surface,
            "id_position": [id_start, id_end],
            "surface_position": [primary_s_start, primary_s_end],
            "additional_surfaces": [s[2] for s in nearby[1:]],
        })

    return violations


def main() -> int:
    parser = argparse.ArgumentParser(
        description="G4.3 canary: reproduction-surface composite-ID violation counter (±50 chars)."
    )
    parser.add_argument("--config", type=str, default=None)
    parser.add_argument("--pid", type=str, default=None)
    parser.add_argument("--cp-root", type=str, default=None)
    parser.add_argument("--prompt-version", type=str, required=True)
    parser.add_argument("--role", type=str, default="candidate",
                        choices=["baseline", "candidate"])
    parser.add_argument("--scene-index-list", type=str, default=None)
    parser.add_argument("--shot-index-list-per-scene", type=str, default=None)
    parser.add_argument("--model-routing", type=str, default=None)
    parser.add_argument("--prompt-source-mode", type=str, default=None,
                        choices=[None, "file", "db"])
    parser.add_argument("--card-commit-hash", type=str, default=None)
    parser.add_argument("--id-policy-card-snapshot-hash", type=str, default=None)
    parser.add_argument("--output", type=str, required=True)
    args = parser.parse_args()

    pinning = load_pinning_from_args(args)
    validate_pinning(pinning)

    cp_root_str: str | None = None
    if args.config:
        with Path(args.config).open("r", encoding="utf-8") as fp:
            cfg = json.load(fp)
        cp_root_str = cfg.get("cp_root")
    if not cp_root_str:
        cp_root_str = args.cp_root
    if not cp_root_str:
        print("[g4_3_reproduction_surface] ERROR: --cp-root required",
              file=sys.stderr)
        return 1
    cp_root = Path(cp_root_str)
    if not cp_root.exists():
        print(f"[g4_3_reproduction_surface] ERROR: --cp-root not found: {cp_root}",
              file=sys.stderr)
        return 1

    cp = load_cp_manifest(cp_root)
    measurement_failures: list[str] = []

    scenes = _load_scenes_or_fail(cp, measurement_failures)
    if scenes is None:
        out = {
            "timestamp": iso8601_now(),
            "prompt_version": args.prompt_version,
            "role": args.role,
            "pinning": pinning,
            "scene_set": [],
            "metrics": {
                "total_shots": 0,
                "shots_with_reproduction_surface": 0,
                "reproduction_surface_composite_id_count": 0,
                "reproduction_surface_composite_id_per_shot": {},
                "violations_detail": [],
            },
            "measurement_failures": measurement_failures,
            "measurement_scripts": {
                "reproduction_surface": (
                    "scripts/canary/g4_3_reproduction_surface.py "
                    "(all shots, ±50 chars window — R1R2-B3 / R3-I1 / R4-M5 dedup)"
                )
            },
        }
        _emit(out, args)
        return 1

    pinning_set = build_pinning_scene_set(pinning)
    seen_in_cp: set[tuple[int, int]] = set()

    total_shots = 0
    shots_with_surface = 0
    violations_total = 0
    violations_per_shot: dict[str, int] = {}
    violations_detail: list[dict] = []
    scene_set_out: list[dict] = []

    for scene in scenes:
        if not isinstance(scene, dict):
            measurement_failures.append("scene entry is not dict")
            continue
        si = scene.get("scene_index")
        shi = scene.get("_shot_index")
        if si is None or shi is None:
            continue
        if (int(si), int(shi)) not in pinning_set:
            continue
        seen_in_cp.add((int(si), int(shi)))

        variations = _validate_scene_preflight(scene, si, shi, measurement_failures)
        if variations is None:
            continue

        cf_flag, cf_err = is_close_framing(scene)
        if cf_err is not None:
            measurement_failures.append(f"s{si}_sh{shi}: {cf_err}")

        # Diagnostic flags.
        has_composite_id = False
        has_reproduction_surface = False
        for var in variations:
            if not isinstance(var, dict):
                continue
            p = var.get("t2i_prompt")
            if isinstance(p, str):
                if _ID_REGEX.search(p):
                    has_composite_id = True
                pl = p.lower()
                for surf in _ID_REPRODUCTION_SURFACES:
                    if surf.lower() in pl:
                        has_reproduction_surface = True
                        break

        scene_set_out.append({
            "scene_index": si,
            "shot_index": shi,
            "is_close_framing": cf_flag,
            "has_composite_id": has_composite_id,
            "has_reproduction_surface": has_reproduction_surface,
        })

        total_shots += 1
        if has_reproduction_surface:
            shots_with_surface += 1

        shot_count = 0
        for vi, var in enumerate(variations):
            if not isinstance(var, dict):
                measurement_failures.append(
                    f"s{si}_sh{shi}_v{vi}: variation not a dict"
                )
                continue
            prompt = var.get("t2i_prompt")
            if not isinstance(prompt, str) or not prompt.strip():
                measurement_failures.append(
                    f"s{si}_sh{shi}_v{vi}: t2i_prompt missing, empty, or not str"
                )
                continue
            v_violations = find_reproduction_violations(prompt)
            shot_count += len(v_violations)
            for v in v_violations:
                violations_detail.append({
                    "scene_index": si,
                    "shot_index": shi,
                    "variation_index": vi,
                    **v,
                })
        violations_total += shot_count
        violations_per_shot[f"{si}_{shi}"] = shot_count

    for (si, shi) in pinning_set:
        if (si, shi) not in seen_in_cp:
            measurement_failures.append(
                f"pinning (s{si}_sh{shi}): not present in CP data.scenes"
            )

    out = {
        "timestamp": iso8601_now(),
        "prompt_version": args.prompt_version,
        "role": args.role,
        "pinning": pinning,
        "scene_set": scene_set_out,
        "metrics": {
            "total_shots": total_shots,
            "shots_with_reproduction_surface": shots_with_surface,
            "reproduction_surface_composite_id_count": violations_total,
            "reproduction_surface_composite_id_per_shot": violations_per_shot,
            "violations_detail": violations_detail,
        },
        "measurement_failures": measurement_failures,
        "measurement_scripts": {
            "reproduction_surface": (
                "scripts/canary/g4_3_reproduction_surface.py "
                "(all shots, ±50 chars window — R1R2-B3 / R3-I1 / R4-M5 dedup)"
            )
        },
    }

    _emit(out, args)

    print(
        f"[g4_3_reproduction_surface] role={args.role} prompt_version={args.prompt_version}\n"
        f"  total_shots={total_shots}  shots_with_surface={shots_with_surface}\n"
        f"  reproduction_surface_composite_id_count={violations_total}\n"
        f"  measurement_failures: {len(measurement_failures)}\n"
        f"Exit criteria (STRICT): candidate.metrics.reproduction_surface_composite_id_count == 0",
        file=sys.stderr,
    )

    if measurement_failures:
        return 1
    if args.role == "candidate" and violations_total > 0:
        return 1
    return 0


def _emit(out: dict, args: argparse.Namespace) -> None:
    if args.output == "-":
        json.dump(out, sys.stdout, ensure_ascii=False, indent=2, sort_keys=True)
        sys.stdout.write("\n")
    else:
        out_path = Path(args.output)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        with out_path.open("w", encoding="utf-8") as fp:
            json.dump(out, fp, ensure_ascii=False, indent=2, sort_keys=True)


if __name__ == "__main__":
    sys.exit(main())