"""
validate.py — AdaptiveWorld Submission Validator
=================================================
Runs a comprehensive offline validation of the adaptive-world-env submission.

Checks:
  1. Scenario registry  — all 12 scenarios, all difficulty levels
  2. Drift configs      — every scenario has a matching DriftInjector config
  3. DriftInjector      — inject() mutates world correctly
  4. AdaptiveGrader     — grade_task / grade_belief / infer_belief_from_actions
  5. DriftDifficultyController — escalation logic
  6. Episode lifecycle  — simulated reset → step → done for each difficulty
  7. inference.py       — import check (no LLM call needed)

Run from the adaptive-world-env directory:
    python validate.py
"""

import sys
import os
import copy
import traceback

# ── Bootstrap path ─────────────────────────────────────────────────────────────
ROOT = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, ROOT)

# ── Terminal colours (Windows-safe via ANSI or plain) ──────────────────────────
try:
    import ctypes
    ctypes.windll.kernel32.SetConsoleMode(ctypes.windll.kernel32.GetStdHandle(-11), 7)
    GREEN  = "\033[92m"
    RED    = "\033[91m"
    YELLOW = "\033[93m"
    BOLD   = "\033[1m"
    RESET  = "\033[0m"
except Exception:
    GREEN = RED = YELLOW = BOLD = RESET = ""

PASS = 0
FAIL = 0
ERRORS = []


def ok(msg):
    global PASS
    PASS += 1
    print(f"  {GREEN}✓{RESET} {msg}")


def fail(msg, exc=None):
    global FAIL
    FAIL += 1
    ERRORS.append(msg)
    print(f"  {RED}✗{RESET} {msg}")
    if exc:
        print(f"    {YELLOW}→ {exc}{RESET}")


def section(title):
    print(f"\n{BOLD}{'─'*60}{RESET}")
    print(f"{BOLD}  {title}{RESET}")
    print(f"{BOLD}{'─'*60}{RESET}")


# ══════════════════════════════════════════════════════════════════════════════
# 1. Scenario Registry
# ══════════════════════════════════════════════════════════════════════════════
section("1 / 7 · Scenario Registry")

try:
    from scenarios.registry import SCENARIO_REGISTRY, ALL_SCENARIOS

    DIFFICULTIES = ["easy", "medium", "hard", "expert"]
    REQUIRED_FIELDS = ["id", "description", "domain", "drift_trigger_step",
                       "drift_type", "task_goal", "max_steps"]

    for level in DIFFICULTIES:
        if level in SCENARIO_REGISTRY:
            ok(f"Difficulty '{level}' present")
        else:
            fail(f"Difficulty '{level}' MISSING from registry")

    if len(ALL_SCENARIOS) == 12:
        ok(f"Total scenario count: {len(ALL_SCENARIOS)} (expected 12)")
    else:
        fail(f"Expected 12 scenarios, found {len(ALL_SCENARIOS)}")

    for level, scenarios in SCENARIO_REGISTRY.items():
        if len(scenarios) == 3:
            ok(f"  '{level}' has 3 scenarios")
        else:
            fail(f"  '{level}' has {len(scenarios)} scenarios (expected 3)")

    bad = []
    for s in ALL_SCENARIOS:
        missing = [f for f in REQUIRED_FIELDS if f not in s]
        if missing:
            bad.append(f"{s.get('id', '?')} missing: {missing}")
        if "drift_occurred" in s:
            bad.append(f"{s['id']} has forbidden key 'drift_occurred'")

    if bad:
        for b in bad:
            fail(b)
    else:
        ok(f"All {len(ALL_SCENARIOS)} scenarios have required fields (no forbidden keys)")

except Exception as e:
    fail("Scenario registry import/validation failed", e)
    traceback.print_exc()


# ══════════════════════════════════════════════════════════════════════════════
# 2. Drift Configs
# ══════════════════════════════════════════════════════════════════════════════
section("2 / 7 · Drift Injector Configs")

