#!/usr/bin/env python3

from __future__ import annotations

import json
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path


ROOT = Path(__file__).resolve().parent


class CodeTracePreflightTest(unittest.TestCase):
    def test_audit_and_materialize_512_block_window(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            source = root / "051315-051317.jsonl"
            with source.open("w") as stream:
                for index in range(4501):
                    input_tokens = 513 if index % 2 else 512
                    stream.write(
                        json.dumps(
                            {
                                "chat_id": index,
                                "parent_chat_id": index - 1 if index % 2 else -1,
                                "timestamp": float(index),
                                "input_length": input_tokens,
                                "output_length": 32,
                                "hash_ids": [index // 2]
                                if input_tokens == 512
                                else [index // 2, 100000 + index],
                            }
                        )
                        + "\n"
                    )
            audit = root / "audit.json"
            subprocess.run(
                [
                    sys.executable,
                    str(ROOT / "audit_code_trace.py"),
                    "--source",
                    str(source),
                    "--output",
                    str(audit),
                ],
                check=True,
            )
            payload = json.loads(audit.read_text())
            self.assertEqual(payload["data_gate"], "PASS")
            self.assertEqual(
                payload["selected"]["hash_contract"]["exact_source_block_size"], 512
            )
            self.assertEqual(payload["max_model_len_recommendation"], 40960)
            output = root / "window"
            subprocess.run(
                [
                    sys.executable,
                    str(ROOT / "prepare_code_window.py"),
                    "--audit",
                    str(audit),
                    "--output-root",
                    str(output),
                ],
                check=True,
            )
            manifest = json.loads((output / "window-manifest.json").read_text())
            rows = [
                json.loads(line)
                for line in (output / "code-raw-window.jsonl").read_text().splitlines()
            ]
        self.assertEqual(manifest["requests"], len(rows))
        self.assertGreaterEqual(manifest["duration_s"], 3600)
        self.assertEqual(rows[0]["sampling_u"], rows[1]["sampling_u"])


if __name__ == "__main__":
    unittest.main()