"""LockedIn FDE · Private inference capacity planning · Python 3.11+ · Edition 1.

Single-device, dense grouped-query attention teaching model. Synthetic bytes and
pre-counted tokens only: no model, GPU, packages, tokenizer, network or file writes.
An estimate is neither measured capacity nor permission to admit production work.
Run: python3 fde-private-inference-lab.py test | demo | <case> [--candidate prompt-only]
"""
from copy import deepcopy
import argparse
import json
import sys
import unittest

GIB = 1024 ** 3
CASES = ("baseline", "output-pressure", "context-overflow", "block-rounding",
         "weights-only-quantization", "half-width-cache", "multiple-devices")
EXPECTED = ("fits_estimate", "memory_budget_exceeded", "context_limit_exceeded",
            "fits_estimate", "fits_estimate", "fits_estimate", "unsupported_topology")


def fixture(name):
    if name not in CASES:
        raise ValueError("unknown-case")
    profile = dict(device_bytes=12 * GIB, weights_bytes=6 * GIB, reserve_bytes=2 * GIB,
                   layers=32, query_heads=64, kv_heads=8, head_dim=128,
                   cache_element_bytes=2, context_limit=16384, block_tokens=16, devices=1)
    requests = [dict(prompt_tokens=4096, max_output_tokens=1024) for _ in range(4)]
    if name == "output-pressure":
        for request in requests:
            request["max_output_tokens"] = 8192
    if name == "context-overflow":
        requests = [dict(prompt_tokens=15360, max_output_tokens=2048)]
    if name == "block-rounding":
        for request in requests:
            request["prompt_tokens"] += 1
    if name in ("weights-only-quantization", "half-width-cache"):
        requests += [deepcopy(requests[0]) for _ in range(3)]
        if name == "weights-only-quantization":
            profile["weights_bytes"] = 3 * GIB
        else:
            profile["cache_element_bytes"] = 1
    if name == "multiple-devices":
        profile["devices"] = 2
    return profile, requests


