#!/usr/bin/env python3
"""Original LockedIn Labs synthetic trace-projection lab, Python 3.11+.
Run: python3 fde-agent-trace-lab.py test
     python3 fde-agent-trace-lab.py demo
No exporter, provider, network, credentials or real payloads. This is a local
application event format, not an OpenTelemetry SDK or compliance mechanism.
"""
from dataclasses import dataclass
import argparse
import json
import re
import socket
import unittest
from unittest.mock import patch


class TraceRefusal(ValueError):
    """Only fixed policy codes may reach the diagnostic boundary."""


@dataclass(frozen=True)
class HostContext:
    tenant: str
    run: str


# These schemas are deliberately closed, small and application-owned.
SCHEMAS = {
    "request_admitted": {"source_version": (int, 1, 1_000_000)},
    "tool_finished": {"tool": (str, {"lookup_procedure"}),
                      "outcome": (str, {"returned", "withheld", "failed"}),
                      "elapsed_ms": (int, 0, 60_000)},
    "model_claimed": {"claim": (str, {"ready_for_review"})},
    "operation_failed": {"error_class": (str, {"timeout", "unavailable", "invalid_result"})},
}


def project_event(kind, fields, context, max_bytes=256):
    """Call at the host boundary before any logging/export; never log fields raw.
    Trusted host code chooses kind/context. A model cannot call this API directly.
    Synthetic context constraints are teaching fixtures, not authentication.
    """
    if (type(context) is not HostContext or context.tenant not in ("north", "south")
            or type(context.run) is not str or not re.fullmatch(r"R[0-9]{1,6}", context.run)):
        raise TraceRefusal("invalid_context")
    if type(kind) is not str or kind not in SCHEMAS:
        raise TraceRefusal("unknown_event")
    if type(fields) is not dict or any(type(key) is not str for key in fields):
        raise TraceRefusal("invalid_fields")
    schema = SCHEMAS[kind]
    if set(fields) != set(schema):
        raise TraceRefusal("unsupported_fields")
    safe = {}
    for key, rule in schema.items():
        value = fields[key]
        if type(value) is not rule[0]:  # bool is not a measured integer.
            raise TraceRefusal("invalid_value")
        if rule[0] is str and value not in rule[1]:
            raise TraceRefusal("invalid_value")
        if rule[0] is int and not rule[1] <= value <= rule[2]:
            raise TraceRefusal("invalid_value")
        safe[key] = value
    event = {"schema_version": 1, "tenant": context.tenant, "run": context.run,
             "event": kind, "evidence": "model_proposal" if kind == "model_claimed" else "application_observation",
             "attributes": safe}
    if type(max_bytes) is not int or not 1 <= max_bytes <= 512:
        raise TraceRefusal("invalid_budget")
    if len(json.dumps(event, separators=(",", ":")).encode("utf-8")) > max_bytes:
        raise TraceRefusal("event_too_large")
    return event


def demo_report():
    context = HostContext("north", "R1")
    cases = [
        ("T1", "request_admitted", {"source_version": 3}),
        ("T2", "tool_finished", {"tool": "lookup_procedure", "outcome": "returned", "elapsed_ms": 12}),
        ("T3", "model_claimed", {"claim": "ready_for_review"}),
        ("T4", "request_admitted", {"source_version": 3, "source_text": "SYNTHETIC_SOURCE_BODY"}),
        ("T5", "request_admitted", {"source_version": 3, "tenant": "south"}),
        ("T6", "operation_failed", {"error_class": "timeout", "message": "SYNTHETIC person@example.invalid"}),
        ("T7", "tool_finished", {"tool": "lookup_procedure", "outcome": "returned", "elapsed_ms": True}),
        ("T8", "tool_finished", {"tool": "SYNTHETIC_UNBOUNDED_" * 100, "outcome": "returned", "elapsed_ms": 12}),
    ]
    reports = []
    for case_id, kind, fields in cases:
        try:
            event = project_event(kind, fields, context)
            reports.append({"case": case_id, "status": "projected", "event": event})
        except TraceRefusal as error:
            reports.append({"case": case_id, "status": "refused", "code": str(error)})
    return {"cases": reports,
            "limits": "Eight visible synthetic inputs. Local event projection only; no inference, exporter, retention service, access control, deletion, sampling or production logging is tested."}


