#!/usr/bin/env python3
import json
import subprocess
import unittest
from pathlib import Path
from tempfile import TemporaryDirectory

from context_packet import PacketError, generate_packet
from validate_model_budget import validate_request


class ContextPacketTests(unittest.TestCase):
    def base(self, candidates):
        return {
            "task": "bounded task",
            "header": {"milestone": "M1"},
            "baseline": {"source_sha": "abc"},
            "scope": {"include": ["feature/1"]},
            "contracts": ["protocol.v1"],
            "risks": [],
            "next": ["run test"],
            "forbidden": ["real hardware"],
            "candidates": candidates,
        }

    def test_filters_by_evidence_bytes_before_final_write(self):
        packet = generate_packet(
            self.base([
                {"path": "small.log", "size": 6, "priority": 2},
                {"path": "large.log", "size": 7, "priority": 1},
            ]),
            context_id="c1",
            packet_id="p1",
            total_bytes=4096,
            evidence_bytes=10,
        )
        self.assertEqual([item["path"] for item in packet["evidence"]["selected"]], ["small.log"])
        self.assertEqual(packet["evidence"]["excluded"][0]["reason_code"], "E4")
        self.assertLessEqual(packet["budget"]["packet_bytes"], 4096)

    def test_filters_by_remaining_context_space(self):
        packet = generate_packet(
            self.base([{"path": "evidence.log", "size": 1, "priority": 1}]),
            context_id="c1",
            packet_id="p1",
            total_bytes=550,
            evidence_bytes=100,
        )
        self.assertEqual(packet["evidence"]["selected"], [])
        self.assertEqual(packet["evidence"]["excluded"][0]["reason_code"], "E5")
        self.assertLessEqual(packet["budget"]["packet_bytes"], 550)

    def test_evidence_payload_counts_against_total_context_budget(self):
        packet = generate_packet(
            self.base([{"path": "payload.bin", "size": 40, "priority": 1}]),
            context_id="c1",
            packet_id="p1",
            total_bytes=570,
            evidence_bytes=100,
        )
        self.assertEqual(packet["evidence"]["selected"], [])
        self.assertEqual(packet["evidence"]["excluded"][0]["reason_code"], "E5")
        self.assertLessEqual(
            packet["budget"]["packet_bytes"]
            + packet["budget"]["selected_evidence_bytes"],
            570,
        )
        self.assertEqual(
            packet["budget"]["packet_bytes"],
            len(json.dumps(packet, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")),
        )

    def test_skeleton_over_budget_fails_before_selection(self):
        with self.assertRaisesRegex(PacketError, "context skeleton exceeds byte budget"):
            generate_packet(
                self.base([]),
                context_id="c1",
                packet_id="p1",
                total_bytes=100,
                evidence_bytes=100,
            )

    def test_required_item_fails_without_writing_an_over_budget_packet(self):
        with self.assertRaisesRegex(PacketError, "required evidence cannot fit"):
            generate_packet(
                self.base([{"path": "required.log", "size": 101, "required": True}]),
                context_id="c1",
                packet_id="p1",
                total_bytes=4096,
                evidence_bytes=100,
            )

    def test_required_item_fails_when_payload_consumes_remaining_context(self):
        with self.assertRaisesRegex(PacketError, r"required evidence cannot fit \(E5\)"):
            generate_packet(
                self.base([{"path": "required.log", "size": 40, "required": True}]),
                context_id="c1",
                packet_id="p1",
                total_bytes=570,
                evidence_bytes=100,
            )

    def test_selection_is_deterministic_and_count_is_enforced(self):
        packet = generate_packet(
            self.base([
                {"path": "z.log", "size": 1},
                {"path": "a.log", "size": 1},
                {"path": "required.log", "size": 1, "required": True},
            ]),
            context_id="c1",
            packet_id="p1",
            total_bytes=4096,
            evidence_bytes=100,
            max_files=2,
        )
        self.assertEqual([item["path"] for item in packet["evidence"]["selected"]], ["required.log", "a.log"])
        self.assertEqual(packet["evidence"]["excluded"][0]["reason_code"], "E3")

    def test_token_budget_is_enforced_during_generation(self):
        packet = generate_packet(
            self.base([{"path": "large.log", "size": 1800, "priority": 1}]),
            context_id="c1",
            packet_id="p1",
            profile="doc",
            total_bytes=8192,
            evidence_bytes=2048,
        )
        self.assertEqual(packet["evidence"]["selected"], [])
        self.assertEqual(packet["evidence"]["excluded"][0]["reason_code"], "E6")

    def test_required_item_fails_on_token_budget(self):
        with self.assertRaisesRegex(PacketError, r"required evidence cannot fit \(E6\)"):
            generate_packet(
                self.base([{"path": "required.log", "size": 1800, "required": True}]),
                context_id="c1",
                packet_id="p1",
                profile="doc",
                total_bytes=8192,
                evidence_bytes=2048,
            )

    def test_overrides_cannot_expand_profile(self):
        with self.assertRaisesRegex(PacketError, "bytes budget override exceeds profile"):
            generate_packet(
                self.base([]),
                context_id="c1",
                packet_id="p1",
                profile="doc",
                total_bytes=8193,
            )

    def test_token_limit_override_can_only_shrink_profile(self):
        packet = generate_packet(
            self.base([]),
            context_id="c1",
            packet_id="p1",
            profile="feature",
            token_limit=512,
        )
        self.assertEqual(packet["budget"]["token_limit"], 512)
        with self.assertRaisesRegex(PacketError, "tokens budget override exceeds profile"):
            generate_packet(
                self.base([]),
                context_id="c1",
                packet_id="p1",
                profile="feature",
                token_limit=4097,
            )

    def test_generated_packet_passes_model_preflight(self):
        packet = generate_packet(
            self.base([]),
            context_id="c1",
            packet_id="p1",
            profile="feature",
        )
        result = validate_request(packet, context_window=16384, overhead_tokens=256)
        self.assertLessEqual(result["reserved_total_tokens"], 16384)

    def test_cli_does_not_replace_existing_output_on_generation_failure(self):
        with TemporaryDirectory() as directory:
            root = Path(directory)
            input_path = root / "input.json"
            output_path = root / "packet.json"
            input_path.write_text(
                json.dumps(self.base([{"path": "required.log", "size": 101, "required": True}]))
            )
            output_path.write_text("previous packet")
            result = subprocess.run(
                [
                    "python3",
                    "scripts/context_packet.py",
                    "--input",
                    str(input_path),
                    "--output",
                    str(output_path),
                    "--context-id",
                    "c1",
                    "--packet-id",
                    "p1",
                    "--total-bytes",
                    "4096",
                    "--evidence-bytes",
                    "100",
                ],
                check=False,
                capture_output=True,
                text=True,
            )
            self.assertEqual(result.returncode, 2)
            self.assertEqual(output_path.read_text(), "previous packet")


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