#!/usr/bin/env python3
"""Audit telemetry responses over every uncensored replay decile.

This corrective analysis keeps the frozen P1 pairs and thresholds, but replaces
the absolute 5/10-second cutoff with cumulative and non-overlapping 10%-of-trace
windows.  It deliberately reports every common decile instead of selecting the
best-looking horizon.
"""

from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
import math
from pathlib import Path
from statistics import median
from typing import Any, Iterable, Mapping


HERE = Path(__file__).resolve().parent
P1_PATH = HERE.parent / "intervention-response-v0" / "analyze_p1.py"
SCHEMA = "intervention-response-phase-aware-existing-v2"
DECILE_FRACTION = 0.1
MAX_DECILES = 10


def _load_p1():
    spec = importlib.util.spec_from_file_location("intervention_response_p1", P1_PATH)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


P1 = _load_p1()


def numeric(values: Iterable[float | int]) -> dict[str, Any]:
    finite = [float(value) for value in values]
    result = P1.V0.numeric(finite)
    result["median"] = median(finite)
    return result


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def trial_directories(run_root: Path) -> list[Path]:
    result = []
    for cell in sorted((run_root / "cells").iterdir()):
        if not cell.is_dir():
            continue
        for candidate in sorted(cell.iterdir()):
            if candidate.is_dir() and P1.RUN_PATTERN.match(candidate.name):
                result.append(candidate)
    if not result:
        raise ValueError("P1 run root contains no measured trial directories")
    return result


def load_metadata(run_root: Path) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    metadata = []
    streams = []
    for cell in sorted((run_root / "cells").iterdir()):
        if not cell.is_dir():
            continue
        stream_paths = sorted((cell / "opprof").glob("*.jsonl"))
        if len(stream_paths) != 1:
            raise ValueError(f"{cell}: expected one Layer-1 stream")
        stream_path = stream_paths[0]
        streams.append(
            {
                "cell": cell.name,
                "path": str(stream_path.resolve()),
                "sha256": sha256_file(stream_path),
                "bytes": stream_path.stat().st_size,
            }
        )
    for run_dir in trial_directories(run_root):
        match = P1.RUN_PATTERN.match(run_dir.name)
        assert match is not None
        level, replicate_text = match.groups()
        result_path = run_dir / "result.json"
        requests_path = run_dir / "requests.jsonl"
        result = json.loads(result_path.read_text(encoding="utf-8"))
        selected = int(result["selection"]["count"])
        offered = float(result["selection"]["offered_req_s"])
        if selected <= 0 or offered <= 0.0:
            raise ValueError(f"{result_path}: invalid selected count or offered rate")
        metadata.append(
            {
                "trial_id": str(result_path.relative_to(run_root)),
                "cell": str(result["cell"]),
                "tp": int(result["tp"]),
                "mns": int(result["mns"]),
                "level": level,
                "replicate": int(replicate_text),
                "elapsed_s": float(result["interval"]["elapsed_s"]),
                "trace_duration_s": round(selected / offered, 9),
                "early_stopped": bool(result["early_stopped"]),
                "request_count": selected,
                "result_sha256": sha256_file(result_path),
                "requests_sha256": sha256_file(requests_path),
            }
        )
    return metadata, streams


def common_decile_fractions(
    *, trace_duration_s: float, minimum_elapsed_s: float
) -> tuple[float, ...]:
    if trace_duration_s <= 0.0 or minimum_elapsed_s <= 0.0:
        raise ValueError("trace duration and elapsed time must be positive")
    supported = min(
        MAX_DECILES,
        int(math.floor((minimum_elapsed_s / trace_duration_s) * 10.0 + 1e-12)),
    )
    return tuple(
        round(index * DECILE_FRACTION, 10) for index in range(1, supported + 1)
    )