class TraceTests(unittest.TestCase):
    def setUp(self):
        self.context = HostContext("north", "R1")

    def test_useful_diagnostics_survive_without_raw_payloads(self):
        fields = {"tool": "lookup_procedure", "outcome": "returned", "elapsed_ms": 12}
        event = project_event("tool_finished", fields, self.context)
        self.assertEqual(event["attributes"], fields)
        self.assertEqual((event["tenant"], event["run"], event["evidence"]), ("north", "R1", "application_observation"))
        self.assertLessEqual(len(json.dumps(event, separators=(",", ":")).encode()), 256)
        fields["elapsed_ms"] = 99
        self.assertEqual(event["attributes"]["elapsed_ms"], 12)

    def test_a_model_claim_is_not_application_acceptance(self):
        event = project_event("model_claimed", {"claim": "ready_for_review"}, self.context)
        self.assertEqual(event["evidence"], "model_proposal")
        with self.assertRaisesRegex(TraceRefusal, "invalid_value"):
            project_event("model_claimed", {"claim": "human_approved"}, self.context)

    def test_unknown_nested_fields_and_context_overrides_are_not_exported(self):
        for key in ("prompt", "response", "arguments", "source_text", "tenant", "run"):
            with self.subTest(key=key), self.assertRaisesRegex(TraceRefusal, "unsupported_fields"):
                project_event("request_admitted", {"source_version": 3, key: {"body": "SYNTHETIC_PRIVATE"}}, self.context)

    def test_types_ranges_and_enum_values_are_closed(self):
        for value in (True, -1, 60_001, "12", [12]):
            with self.subTest(value=value), self.assertRaisesRegex(TraceRefusal, "invalid_value"):
                project_event("tool_finished", {"tool": "lookup_procedure", "outcome": "returned", "elapsed_ms": value}, self.context)
        for value in ("SYNTHETIC_UNBOUNDED_" * 100, "send_message"):
            with self.assertRaisesRegex(TraceRefusal, "invalid_value"):
                project_event("tool_finished", {"tool": value, "outcome": "returned", "elapsed_ms": 12}, self.context)

    def test_failure_diagnostics_never_need_an_exception_message(self):
        event = project_event("operation_failed", {"error_class": "timeout"}, self.context)
        self.assertEqual(event["attributes"], {"error_class": "timeout"})
        with self.assertRaisesRegex(TraceRefusal, "unsupported_fields"):
            project_event("operation_failed", {"error_class": "timeout", "message": "SYNTHETIC_SECRET"}, self.context)

    def test_context_and_event_names_cannot_carry_free_text(self):
        for context in (HostContext("person@example.invalid", "R1"), HostContext("north", "SYNTHETIC_PRIVATE"), {}):
            with self.assertRaisesRegex(TraceRefusal, "invalid_context"):
                project_event("request_admitted", {"source_version": 3}, context)
        with self.assertRaisesRegex(TraceRefusal, "unknown_event"):
            project_event("SYNTHETIC_PRIVATE", {}, self.context)
        for fields in ([], {1: "text"}, {"source_version": 3, "extra": None}, {}):
            with self.assertRaises(TraceRefusal):
                project_event("request_admitted", fields, self.context)

    def test_serialized_event_budget_is_enforced_before_release(self):
        with self.assertRaisesRegex(TraceRefusal, "event_too_large"):
            project_event("request_admitted", {"source_version": 3}, self.context, max_bytes=1)
        for budget in (True, 0, 513):
            with self.assertRaisesRegex(TraceRefusal, "invalid_budget"):
                project_event("request_admitted", {"source_version": 3}, self.context, max_bytes=budget)

    def test_whole_demo_retains_every_case_without_echoing_rejected_input(self):
        report = demo_report()
        self.assertEqual([case["case"] for case in report["cases"]], [f"T{i}" for i in range(1, 9)])
        self.assertEqual([case["status"] for case in report["cases"]], ["projected"] * 3 + ["refused"] * 5)
        output = json.dumps(report)
        for forbidden in ("SYNTHETIC_SOURCE_BODY", "person@example.invalid", "SYNTHETIC_UNBOUNDED_"):
            self.assertNotIn(forbidden, output)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=["test", "demo"])
    args = parser.parse_args()
    with patch.object(socket.socket, "connect", side_effect=AssertionError("network_forbidden")):
        if args.command == "test":
            result = unittest.TextTestRunner(verbosity=2).run(unittest.defaultTestLoader.loadTestsFromTestCase(TraceTests))
            raise SystemExit(0 if result.wasSuccessful() else 1)
        print(json.dumps(demo_report(), indent=2))
