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

import ast
import hashlib
import json
import socket
import sys
import urllib.error
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent

if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from services.design_academy_practice_service import (
    AcademyBrainResponseError,
    practice_service,
)

BRAIN = ROOT / "ai" / "brain_server.py"
SERVICE = ROOT / "services" / "design_academy_practice_service.py"

brain_source = BRAIN.read_text(encoding="utf-8")
service_source = SERVICE.read_text(encoding="utf-8")
compile(brain_source, str(BRAIN), "exec")
compile(service_source, str(SERVICE), "exec")


def function_namespace(source: str, names: set[str], namespace: dict) -> dict:
    tree = ast.parse(source)
    selected = [
        node for node in tree.body
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
        and node.name in names
    ]
    module = ast.Module(body=selected, type_ignores=[])
    ast.fix_missing_locations(module)
    exec(compile(module, "<selected-functions>", "exec"), namespace)
    return namespace

brain_ns = function_namespace(
    brain_source,
    {
        "_academy_timeout_for_round",
        "_academy_timeout_error",
        "_academy_timeout_payload",
    },
    {
        "ACADEMY_PRACTICE_TIMEOUT": 220,
        "ACADEMY_ROUND3_TIMEOUT": 360,
        "ACADEMY_OUTPUT_CONTRACT": "academy_round_json_v6",
        "TimeoutError": TimeoutError,
        "socket": socket,
        "urllib": __import__("urllib"),
        "hashlib": hashlib,
    },
)

assert brain_ns["_academy_timeout_for_round"](1) == 220
assert brain_ns["_academy_timeout_for_round"](2) == 220
assert brain_ns["_academy_timeout_for_round"](3) == 360
assert brain_ns["_academy_timeout_error"](TimeoutError("timed out")) is True
assert brain_ns["_academy_timeout_error"](socket.timeout("timed out")) is True
assert brain_ns["_academy_timeout_error"](
    urllib.error.URLError(TimeoutError("timed out"))
) is True
assert brain_ns["_academy_timeout_error"](RuntimeError("different failure")) is False

payload = brain_ns["_academy_timeout_payload"](
    3,
    360,
    raw_answer="",
    model_retry_count=0,
)
assert payload == {
    "ok": False,
    "error": "Academy model request timed out",
    "retryable": True,
    "infrastructure_failure": True,
    "output_contract": "academy_round_json_v6",
    "exam_round": 3,
    "timeout_seconds": 360,
    "timeout_stage": "ollama_generation",
    "raw_answer": "",
    "raw_answer_sha256": hashlib.sha256(b"").hexdigest(),
    "raw_answer_audited": True,
    "model_retry_count": 0,
    "model_retry_reason": "",
    "initial_raw_answer": "",
    "initial_raw_answer_sha256": "",
}

service_ns = function_namespace(
    service_source,
    {"_academy_client_timeout_for_round"},
    {
        "ACADEMY_MICRO_EXAM_CLIENT_TIMEOUT": 300,
        "ACADEMY_ROUND3_CLIENT_TIMEOUT": 420,
    },
)
assert service_ns["_academy_client_timeout_for_round"](1) == 300
assert service_ns["_academy_client_timeout_for_round"](2) == 300
assert service_ns["_academy_client_timeout_for_round"](3) == 420

timeout_failure = AcademyBrainResponseError(
    "Brain API HTTP 504: Academy model request timed out",
    {
        "retryable": True,
        "infrastructure_failure": True,
        "timeout_stage": "ollama_generation",
    },
)

assert (
    practice_service._transient_brain_failure(
        timeout_failure
    )
    is True
)

contract_failure = AcademyBrainResponseError(
    "Brain API HTTP 422: Academy structured output invalid",
    {
        "retryable": False,
        "infrastructure_failure": False,
        "content_format_failure": True,
        "failure_class": "academy_output_contract_failure",
    },
)

assert (
    practice_service._transient_brain_failure(
        contract_failure
    )
    is False
)

required_brain_tokens = [
    'ACADEMY_OUTPUT_CONTRACT = "academy_round_json_v6"',
    'os.environ.get("SHIRE_ACADEMY_ROUND3_TIMEOUT", "360")',
    '"round3_timeout_seconds": ACADEMY_ROUND3_TIMEOUT',
    '"round3_structured_timeout_response": True',
    'status=504',
    '"schema_name": f"academy_round_{exam_round}_json_v6"',
    'SHIRE ACADEMY ROUND 3: TRANSFER AND MEASURABLE ACCEPTANCE.',
]
for token in required_brain_tokens:
    assert token in brain_source, token

required_service_tokens = [
    'ACADEMY_ROUND3_CLIENT_TIMEOUT = 420',
    'timeout=_academy_client_timeout_for_round(round_number)',
    '"practice_round3_client_timeout_seconds": ACADEMY_ROUND3_CLIENT_TIMEOUT',
    'payload.get("infrastructure_failure") is True',
    'return payload.get("retryable") is True',
    '"brain api practice request failed"',
    'if payload.get("output_contract") != "academy_round_json_v6"',
    'if payload.get("context_profile") != "academy_round_specific_v5"',
]
for token in required_service_tokens:
    assert token in service_source, token

# Ensure Round 3 content contract was not weakened.
for token in (
    '"transfer_project"',
    '"boundary_cases"',
    '"provenance_review"',
    '"safety_gate"',
    '"limitations"',
    '"forge_execution_allowed": {"type": "boolean", "const": False}',
):
    assert token in brain_source, token

print(json.dumps({
    "ok": True,
    "output_contract": "academy_round_json_v6",
    "round3_server_timeout_seconds": 360,
    "round3_client_timeout_seconds": 420,
    "round1_and_round2_server_timeout_seconds": 220,
    "round1_and_round2_client_timeout_seconds": 300,
    "structured_timeout_http_status": 504,
    "structured_timeout_retryable": True,
    "structured_timeout_infrastructure_failure": True,
    "structured_timeout_classifier_verified": True,
    "contract_failure_not_classified_as_timeout": True,
    "empty_raw_answer_sha256_audited": True,
    "round3_schema_preserved": True,
    "round2_completeness_guard_preserved": True,
    "final_counted_attempt_protected_on_timeout": True,
    "blender_required": False,
}, indent=2))