#!/usr/bin/env python3
"""Materialize a fixed-load, greedy subset of an immutable replay window.

The output preserves prompt/arrival/sampling provenance while setting every
selected request's temperature to zero.  It is intentionally a controlled
mechanism workload, not a replacement for a production sampling trace.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
from pathlib import Path
from typing import Any


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def resolve_trace_path(windows_path: Path, trace_file: str) -> Path:
    candidate = Path(trace_file)
    if candidate.is_absolute():
        return candidate
    direct = (windows_path.parent / candidate).resolve()
    if direct.exists():
        return direct
    if candidate.parts and candidate.parts[0] == "trace_windows":
        return (windows_path.parent / Path(*candidate.parts[1:])).resolve()
    return direct


def numeric(value: Any, *, field: str, row_index: int) -> float:
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise SystemExit(f"row {row_index}: {field} must be numeric")
    value = float(value)
    if not math.isfinite(value):
        raise SystemExit(f"row {row_index}: {field} must be finite")
    return value


def stable_session_key(value: Any) -> str:
    """Render a session identifier without conflating JSON types."""
    return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--windows-path", type=Path, required=True)
    parser.add_argument("--window-id", required=True)
    parser.add_argument("--sampling-u-threshold", type=float, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--output-window-id", required=True)
    args = parser.parse_args()

    if not args.windows_path.is_file():
        raise SystemExit(f"windows file does not exist: {args.windows_path}")
    threshold = float(args.sampling_u_threshold)
    if not math.isfinite(threshold) or not 0.0 <= threshold <= 1.0:
        raise SystemExit("--sampling-u-threshold must be finite and in [0, 1]")

    source_windows = json.loads(args.windows_path.read_text(encoding="utf-8"))
    windows = source_windows.get("windows") if isinstance(source_windows, dict) else source_windows
    if not isinstance(windows, list):
        raise SystemExit("windows payload must contain a list")
    source_window = next(
        (
            item
            for item in windows
            if isinstance(item, dict) and item.get("window_id") == args.window_id
        ),
        None,
    )
    if source_window is None:
        raise SystemExit(f"window id not found: {args.window_id}")
    trace_file = source_window.get("trace_file")
    if not isinstance(trace_file, str) or not trace_file:
        raise SystemExit(f"window {args.window_id} has no trace_file")
    source_trace = resolve_trace_path(args.windows_path, trace_file)
    if not source_trace.is_file():
        raise SystemExit(f"source trace does not exist: {source_trace}")
    source_sampling_strategy = source_window.get("sampling_strategy")
    full_session_source = source_sampling_strategy == "session_coherent_uniform_score"
    window_session_source = source_sampling_strategy == "window_session_closed_uniform_score"
    session_closed_source = full_session_source or window_session_source

    args.output_dir.mkdir(parents=True, exist_ok=True)
    output_trace = args.output_dir / f"{args.output_window_id}.jsonl"
    selected = 0
    sampling_values: list[float] = []
    timestamps: list[float] = []
    session_counts: dict[str, dict[str, int]] = {}
    with source_trace.open("r", encoding="utf-8") as src, output_trace.open(
        "w", encoding="utf-8"
    ) as dest:
        for row_index, line in enumerate(src):
            if not line.strip():
                continue
            row = json.loads(line)
            if not isinstance(row, dict):
                raise SystemExit(f"row {row_index}: expected an object")
            sampling_u = numeric(row.get("sampling_u"), field="sampling_u", row_index=row_index)
            session_entry: dict[str, int] | None = None
            if session_closed_source:
                if "session_root" not in row:
                    raise SystemExit(
                        f"row {row_index}: session-closed source has no session_root"
                    )
                session_key = stable_session_key(row["session_root"])
                session_entry = session_counts.setdefault(
                    session_key, {"total_turns": 0, "selected_turns": 0}
                )
                session_entry["total_turns"] += 1
            if sampling_u > threshold:
                continue
            timestamp = numeric(row.get("timestamp"), field="timestamp", row_index=row_index)
            row["temperature"] = 0.0
            dest.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")
            selected += 1
            sampling_values.append(sampling_u)
            timestamps.append(timestamp)
            if session_entry is not None:
                session_entry["selected_turns"] += 1
    if selected == 0:
        output_trace.unlink(missing_ok=True)
        raise SystemExit("threshold selected zero requests")
    split_sessions = [
        entry
        for entry in session_counts.values()
        if 0 < entry["selected_turns"] < entry["total_turns"]
    ]
    if split_sessions:
        output_trace.unlink(missing_ok=True)
        raise SystemExit(
            "session-coherent source produced a partial selected session; "
            f"split_session_count={len(split_sessions)}"
        )
    selected_session_turn_counts = [
        entry["selected_turns"]
        for entry in session_counts.values()
        if entry["selected_turns"] > 0
    ]

    output_window = dict(source_window)
    output_window.update(
        {
            "window_id": args.output_window_id,
            "trace_file": str(output_trace.resolve()),
            "num_requests": selected,
            "sampling_strategy": (
                "session_coherent_uniform_score_then_greedy_temperature"
                if full_session_source
                else (
                    "window_session_closed_uniform_score_then_greedy_temperature"
                    if window_session_source
                    else "fixed_uniform_score_then_greedy_temperature"
                )
            ),
            "collectivespec_source_window_id": args.window_id,
            "collectivespec_sampling_u_threshold": threshold,
            "collectivespec_temperature_override": 0.0,
            "collectivespec_source_trace_path": str(source_trace),
            "collectivespec_source_trace_sha256": sha256(source_trace),
            "collectivespec_source_sampling_strategy": source_sampling_strategy,
            "collectivespec_session_closed_within_window": session_closed_source,
            "collectivespec_session_closure_scope": (
                "full_session" if full_session_source else "window_only" if window_session_source else None
            ),
        }
    )
    output_windows = {
        "kind": "collectivespec_controlled_greedy_window",
        "source_windows_path": str(args.windows_path.resolve()),
        "source_window_id": args.window_id,
        "sampling_u_threshold": threshold,
        "temperature_override": 0.0,
        "windows": [output_window],
    }
    output_manifest = args.output_dir / "windows.json"
    output_manifest.write_text(
        json.dumps(output_windows, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    sanity = {
        "n": selected,
        "sampling_u": {
            "min": min(sampling_values),
            "max": max(sampling_values),
            "distinct_value_count": len(set(sampling_values)),
            "all_at_or_below_threshold": all(value <= threshold for value in sampling_values),
        },
        "timestamp": {
            "min": min(timestamps),
            "max": max(timestamps),
            "non_negative": all(value >= 0 for value in timestamps),
            "nondecreasing": all(
                timestamps[index] >= timestamps[index - 1]
                for index in range(1, len(timestamps))
            ),
        },
        "temperature": {"distinct_value_count": 1, "all_zero": True},
        "session": {
            "source_is_session_coherent": full_session_source,
            "source_is_window_session_closed": window_session_source,
            "session_root_count": len(session_counts) if session_closed_source else None,
            "selected_session_count": (
                len(selected_session_turn_counts) if session_closed_source else None
            ),
            "split_session_count": len(split_sessions) if session_closed_source else None,
            "all_selected_or_dropped": not split_sessions if session_closed_source else None,
            "selected_turns": {
                "min": min(selected_session_turn_counts)
                if selected_session_turn_counts
                else None,
                "max": max(selected_session_turn_counts)
                if selected_session_turn_counts
                else None,
                "distinct_value_count": len(set(selected_session_turn_counts)),
            },
        },
        "source_trace_sha256": sha256(source_trace),
        "output_trace_sha256": sha256(output_trace),
    }
    print(
        json.dumps(
            {
                "windows_path": str(output_manifest),
                "trace_path": str(output_trace),
                "data_sanity": sanity,
            },
            ensure_ascii=False,
            indent=2,
            sort_keys=True,
        )
    )
    return 0


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