#!/usr/bin/env python3
"""Extract a prompt-bearing routing fixture while emitting a prompt-free manifest."""

from __future__ import annotations

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


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--trace", type=Path, required=True)
    parser.add_argument("--row-ids", type=int, nargs="+", required=True)
    parser.add_argument("--parent-of", type=int, nargs="*", default=[])
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--manifest", type=Path, required=True)
    return parser.parse_args()


def common_prefix(left: list[Any], right: list[Any]) -> int:
    count = 0
    for lhs, rhs in zip(left, right):
        if lhs != rhs:
            break
        count += 1
    return count


def main() -> None:
    args = parse_args()
    target_ids = set(args.row_ids)
    parent_targets = set(args.parent_of)
    if not parent_targets.issubset(target_ids):
        raise SystemExit("--parent-of must be a subset of --row-ids")

    trace_digest = hashlib.sha256()
    by_row: dict[int, dict[str, Any]] = {}
    by_chat_id: dict[str, dict[str, Any]] = {}
    with args.trace.open("rb") as handle:
        row_id = 0
        while True:
            offset = handle.tell()
            line = handle.readline()
            if not line:
                break
            trace_digest.update(line)
            row = json.loads(line)
            meta = {
                "row_id": row_id,
                "offset": offset,
                "chat_id": str(row.get("chat_id")),
                "parent_chat_id": str(row.get("parent_chat_id")),
                "turn": int(row.get("turn", 1)),
                "input_length": int(row["input_length"]),
                "output_length": int(row["output_length"]),
                "hash_ids": row.get("hash_ids") or [],
            }
            by_chat_id[meta["chat_id"]] = meta
            if row_id in target_ids:
                by_row[row_id] = meta
            row_id += 1

    missing = target_ids - set(by_row)
    if missing:
        raise SystemExit(f"missing target row IDs: {sorted(missing)}")
    parent_rows: list[dict[str, Any]] = []
    for row_id in args.parent_of:
        parent_id = by_row[row_id]["parent_chat_id"]
        parent = by_chat_id.get(parent_id)
        if parent is None:
            raise SystemExit(f"row {row_id} parent chat {parent_id} is absent")
        parent_rows.append(parent)

    # Put parents first so online prefix caching can materialize their shared
    # blocks before descendants are admitted.
    ordered_meta = parent_rows + [by_row[row_id] for row_id in args.row_ids]
    if len({row["row_id"] for row in ordered_meta}) != len(ordered_meta):
        raise SystemExit("fixture rows must be unique")

    output_rows: list[dict[str, Any]] = []
    with args.trace.open("rb") as handle:
        for fixture_index, meta in enumerate(ordered_meta):
            handle.seek(meta["offset"])
            source = json.loads(handle.readline())
            prompt = source.get("prompt")
            if not isinstance(prompt, str) or not prompt:
                raise SystemExit(f"row {meta['row_id']} has no prompt")
            output_rows.append(
                {
                    "fixture_index": fixture_index,
                    "row_id": str(meta["row_id"]),
                    "prompt": prompt,
                    "prompt_sha256": hashlib.sha256(prompt.encode()).hexdigest(),
                    "input_length": meta["input_length"],
                    "output_length": meta["output_length"],
                    "turn": meta["turn"],
                    "chat_id": meta["chat_id"],
                    "parent_chat_id": meta["parent_chat_id"],
                    "hash_ids": meta["hash_ids"],
                }
            )

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w") as handle:
        for row in output_rows:
            handle.write(json.dumps(row, sort_keys=True) + "\n")

    pair_coverage = []
    by_chat = {row["chat_id"]: row for row in output_rows}
    for child in output_rows:
        parent = by_chat.get(child["parent_chat_id"])
        if parent is None:
            continue
        pair_coverage.append(
            {
                "parent_row_id": parent["row_id"],
                "child_row_id": child["row_id"],
                "parent_turn": parent["turn"],
                "child_turn": child["turn"],
                "trace_hash_common_prefix_blocks": common_prefix(
                    parent["hash_ids"], child["hash_ids"]
                ),
            }
        )
    fixture_digest = hashlib.sha256(args.output.read_bytes()).hexdigest()
    manifest = {
        "schema_version": "qwen30_trace_routing_fixture.v1",
        "source_trace": str(args.trace.resolve()),
        "source_trace_sha256": trace_digest.hexdigest(),
        "fixture_path": str(args.output.resolve()),
        "fixture_sha256": fixture_digest,
        "contains_prompt_text": False,
        "request_count": len(output_rows),
        "rows": [
            {
                key: row[key]
                for key in (
                    "fixture_index",
                    "row_id",
                    "prompt_sha256",
                    "input_length",
                    "output_length",
                    "turn",
                    "chat_id",
                    "parent_chat_id",
                )
            }
            | {"trace_hash_blocks": len(row["hash_ids"])}
            for row in output_rows
        ],
        "prefix_pairs": pair_coverage,
    }
    args.manifest.parent.mkdir(parents=True, exist_ok=True)
    args.manifest.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    print(json.dumps({"requests": len(output_rows), "prefix_pairs": pair_coverage}))


if __name__ == "__main__":
    main()