def estimate(profile, requests, candidate="reference"):
    if candidate not in ("reference", "prompt-only"):
        raise ValueError("unknown-candidate")
    keys = set(fixture("baseline")[0])
    if not isinstance(profile, dict) or set(profile) != keys:
        raise ValueError("invalid-profile")
    if any(type(value) is not int or value < 1 for value in profile.values()):
        raise ValueError("invalid-profile-scalar")
    if profile["cache_element_bytes"] not in (1, 2, 4):
        raise ValueError("unsupported-cache-width")
    if profile["query_heads"] % profile["kv_heads"]:
        raise ValueError("incompatible-head-counts")
    if not isinstance(requests, list) or not 1 <= len(requests) <= 64:
        raise ValueError("invalid-request-batch")
    for request in requests:
        if not isinstance(request, dict) or set(request) != {"prompt_tokens", "max_output_tokens"}:
            raise ValueError("invalid-request")
        if any(type(value) is not int or value < 1 for value in request.values()):
            raise ValueError("invalid-request-scalar")
    boundary = dict(model_executed=False, capacity_measured=False, production_admission=False)
    if profile["devices"] != 1:
        return dict(status="unsupported_topology", **boundary)
    # Full-attention dense GQA: K and V for every layer, KV head, head element.
    per_token = 2 * profile["layers"] * profile["kv_heads"] * profile["head_dim"] * profile["cache_element_bytes"]
    headroom = profile["device_bytes"] - profile["weights_bytes"] - profile["reserve_bytes"]
    if headroom <= 0:
        return dict(status="no_cache_headroom", kv_bytes_per_token=per_token, **boundary)
    block = profile["block_tokens"]
    raw = [request["prompt_tokens"] + request["max_output_tokens"] for request in requests]
    # The deliberately broken candidate retains the full context check but drops
    # output capacity from its memory accounting. It does not change the fixture.
    counted = raw if candidate == "reference" else [request["prompt_tokens"] for request in requests]
    rounded = [((tokens + block - 1) // block) * block for tokens in counted]
    required_tokens = sum(rounded)
    cache_blocks = headroom // (per_token * block)
    result = dict(kv_bytes_per_token=per_token, cache_headroom_bytes=headroom,
                  token_capacity=cache_blocks * block, required_tokens=required_tokens,
                  reserved_tokens_per_request=rounded, required_cache_bytes=required_tokens * per_token,
                  request_count=len(requests), **boundary)
    if any(tokens > profile["context_limit"] for tokens in raw):
        return dict(status="context_limit_exceeded", **result)
    return dict(status="fits_estimate" if required_tokens <= result["token_capacity"] else "memory_budget_exceeded", **result)


def observe(name, candidate="reference"):
    profile, requests = fixture(name)
    result = estimate(profile, requests, candidate)
    return dict(case=name, candidate=candidate, expected_status=EXPECTED[CASES.index(name)],
                matches_expected=result["status"] == EXPECTED[CASES.index(name)], profile=profile,
                requests=requests, **result)


class CapacityTests(unittest.TestCase):
    def test_gqa_geometry_and_binary_units(self):
        row = observe("baseline")
        self.assertEqual(row["kv_bytes_per_token"], 131072)
        self.assertEqual(row["cache_headroom_bytes"], 4294967296)
        self.assertEqual(row["token_capacity"], 32768)
        self.assertEqual(row["required_tokens"], 20480)
        self.assertEqual(row["required_cache_bytes"], 2684354560)

    def test_generated_output_changes_the_capacity_decision(self):
        row = observe("output-pressure")
        self.assertEqual(row["required_tokens"], 49152)
        self.assertEqual(row["status"], "memory_budget_exceeded")
        broken = observe("output-pressure", "prompt-only")
        self.assertEqual(broken["status"], "fits_estimate")
        self.assertFalse(broken["matches_expected"])

    def test_context_is_prompt_plus_output_before_block_rounding(self):
        self.assertEqual(observe("context-overflow")["status"], "context_limit_exceeded")
        profile, _ = fixture("baseline")
        profile["context_limit"] = 5121
        self.assertEqual(estimate(profile, [dict(prompt_tokens=4097, max_output_tokens=1024)])["status"], "fits_estimate")

    def test_each_request_rounds_up_separately(self):
        row = observe("block-rounding")
        self.assertEqual(row["reserved_tokens_per_request"], [5136] * 4)
        self.assertEqual(row["required_tokens"], 20544)

    def test_cache_capacity_rounds_down_to_whole_blocks(self):
        profile, requests = fixture("baseline")
        profile["device_bytes"] += 15 * 131072
        self.assertEqual(estimate(profile, requests)["token_capacity"], 32768)

    def test_weight_and_cache_width_are_independent(self):
        weights = observe("weights-only-quantization")
        cache = observe("half-width-cache")
        self.assertEqual(weights["kv_bytes_per_token"], 131072)
        self.assertEqual(weights["token_capacity"], 57344)
        self.assertEqual(cache["kv_bytes_per_token"], 65536)
        self.assertEqual(cache["token_capacity"], 65536)
        self.assertEqual(weights["required_tokens"], cache["required_tokens"])

    def test_no_automatic_pooling_or_negative_headroom(self):
        self.assertEqual(observe("multiple-devices")["status"], "unsupported_topology")
        profile, requests = fixture("baseline")
        profile["weights_bytes"] = 11 * GIB
        self.assertEqual(estimate(profile, requests)["status"], "no_cache_headroom")

    def test_rejects_boolean_float_extra_and_missing_contract_fields(self):
        for bad in (True, 3.5, -1, "12"):
            profile, requests = fixture("baseline")
            profile["kv_heads"] = bad
            with self.assertRaises(ValueError):
                estimate(profile, requests)
        profile, requests = fixture("baseline")
        requests[0]["max_output_tokens"] = False
        with self.assertRaises(ValueError):
            estimate(profile, requests)
        for key in ("tenant", "gpu_memory_utilization"):
            profile, requests = fixture("baseline")
            profile[key] = 1
            with self.assertRaises(ValueError):
                estimate(profile, requests)
        profile, requests = fixture("baseline")
        del requests[0]["max_output_tokens"]
        with self.assertRaises(ValueError):
            estimate(profile, requests)

    def test_useful_estimate_never_claims_execution_or_deployment(self):
        for name in CASES:
            row = observe(name)
            self.assertTrue(row["matches_expected"])
            self.assertFalse(row["model_executed"])
            self.assertFalse(row["capacity_measured"])
            self.assertFalse(row["production_admission"])


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("case", choices=("test", "demo", *CASES))
    parser.add_argument("--candidate", choices=("reference", "prompt-only"), default="reference")
    args = parser.parse_args()
    if args.case == "test" and args.candidate != "reference":
        parser.error("candidate selection applies to demo or individual cases; tests retain the reference contract")
    if args.case == "test":
        suite = unittest.defaultTestLoader.loadTestsFromTestCase(CapacityTests)
        return 0 if unittest.TextTestRunner(verbosity=2).run(suite).wasSuccessful() else 1
    if args.case == "demo":
        rows = [observe(name, args.candidate) for name in CASES]
        complete = all(row["matches_expected"] for row in rows)
        print(json.dumps(dict(comparison_complete=complete, observations=rows), indent=2))
        return 0 if complete else 2
    row = observe(args.case, args.candidate)
    print(json.dumps(row, indent=2))
    return 0 if row["matches_expected"] and row["status"] == "fits_estimate" else 2


if __name__ == "__main__":
    sys.exit(main())