try:
    from server.drift_injector import DriftInjector, DRIFT_CONFIGS

    missing_configs = [s["id"] for s in ALL_SCENARIOS if s["id"] not in DRIFT_CONFIGS]
    if missing_configs:
        for sid in missing_configs:
            fail(f"No DRIFT_CONFIG entry for scenario '{sid}'")
    else:
        ok(f"All {len(ALL_SCENARIOS)} scenarios have DRIFT_CONFIG entries")

    # Quick smoke-test inject on a few known scenarios
    for sid in ["easy_field_rename", "easy_endpoint_version", "hard_status_meaning"]:
        try:
            inj = DriftInjector(sid)
            before = copy.deepcopy(inj.get_world())
            inj.inject()
            after = inj.get_world()
            if before != after:
                ok(f"  DriftInjector({sid!r}).inject() mutates world ✓")
            else:
                fail(f"  DriftInjector({sid!r}).inject() did NOT change world")
        except Exception as ex:
            fail(f"  DriftInjector({sid!r}) error", ex)

    # Expert-specific helpers
    try:
        inj = DriftInjector("expert_transient_vs_real")
        step = inj.get_transient_error_step()
        if step == 2:
            ok(f"  expert_transient_vs_real: transient_error_step == 2")
        else:
            fail(f"  expected transient_error_step 2, got {step}")
    except Exception as ex:
        fail("  expert_transient_vs_real transient error step check", ex)

    try:
        inj = DriftInjector("expert_cross_service")
        step = inj.get_secondary_drift_step()
        if step == 6:
            ok(f"  expert_cross_service: secondary_drift_step == 6")
        else:
            fail(f"  expected secondary_drift_step 6, got {step}")
    except Exception as ex:
        fail("  expert_cross_service secondary drift step check", ex)

    try:
        DriftInjector("this_scenario_does_not_exist")
        fail("  Unknown scenario should raise ValueError but did not")
    except ValueError:
        ok("  Unknown scenario raises ValueError ✓")
    except Exception as ex:
        fail("  Unknown scenario check raised unexpected exception", ex)

except Exception as e:
    fail("Drift injector section failed", e)
    traceback.print_exc()


# ══════════════════════════════════════════════════════════════════════════════
# 3. AdaptiveGrader
# ══════════════════════════════════════════════════════════════════════════════
section("3 / 7 · AdaptiveGrader")

try:
    from graders.grader import AdaptiveGrader
    g = AdaptiveGrader()

    # grade_task
    r = g.grade_task(task_completed=False, steps_taken=5, max_steps=8, drift_detected=False)
    if r == 0.001:
        ok("grade_task(failed, no detection) == 0.001")
    else:
        fail(f"grade_task(failed, no detection) expected 0.001 got {r}")

    r = g.grade_task(task_completed=False, steps_taken=5, max_steps=8, drift_detected=True)
    if r == 0.150:
        ok("grade_task(failed, detected) == 0.150")
    else:
        fail(f"grade_task(failed, detected) expected 0.150 got {r}")

    r_no = g.grade_task(True, 4, 8, False)
    r_yes = g.grade_task(True, 4, 8, True)
    if r_yes > r_no:
        ok("Proactive drift detection gives bonus reward ✓")
    else:
        fail(f"Bonus expected: r_with_detect={r_yes:.4f} vs r_without={r_no:.4f}")

    r = g.grade_task(True, 1, 8, True)
    if r <= 0.999:
        ok(f"Task reward capped at 0.999 (got {r})")
    else:
        fail(f"Task reward not capped: {r}")

    # grade_belief — field_rename
    score = g.grade_belief(
        {"order_field": "quantity", "required_extra": "customer_id"},
        {"order_field": "quantity", "required_extra": "customer_id"},
        "field_rename"
    )
    if score >= 0.9:
        ok(f"grade_belief: correct field_rename → {score:.3f} (≥0.9) ✓")
    else:
        fail(f"grade_belief: correct field_rename expected ≥0.9, got {score:.3f}")

    score = g.grade_belief({"order_field": "qty"}, {"order_field": "quantity", "required_extra": "customer_id"}, "field_rename")
    if score < 0.5:
        ok(f"grade_belief: stale field_rename → {score:.3f} (<0.5) ✓")
    else:
        fail(f"grade_belief: stale belief expected <0.5, got {score:.3f}")

    # grade_belief — endpoint_version
    score = g.grade_belief({"endpoint": "/mock_api/v2/rooms/book"}, {"rooms_endpoint": "/mock_api/v2/rooms/book"}, "endpoint_version")
    if score == 1.0:
        ok(f"grade_belief: correct endpoint_version → 1.0 ✓")
    else:
        fail(f"grade_belief: correct endpoint expected 1.0, got {score}")

    # grade_belief — None
    score = g.grade_belief(None, {"x": 1}, "field_rename")
    if score == 0.0:
        ok("grade_belief(None, ...) == 0.0 ✓")
    else:
        fail(f"grade_belief(None) expected 0.0, got {score}")

    # infer_belief_from_actions
    log = [
        {"step": 1, "url": "/mock_api/orders", "status": 200, "response": "{}"},
        {"step": 2, "url": "/mock_api/orders", "status": 422, "response": '{"detail": "..."}'},
        {"step": 3, "url": "/openapi.json",    "status": 200, "response": "{}"},
        {"step": 4, "url": "/mock_api/orders", "status": 200, "response": '{"order_id": "x"}'},
    ]
    score = g.infer_belief_from_actions(log, "field_rename")
    if score == 0.6:
        ok(f"infer_belief_from_actions: probed+recovered → 0.6 ✓")
    else:
        fail(f"infer_belief_from_actions expected 0.6, got {score}")

    score = g.infer_belief_from_actions([], "field_rename")
    if score == 0.0:
        ok("infer_belief_from_actions([]) == 0.0 ✓")
    else:
        fail(f"infer_belief_from_actions([]) expected 0.0, got {score}")