def _trial_record(
    *,
    run_root: Path,
    run_dir: Path,
    result: Mapping[str, Any],
    state: dict[str, float],
    outcome: dict[str, float],
) -> dict[str, Any]:
    match = P1.RUN_PATTERN.match(run_dir.name)
    assert match is not None
    level, replicate_text = match.groups()
    result_path = run_dir / "result.json"
    requests_path = run_dir / "requests.jsonl"
    return {
        "trial_id": str(result_path.relative_to(run_root)),
        "cell": str(result["cell"]),
        "tp": int(result["tp"]),
        "mns": int(result["mns"]),
        "level": level,
        "replicate": int(replicate_text),
        "offered_rate_per_gpu": float(
            result["selection"]["offered_req_s_per_gpu"]
        ),
        "request_hash": str(result["selection"]["request_id_order_sha256"]),
        "request_count": int(result["selection"]["count"]),
        "result_sha256": sha256_file(result_path),
        "requests_sha256": sha256_file(requests_path),
        "full_pass_rate": float(result["pass_rate"]),
        "full_feasible": bool(result["feasible"]),
        "early_stopped": bool(result["early_stopped"]),
        "state": state,
        "outcome": outcome,
    }


def load_interval_trials(
    run_root: Path,
    intervals_s: tuple[tuple[float, float], ...],
) -> tuple[dict[tuple[float, float], list[dict[str, Any]]], list[dict[str, Any]]]:
    by_interval = {interval: [] for interval in intervals_s}
    stream_provenance = []
    for cell in sorted((run_root / "cells").iterdir()):
        if not cell.is_dir():
            continue
        stream_paths = sorted((cell / "opprof").glob("*.jsonl"))
        if len(stream_paths) != 1:
            raise ValueError(f"{cell}: expected one Layer-1 stream")
        stream_path = stream_paths[0]
        stream = P1.load_jsonl(stream_path)
        stream_provenance.append(
            {
                "cell": cell.name,
                "path": str(stream_path.resolve()),
                "sha256": sha256_file(stream_path),
                "bytes": stream_path.stat().st_size,
            }
        )
        for run_dir in sorted(cell.iterdir()):
            if not run_dir.is_dir() or P1.RUN_PATTERN.match(run_dir.name) is None:
                continue
            result_path = run_dir / "result.json"
            requests_path = run_dir / "requests.jsonl"
            result = json.loads(result_path.read_text(encoding="utf-8"))
            requests = P1.load_jsonl(requests_path)
            start_ns = int(result["interval"]["start_mono_ns"])
            elapsed_s = float(result["interval"]["elapsed_s"])
            for interval in intervals_s:
                start_s, end_s = interval
                if start_s < 0.0 or end_s <= start_s:
                    raise ValueError(f"invalid analysis interval: {interval}")
                if elapsed_s + 1e-9 < end_s:
                    raise ValueError(
                        f"{result_path}: elapsed {elapsed_s} shorter than {end_s}s"
                    )
                state = P1.V0.flatten_state(
                    P1.summarize_engine(
                        stream,
                        start_ns=start_ns + int(start_s * 1e9),
                        end_ns=start_ns + int(end_s * 1e9),
                        request_count=int(result["selection"]["count"]),
                    )
                )
                outcome = P1._prefix_outcome(result, requests, end_s)
                by_interval[interval].append(
                    _trial_record(
                        run_root=run_root,
                        run_dir=run_dir,
                        result=result,
                        state=state,
                        outcome=outcome,
                    )
                )
    return by_interval, stream_provenance


def coverage(trials: list[dict[str, Any]]) -> dict[str, Any]:
    admitted = [float(trial["outcome"]["admitted_fraction"]) for trial in trials]
    completed = [
        float(trial["outcome"]["admitted_fraction"])
        * float(trial["outcome"]["completed_over_admitted"])
        for trial in trials
    ]
    return {
        "admitted_fraction_of_total": numeric(admitted),
        "completed_fraction_of_total": numeric(completed),
    }


