"""Synthetic cost-model and isolated router-gate replay. No provider calls.

Usage: python cost-probe.py /path/to/pinned/headroom-source
The source root must contain the four files in DIGESTS at PIN.
Package namespaces bypass SDK initializers; policy modules execute unchanged.
The router method and its log-bucket helper are compiled from their original
AST nodes, without importing the native compressor or changing their bodies.
"""
import ast
import hashlib
import importlib
import logging
import math
import os
from pathlib import Path
import platform
import sys
import types

PIN = "e995eb9f1dacac7ed548758406de6c1508caef95"
DIGESTS = {
    "headroom/proxy/auth_policy.py": "ebddc60db80e901275c17b1757835f7231bb13802d34e76fe027c0bb466f1b4a",
    "headroom/proxy/auth_mode.py": "7571bd6793efb451651413313338e4c7f3d6b4a91364396a4f4b5fb496f35b27",
    "headroom/transforms/compression_policy.py": "a51cfadbfe69bd0300cca1cdb504dc1b6222e8dc4941beb536c6a20606ebb790",
    "headroom/transforms/content_router.py": "d6f26cdcd2dcacc3f41a50494bb4dc11ad508dcb19a657416f1eb6772d4e9842",
}
root = Path(sys.argv[1]).resolve()
for name, digest in DIGESTS.items():
    assert hashlib.sha256((root / name).read_bytes()).hexdigest() == digest, name
for name in ("headroom", "headroom.proxy", "headroom.transforms"):
    assert name not in sys.modules, "Run this probe in a fresh Python process"
    package = types.ModuleType(name)
    package.__path__ = [str(root.joinpath(*name.split(".")))]
    sys.modules[name] = package
cp = importlib.import_module("headroom.transforms.compression_policy")
policy = cp.policy_default_payg()

router_path = root / "headroom/transforms/content_router.py"
tree = ast.parse(router_path.read_text())
router_class = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "ContentRouter")
method = next(n for n in router_class.body if isinstance(n, ast.FunctionDef) and n.name == "_net_cost_allows")
bucket = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "_gain_bucket")
namespace = {"math": math, "os": os, "logger": logging.getLogger("cost-probe")}
exec(compile(ast.Module(body=[bucket, method], type_ignores=[]), str(router_path), "exec"), namespace)
gate = namespace["_net_cost_allows"]
router = types.SimpleNamespace(_runtime_compression_policy=policy)

print(f"Python {platform.python_version()}; source {PIN}")
print("Synthetic model: gains are plain-input-token cost units, not measured bills")
cases = [
    ("small-warm", 2000, 50000, 10, 1, 300, -55500, False),
    ("small-cold", 2000, 50000, 10, 0, 300, 4300, True),
    ("large-warm", 50000, 10000, 10, 1, 300, 38500, True),
    ("near-5m", 2000, 1500, 10, 1, 300, 275, True),
    ("near-1h", 2000, 1500, 10, 1, 3600, -850, False),
]
for label, dt, suffix, reads, alive, ttl, expected_gain, expected_allow in cases:
    w = cp.cache_write_multiplier_for_ttl(ttl)
    gain = policy.net_mutation_gain(dt, suffix, reads, alive, write_multiplier=w)
    allow = policy.should_mutate_deep(dt, suffix, reads, alive, write_multiplier=w)
    assert math.isclose(gain, expected_gain, abs_tol=1e-8)
    assert allow is expected_allow
    if alive == 1:
        # Independent warm-cache reduction of the general expression.
        assert math.isclose(gain, dt * 0.1 * reads - (w - 0.1) * suffix, abs_tol=1e-8)
    print(f"{label}: gain={gain:.1f}; mutate={allow}")
for ttl, expected in ((300, 287.5), (3600, 475.0)):
    w = cp.cache_write_multiplier_for_ttl(ttl)
    boundary = policy.break_even_reads(2000, 50000, write_multiplier=w)
    assert math.isclose(boundary, expected)
    print(f"small-warm break-even, ttl={ttl}: R={boundary:.1f}")

previous = {key: os.environ.get(key) for key in (
    "HEADROOM_NET_COST_EXPECTED_READS", "HEADROOM_NET_COST_P_ALIVE"
)}
try:
    os.environ["HEADROOM_NET_COST_EXPECTED_READS"] = "10"
    os.environ["HEADROOM_NET_COST_P_ALIVE"] = "1"
    counts, markers = {}, []

    def call(slot, dt, suffix, state, **kwargs):
        return gate(router, slot_idx=slot, original_tokens=dt + 100,
                    compressed_tokens=100, suffix_tokens=[0] * (slot + 1) + [suffix],
                    route_counts=counts, transforms_applied=markers,
                    batch_state=state, **kwargs)

    # A rejected mutation must not open the batch floor.
    rejected = {"floor": None}
    assert not call(0, 2000, 50000, rejected)
    assert rejected["floor"] is None
    print("rejected edit: batch floor remains None")

    state = {"floor": None}
    assert call(0, 50000, 10000, state)
    assert state["floor"] == 0
    assert call(1, 2000, 50000, state)
    assert state["floor"] == 0
    assert counts["netcost_batch_admitted"] == 1
    assert "router:netcost_batch_admit" in markers
    print("batch: shallow edit admitted; deeper small edit admitted; floor=0")

    # Later slots cannot subsidize an earlier edit.
    earlier = {"floor": 1}
    assert not call(0, 2000, 50000, earlier)
    print("earlier edit after a deeper floor: rejected")

    assert call(0, 2000, 50000, {"floor": None}, p_alive_override=0.0)
    assert counts["netcost_idle_admitted"] == 1
    assert "router:netcost_idle_compaction" in markers
    print("idle override P_alive=0: admitted with idle-compaction marker")
finally:
    for key, value in previous.items():
        if value is None:
            os.environ.pop(key, None)
        else:
            os.environ[key] = value
print("PASS: 5 cost cases, 2 break-even boundaries, 4 isolated gate scenarios")