except Exception as e:
    fail("AdaptiveGrader section failed", e)
    traceback.print_exc()


# ══════════════════════════════════════════════════════════════════════════════
# 4. DriftDifficultyController
# ══════════════════════════════════════════════════════════════════════════════
section("4 / 7 · DriftDifficultyController")

try:
    from server.difficulty_controller import DriftDifficultyController

    ctrl = DriftDifficultyController()
    if ctrl.level == 0:
        ok("Initial level == 0 ✓")
    else:
        fail(f"Expected initial level 0, got {ctrl.level}")

    # not enough data
    ctrl2 = DriftDifficultyController()
    for _ in range(4):
        ctrl2.record("field_rename", 0.95)
        ctrl2.record("endpoint_version", 0.95)
        ctrl2.record("policy_change", 0.95)
    if ctrl2.level == 0:
        ok("No escalation with < 5 window data ✓")
    else:
        fail(f"Expected no escalation, level={ctrl2.level}")

    # escalation happens at window=5
    ctrl3 = DriftDifficultyController()
    for _ in range(5):
        ctrl3.record("field_rename", 0.90)
        ctrl3.record("endpoint_version", 0.90)
        ctrl3.record("policy_change", 0.90)
    if ctrl3.level == 1:
        ok("Escalates to level 1 after 5 high-accuracy episodes ✓")
    else:
        fail(f"Expected level 1 after escalation, got {ctrl3.level}")

    # no escalation if one type below threshold
    ctrl4 = DriftDifficultyController()
    for _ in range(5):
        ctrl4.record("field_rename", 0.90)
        ctrl4.record("endpoint_version", 0.50)  # below threshold
        ctrl4.record("policy_change", 0.90)
    if ctrl4.level == 0:
        ok("No escalation when one drift type below threshold ✓")
    else:
        fail(f"Expected no escalation, got level {ctrl4.level}")

    # reset
    ctrl5 = DriftDifficultyController()
    ctrl5._escalation_level = 2
    ctrl5.reset()
    if ctrl5.level == 0 and len(ctrl5._history) == 0:
        ok("reset() clears level and history ✓")
    else:
        fail(f"reset() failed: level={ctrl5.level}, history_len={len(ctrl5._history)}")

    # get_scenario_params at level 0 is identity
    sc = {"id": "easy_field_rename", "drift_trigger_step": 3, "drift_type": "field_rename"}
    result = ctrl.get_scenario_params(sc)
    if result == sc:
        ok("get_scenario_params at level 0 returns unchanged scenario ✓")
    else:
        fail(f"get_scenario_params unexpected change: {result}")

except Exception as e:
    fail("DriftDifficultyController section failed", e)
    traceback.print_exc()


# ══════════════════════════════════════════════════════════════════════════════
# 5. Models
# ══════════════════════════════════════════════════════════════════════════════
section("5 / 7 · Models (AdaptiveAction / AdaptiveObservation / AdaptiveState)")

try:
    from models import AdaptiveAction, AdaptiveObservation, AdaptiveState

    a = AdaptiveAction()
    if a.action_type == "call_api" and a.method == "GET":
        ok("AdaptiveAction defaults: action_type='call_api', method='GET' ✓")
    else:
        fail(f"AdaptiveAction defaults wrong: {a.action_type}, {a.method}")

    obs = AdaptiveObservation()
    if not hasattr(obs, "drift_occurred") and not hasattr(obs, "drift_type"):
        ok("AdaptiveObservation has NO drift_occurred / drift_type (v2 compliant) ✓")
    else:
        fail("AdaptiveObservation still has forbidden field drift_occurred / drift_type")

    for field in ["prior_world_model", "episode_history", "belief_accuracy", "difficulty_level"]:
        if hasattr(obs, field):
            ok(f"  AdaptiveObservation has v2 field '{field}' ✓")
        else:
            fail(f"  AdaptiveObservation missing v2 field '{field}'")

    state = AdaptiveState()
    if not state.drift_injected and state.agent_belief == {} and state.world_truth == {}:
        ok("AdaptiveState defaults correct ✓")
    else:
        fail(f"AdaptiveState defaults wrong: {state}")

except Exception as e:
    fail("Models section failed", e)
    traceback.print_exc()


# ══════════════════════════════════════════════════════════════════════════════
# 6. Simulated Episode Lifecycle (no live server)
# ══════════════════════════════════════════════════════════════════════════════
section("6 / 7 · Simulated Episode Lifecycle (offline, all difficulties)")