def slim_window_analysis(
    trials: list[dict[str, Any]], *, start_s: float, end_s: float, fraction: float
) -> dict[str, Any]:
    analysis = P1.analyze_horizon(trials, end_s)
    return {
        "start_s": start_s,
        "end_s": end_s,
        "end_fraction": fraction,
        "coverage_at_end": coverage(trials),
        "action_pairs": len(analysis["actions"]),
        "repeat_pairs": len(analysis["repeats"]),
        "response_statistics": analysis["response_statistics"],
        "qualifying_response_features": analysis["qualifying_response_features"],
        "efficacy": analysis["efficacy"],
        "sanity": analysis["sanity"],
    }


def _pearson(left: list[float], right: list[float]) -> float | None:
    if len(left) != len(right) or not left:
        raise ValueError("Pearson inputs must be non-empty and have equal length")
    left_mean = sum(left) / len(left)
    right_mean = sum(right) / len(right)
    numerator = sum(
        (x - left_mean) * (y - right_mean)
        for x, y in zip(left, right, strict=True)
    )
    left_ss = sum((x - left_mean) ** 2 for x in left)
    right_ss = sum((y - right_mean) ** 2 for y in right)
    if left_ss == 0.0 or right_ss == 0.0:
        return None
    return numerator / math.sqrt(left_ss * right_ss)


def trajectory_summary(
    block_trials: list[tuple[tuple[float, float], list[dict[str, Any]]]]
) -> dict[str, Any]:
    if not block_trials:
        raise ValueError("trajectory requires at least one block")
    identities = []
    states_by_block = []
    for interval, trials in block_trials:
        ordered = sorted(
            trials,
            key=lambda trial: (trial["cell"], trial["level"], trial["replicate"]),
        )
        current_identities = [
            (trial["cell"], trial["level"], trial["replicate"]) for trial in ordered
        ]
        if identities and current_identities != identities:
            raise ValueError("trajectory blocks do not contain identical trials")
        identities = current_identities
        states_by_block.append((interval, [trial["state"] for trial in ordered]))

    features = {}
    for feature in P1.V0.ALL_FEATURES:
        block_values = [
            [float(state[feature]) for state in states]
            for _interval, states in states_by_block
        ]
        first = block_values[0]
        last = block_values[-1]
        delta = [right - left for left, right in zip(first, last, strict=True)]
        features[feature] = {
            "block_medians": [median(values) for values in block_values],
            "first_to_last_delta": numeric(delta),
            "first_to_last_abs_delta": numeric(abs(value) for value in delta),
            "first_to_last_pearson": _pearson(first, last),
            "changed_trials": sum(abs(value) > 1e-12 for value in delta),
        }
    return {
        "trial_count": len(identities),
        "blocks": [
            {"start_s": interval[0], "end_s": interval[1]}
            for interval, _states in states_by_block
        ],
        "features": features,
    }


