#!/usr/bin/env python3
from __future__ import annotations

import ast
import json
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
brain_path = ROOT / "ai" / "brain_server.py"
practice_path = ROOT / "services" / "design_academy_practice_service.py"
source = brain_path.read_text(encoding="utf-8")
practice_source = practice_path.read_text(encoding="utf-8")
tree = ast.parse(source, filename=str(brain_path))

wanted_assignments = {
    "ACADEMY_ROUND_JSON_SCHEMAS",
    "ACADEMY_ROUND_REQUIRED_KEYS",
    "ACADEMY_OUTPUT_CONTRACT",
    "ACADEMY_ROUND_SYSTEM_CONTEXTS",
    "_STRING",
    "_STRING_8",
    "_STRING_20",
    "_STRING_80",
    "_SAFETY_SCHEMA",
    "_PROVENANCE_SCHEMA",
}
wanted_classes = {"AcademyStructuredOutputError"}
wanted_functions = {
    "_academy_candidate_object",
    "_academy_validate_shape",
    "canonicalise_academy_answer",
}
selected = []
for node in tree.body:
    if isinstance(node, (ast.Import, ast.ImportFrom)):
        names = {alias.name.split(".")[0] for alias in node.names}
        if names & {"ast", "hashlib", "json", "re"}:
            selected.append(node)
    elif isinstance(node, ast.Assign):
        targets = {target.id for target in node.targets if isinstance(target, ast.Name)}
        if targets & wanted_assignments:
            selected.append(node)
    elif isinstance(node, ast.ClassDef) and node.name in wanted_classes:
        selected.append(node)
    elif isinstance(node, ast.FunctionDef) and node.name in wanted_functions:
        selected.append(node)

namespace: dict[str, object] = {}
exec(compile(ast.Module(body=selected, type_ignores=[]), str(brain_path), "exec"), namespace)
canonicalise = namespace["canonicalise_academy_answer"]
error_type = namespace["AcademyStructuredOutputError"]
schemas = namespace["ACADEMY_ROUND_JSON_SCHEMAS"]
contexts = namespace["ACADEMY_ROUND_SYSTEM_CONTEXTS"]

samples = {
    1: {
        "round": 1,
        "summary": "A repeatable parameter architecture defines named inputs, applies constraints consistently, and validates measurable geometry before any export or release decision.",
        "required_inputs": ["target dimensions"],
        "method": ["define named parameters", "apply deterministic constraints", "validate dimensions and solid integrity"],
        "acceptance_criteria": ["all dimensions match the parameters", "the result is one valid solid"],
        "safety_gate": {"risk_class": "low", "action": "continue", "forge_execution_allowed": False, "reason": "This is knowledge practice only and cannot unlock Forge execution."},
        "limitations": ["No physical test performed", "No production release authorised"],
    },
    2: {
        "round": 2,
        "failure_diagnosis": {"failure": "Missing acceptance criteria", "cause": "The workflow did not define measurable conditions for a valid result.", "correction": "Add explicit dimensional, topology, and rejection checks before release."},
        "provenance_review": {"licence_status": "unknown", "decision": "quarantine", "commercial_use": "blocked", "reason": "Unknown licensing cannot support safe commercial use or distribution."},
        "safety_gate": {"risk_class": "low", "action": "continue", "forge_execution_allowed": False, "reason": "This remains knowledge practice and no tool execution is authorised."},
        "limitations": ["No external source verified", "No physical result inspected"],
    },
    3: {
        "round": 3,
        "transfer_project": {"title": "Desk organiser", "goal": "Apply parameter architecture to a harmless organiser with measurable dimensions.", "inputs": ["width and depth"], "steps": ["define dimensions", "build constrained geometry", "validate outputs"], "outputs": ["editable model"], "acceptance_criteria": ["dimensions match inputs", "invalid values are rejected"]},
        "boundary_cases": {"minimum": "small valid dimensions", "nominal": "normal desk dimensions", "maximum": "largest allowed dimensions", "invalid": "negative dimension rejected"},
        "provenance_review": {"licence_status": "unknown", "decision": "quarantine", "commercial_use": "blocked", "reason": "Unknown licensing remains quarantined and cannot be used commercially."},
        "safety_gate": {"risk_class": "low", "action": "continue", "forge_execution_allowed": False, "reason": "The transfer remains knowledge-only until evidence and certification pass."},
        "limitations": ["No tool build performed", "No production readiness claimed"],
    },
}

for round_number, sample in samples.items():
    strict = canonicalise(json.dumps(sample), round_number)
    if strict["repaired"] is not False or strict["method"] != "strict_json":
        raise SystemExit(f"STOP: Round {round_number} strict JSON failed")
    if json.loads(strict["canonical"])["round"] != round_number:
        raise SystemExit(f"STOP: Round {round_number} canonical output mismatch")
    trailing = json.dumps(sample, separators=(",", ":"))[:-1] + ",}"
    repaired = canonicalise(trailing, round_number)
    if repaired["repaired"] is not True:
        raise SystemExit(f"STOP: Round {round_number} trailing comma not repaired")
    unsafe = json.loads(json.dumps(sample))
    unsafe["safety_gate"]["forge_execution_allowed"] = True
    try:
        canonicalise(json.dumps(unsafe), round_number)
    except error_type:
        pass
    else:
        raise SystemExit(f"STOP: Round {round_number} unsafe Forge flag accepted")

if set(schemas) != {1, 2, 3} or set(contexts) != {1, 2, 3}:
    raise SystemExit("STOP: Three round contracts were not installed")

required_source = (
    '"exam_round": int(round_number)',
    'payload.get("output_contract") != "academy_round_json_v3"',
    'payload.get("context_profile") != "academy_round_specific_v2"',
    'evaluate_response(skill_id, bundle, response, round_number)',
    'raw_model_answer.txt',
    'AcademyBrainResponseError',
    'ROUND_REQUIRED_RESPONSE_KEYS',
)
for marker in required_source:
    if marker not in practice_source:
        raise SystemExit(f"STOP: Missing practice marker: {marker}")

for marker in (
    'ACADEMY_OUTPUT_CONTRACT = "academy_round_json_v3"',
    '"round_specific_contracts": [1, 2, 3]',
    '"invalid_raw_answer_returned_on_502": True',
    'canonicalise_academy_answer(raw_answer, exam_round)',
    'format_override=ACADEMY_ROUND_JSON_SCHEMAS[exam_round]',
):
    if marker not in source:
        raise SystemExit(f"STOP: Missing brain marker: {marker}")

print(json.dumps({
    "ok": True,
    "output_contract": "academy_round_json_v3",
    "round_specific_contracts": [1, 2, 3],
    "max_tokens": 420,
    "num_ctx": 1536,
    "raw_invalid_answer_audited": True,
    "missing_content_invented": False,
    "unsafe_forge_flag_blocked": True,
    "normal_ask_preserved": 'self.path not in {"/ask", "/academy/practice/ask"}' in source,
}, indent=2))