try:
    from unittest.mock import patch, MagicMock
    from models import AdaptiveAction, AdaptiveObservation

    def make_mock_http(status=200, body="{}"):
        mock_resp = MagicMock()
        mock_resp.status_code = status
        mock_resp.text = body
        mock_resp.headers = {"content-type": "application/json"}
        return mock_resp

    for difficulty in ["easy", "medium", "hard", "expert"]:
        try:
            with patch("httpx.Client") as MockClient:
                cm = MockClient.return_value.__enter__.return_value
                # Mock the admin mutate call
                cm.post.return_value = make_mock_http(200, "{}")
                # Mock API calls — return a successful order
                cm.request.return_value = make_mock_http(
                    200, '{"order_id": "abc123", "status": "confirmed"}'
                )
                cm.get.return_value = make_mock_http(200, '{"endpoints": []}')

                import importlib
                import server.adaptive_world_environment as awe_module
                importlib.reload(awe_module)
                env = awe_module.AdaptiveWorldEnvironment()

                # reset
                obs = env.reset(scenario_id="auto", difficulty=difficulty)
                assert isinstance(obs, AdaptiveObservation), "reset() must return AdaptiveObservation"
                assert not obs.done, "done should be False after reset"
                assert obs.task_description, "task_description should not be empty"

                # step — probe schema
                action = AdaptiveAction(action_type="probe_schema")
                obs2 = env.step(action)
                assert isinstance(obs2, AdaptiveObservation)

                # step — call_api
                action2 = AdaptiveAction(
                    action_type="call_api",
                    method="POST",
                    url="/mock_api/orders",
                    body={"qty": 1, "product_id": 5},
                )
                obs3 = env.step(action2)
                assert isinstance(obs3, AdaptiveObservation)
                assert env.state.step_count >= 1

                # submit_result
                submit = AdaptiveAction(
                    action_type="submit_result",
                    belief_state={"order_field": "qty", "drift_detected": False},
                )
                obs_final = env.step(submit)
                assert obs_final.done, "done must be True after submit_result"
                assert 0.0 <= obs_final.reward <= 1.0, f"reward out of range: {obs_final.reward}"

                ok(f"  Difficulty '{difficulty}': reset→step→submit OK "
                   f"(reward={obs_final.reward:.4f})")

        except Exception as ex:
            fail(f"  Difficulty '{difficulty}' lifecycle failed", ex)
            traceback.print_exc()

except Exception as e:
    fail("Episode lifecycle section failed", e)
    traceback.print_exc()


# ══════════════════════════════════════════════════════════════════════════════
# 7. Inference.py Import Check
# ══════════════════════════════════════════════════════════════════════════════
section("7 / 7 · inference.py Import / Parse Check")

try:
    import ast
    inf_path = os.path.join(ROOT, "inference.py")
    if os.path.exists(inf_path):
        with open(inf_path, "r", encoding="utf-8") as f:
            src = f.read()
        try:
            tree = ast.parse(src)
            ok("inference.py parses without syntax errors ✓")

            # Check key functions exist
            funcs = {n.name for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)}
            for fn in ["run_episode", "run_evaluation", "build_user_message", "parse_action"]:
                if fn in funcs:
                    ok(f"  Function '{fn}' found in inference.py ✓")
                else:
                    fail(f"  Function '{fn}' MISSING from inference.py")

            # Check all difficulty choices are listed
            src_lower = src.lower()
            for d in ["easy", "medium", "hard", "expert"]:
                if d in src_lower:
                    ok(f"  Difficulty '{d}' referenced in inference.py ✓")
                else:
                    fail(f"  Difficulty '{d}' NOT referenced in inference.py")

        except SyntaxError as se:
            fail(f"inference.py has syntax error: {se}")
    else:
        fail(f"inference.py not found at {inf_path}")

except Exception as e:
    fail("inference.py check failed", e)
    traceback.print_exc()


# ══════════════════════════════════════════════════════════════════════════════
# Summary
# ══════════════════════════════════════════════════════════════════════════════
total = PASS + FAIL
print(f"\n{BOLD}{'═'*60}{RESET}")
if FAIL == 0:
    print(f"{GREEN}{BOLD}  ALL {PASS}/{total} CHECKS PASSED ✓{RESET}")
    print(f"{GREEN}{BOLD}  Submission looks valid!{RESET}")
else:
    print(f"{RED}{BOLD}  {FAIL} / {total} CHECKS FAILED ✗{RESET}")
    print(f"\n{YELLOW}Failed checks:{RESET}")
    for err in ERRORS:
        print(f"  • {err}")
print(f"{BOLD}{'═'*60}{RESET}\n")

sys.exit(0 if FAIL == 0 else 1)