def audit(*, run_root: Path, manifest_path: Path, output_path: Path) -> dict[str, Any]:
    metadata, metadata_streams = load_metadata(run_root)
    durations = [float(item["trace_duration_s"]) for item in metadata]
    elapsed = [float(item["elapsed_s"]) for item in metadata]
    duration = median(durations)
    deciles = common_decile_fractions(
        trace_duration_s=duration, minimum_elapsed_s=min(elapsed)
    )
    if not deciles:
        raise ValueError("no complete replay decile is shared by all trials")
    cumulative_intervals = tuple(
        (0.0, round(duration * fraction, 9)) for fraction in deciles
    )
    block_intervals = tuple(
        (
            round(duration * (fraction - DECILE_FRACTION), 9),
            round(duration * fraction, 9),
        )
        for fraction in deciles
    )
    all_intervals = tuple(dict.fromkeys([*cumulative_intervals, *block_intervals]))
    trials_by_interval, streams = load_interval_trials(run_root, all_intervals)
    manifest_validation = P1.validate_manifest(
        trials_by_interval[cumulative_intervals[0]], manifest_path
    )

    cumulative = []
    blocks = []
    for fraction, cumulative_interval, block_interval in zip(
        deciles, cumulative_intervals, block_intervals, strict=True
    ):
        cumulative.append(
            slim_window_analysis(
                trials_by_interval[cumulative_interval],
                start_s=cumulative_interval[0],
                end_s=cumulative_interval[1],
                fraction=fraction,
            )
        )
        blocks.append(
            slim_window_analysis(
                trials_by_interval[block_interval],
                start_s=block_interval[0],
                end_s=block_interval[1],
                fraction=fraction,
            )
        )

    invariants = {
        "expected_trial_count": len(metadata) == 36,
        "trace_duration_consistent": max(durations) - min(durations) <= 1e-9,
        "all_intervals_uncensored": all(
            item["elapsed_s"] + 1e-9 >= cumulative_intervals[-1][1]
            for item in metadata
        ),
        "stream_provenance_consistent": metadata_streams == streams,
        "manifest_trials_match": (
            manifest_validation["expected_trials"]
            == manifest_validation["matched_trials"]
            == len(metadata)
        ),
        "all_window_sanity_pass": all(
            not item["sanity"]["red_flags"] for item in [*cumulative, *blocks]
        ),
    }
    red_flags = [name for name, passed in invariants.items() if not passed]
    complete_full_trajectory = min(elapsed) + 1e-9 >= duration
    if red_flags:
        decision = "STOP_DATA_INVALID"
    elif not complete_full_trajectory:
        decision = "REQUIRES_UNCENSORED_PHASE_AWARE_PILOT"
    else:
        decision = "FULL_TRAJECTORY_AVAILABLE"

    payload = {
        "schema": SCHEMA,
        "status": "COMPLETE",
        "decision": decision,
        "claim_boundary": (
            "Post-hoc corrective audit over every common replay decile. It can "
            "diagnose horizon sensitivity but cannot establish a held-out tuning claim."
        ),
        "design": {
            "decile_fraction": DECILE_FRACTION,
            "available_deciles": list(deciles),
            "trace_duration_s": duration,
            "maximum_common_end_s": cumulative_intervals[-1][1],
            "maximum_common_fraction": deciles[-1],
            "select_best_horizon": False,
            "cumulative_and_nonoverlapping_blocks": True,
        },
        "cumulative": cumulative,
        "blocks": blocks,
        "trajectory": trajectory_summary(
            [(interval, trials_by_interval[interval]) for interval in block_intervals]
        ),
        "provenance": {
            "analysis_script": str(Path(__file__).resolve()),
            "analysis_script_sha256": sha256_file(Path(__file__).resolve()),
            "p1_analysis_script": str(P1_PATH.resolve()),
            "p1_analysis_script_sha256": sha256_file(P1_PATH),
            "run_root": str(run_root.resolve()),
            "manifest": str(manifest_path.resolve()),
            "manifest_sha256": sha256_file(manifest_path),
            "manifest_validation": manifest_validation,
            "streams": streams,
            "trial_inputs": metadata,
        },
        "sanity": {
            "trials": len(metadata),
            "elapsed_s": numeric(elapsed),
            "trace_duration_s": numeric(durations),
            "early_stopped": sum(bool(item["early_stopped"]) for item in metadata),
            "request_count": numeric(item["request_count"] for item in metadata),
            "stream_bytes": numeric(item["bytes"] for item in streams),
            "invariants": invariants,
            "red_flags": red_flags,
        },
    }
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    return payload


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--run-root", type=Path, required=True)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    payload = audit(
        run_root=args.run_root,
        manifest_path=args.manifest,
        output_path=args.output,
    )
    print(
        json.dumps(
            {
                "decision": payload["decision"],
                "design": payload["design"],
                "sanity": payload["sanity"],
            },
            indent=2,
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()