#!/usr/bin/env python3
"""
Tests for Story 4.03 (Track B): TierExecutor — Tier-Appropriate Dispatcher

Black Box tests (BB): verify public contract from the outside.
White Box tests (WB): verify internal dispatch paths, registry mechanics,
                      observability events, and enriched payload structure.

Story: 4.03
File under test: core/routing/tier_executor.py
"""

from __future__ import annotations

import json
import sys
sys.path.insert(0, '/mnt/e/genesis-system')

import pytest
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch

from core.routing.tier_executor import (
    TierExecutor,
    T0_HANDLER_REGISTRY,
    register_t0_handler,
    EVENTS_DIR,
)
from core.routing.tier_classifier import RoutingDecision


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_decision(tier: str) -> RoutingDecision:
    """Build a RoutingDecision for the given tier."""
    model_map = {
        "T0": "python_function",
        "T1": "gemini-flash",
        "T2": "claude-opus-4-6",
    }
    return RoutingDecision(
        tier=tier,
        model=model_map[tier],
        rationale=f"test decision for {tier}",
    )


# ===========================================================================
# Black Box tests
# ===========================================================================

@pytest.mark.asyncio
async def test_bb1_t0_task_health_check_uses_python_handler():
    """BB1: T0 task 'health_check' → Python handler called, no LLM, correct result."""
    executor = TierExecutor(dispatch_fn=None)
    task_payload = {"type": "health_check", "task_id": "bb1-task"}
    decision = _make_decision("T0")

    result = await executor.execute(task_payload, decision)

    assert result["tier"] == "T0"
    assert result["model"] == "python_function"
    assert result["status"] == "ok"
    assert result["output"]["healthy"] is True


@pytest.mark.asyncio
async def test_bb2_t0_task_without_handler_falls_back_to_t1():
    """BB2: T0 task type with no registered handler → falls back gracefully to T1."""
    # Use a type that won't be in T0_HANDLER_REGISTRY
    unregistered_type = "__unregistered_t0_type_for_bb2__"
    assert unregistered_type not in T0_HANDLER_REGISTRY, \
        "Test prerequisite: unregistered_type must not be in registry"

    dispatch_called_with = []

    async def capture_dispatch(payload: dict) -> dict:
        dispatch_called_with.append(payload)
        return {"status": "ok", "task_id": payload.get("task_id"), "output": "dispatched"}

    executor = TierExecutor(dispatch_fn=capture_dispatch)
    task_payload = {"type": unregistered_type, "task_id": "bb2-task"}
    decision = _make_decision("T0")

    result = await executor.execute(task_payload, decision)

    # dispatch_fn must have been called (T1 fallback)
    assert len(dispatch_called_with) == 1, "dispatch_fn should be called on T0 fallback to T1"
    enriched = dispatch_called_with[0]
    assert enriched["tier"] == "T1"
    assert enriched["model"] == "gemini-flash"


@pytest.mark.asyncio
async def test_bb3_t1_task_calls_dispatch_fn_with_gemini_flash():
    """BB3: T1 task → dispatch_fn called with model='gemini-flash'."""
    dispatch_called_with = []

    async def capture_dispatch(payload: dict) -> dict:
        dispatch_called_with.append(payload)
        return {"status": "ok", "task_id": payload.get("task_id"), "output": "t1-result"}

    executor = TierExecutor(dispatch_fn=capture_dispatch)
    task_payload = {"type": "email_draft", "task_id": "bb3-task"}
    decision = _make_decision("T1")

    result = await executor.execute(task_payload, decision)

    assert len(dispatch_called_with) == 1
    enriched = dispatch_called_with[0]
    assert enriched["model"] == "gemini-flash"
    assert enriched["tier"] == "T1"


@pytest.mark.asyncio
async def test_bb4_t2_task_calls_dispatch_fn_with_opus():
    """BB4: T2 task → dispatch_fn called with model='claude-opus-4-6'."""
    dispatch_called_with = []

    async def capture_dispatch(payload: dict) -> dict:
        dispatch_called_with.append(payload)
        return {"status": "ok", "task_id": payload.get("task_id"), "output": "t2-result"}

    executor = TierExecutor(dispatch_fn=capture_dispatch)
    task_payload = {"type": "novel_task", "task_id": "bb4-task"}
    decision = _make_decision("T2")

    result = await executor.execute(task_payload, decision)

    assert len(dispatch_called_with) == 1
    enriched = dispatch_called_with[0]
    assert enriched["model"] == "claude-opus-4-6"
    assert enriched["tier"] == "T2"


@pytest.mark.asyncio
async def test_bb5_no_dispatch_fn_t1_returns_stub():
    """BB5a: No dispatch_fn provided → T1 returns structured stub result without raising."""
    executor = TierExecutor(dispatch_fn=None)
    task_payload = {"type": "email_draft", "task_id": "bb5a-task"}
    decision = _make_decision("T1")

    result = await executor.execute(task_payload, decision)

    assert result["status"] == "ok"
    assert result["tier"] == "T1"
    assert result["model"] == "gemini-flash"
    assert result["task_id"] == "bb5a-task"
    assert result["output"] is None


@pytest.mark.asyncio
async def test_bb5_no_dispatch_fn_t2_returns_stub():
    """BB5b: No dispatch_fn provided → T2 returns structured stub result without raising."""
    executor = TierExecutor(dispatch_fn=None)
    task_payload = {"type": "novel_reasoning", "task_id": "bb5b-task"}
    decision = _make_decision("T2")

    result = await executor.execute(task_payload, decision)

    assert result["status"] == "ok"
    assert result["tier"] == "T2"
    assert result["model"] == "claude-opus-4-6"
    assert result["task_id"] == "bb5b-task"
    assert result["output"] is None


# ===========================================================================
# White Box tests
# ===========================================================================

@pytest.mark.asyncio
async def test_wb1_t2_enriched_payload_has_tier_guard():
    """WB1: T2 path adds tier='T2' to enriched payload — prevents re-classification loop."""
    captured = []

    async def capture(payload: dict) -> dict:
        captured.append(payload)
        return {"status": "ok"}

    executor = TierExecutor(dispatch_fn=capture)
    task_payload = {"type": "strategic_analysis", "task_id": "wb1"}
    decision = _make_decision("T2")

    await executor.execute(task_payload, decision)

    assert len(captured) == 1
    enriched = captured[0]
    assert enriched["tier"] == "T2", (
        "T2 enriched payload MUST contain tier='T2' guard to prevent re-classification"
    )
    assert enriched["model"] == "claude-opus-4-6"


def test_wb2_t0_handler_registry_is_module_level_dict():
    """WB2: T0_HANDLER_REGISTRY is a module-level dict."""
    from core.routing import tier_executor as mod
    assert hasattr(mod, "T0_HANDLER_REGISTRY"), "T0_HANDLER_REGISTRY must be a module attribute"
    assert isinstance(mod.T0_HANDLER_REGISTRY, dict), \
        f"T0_HANDLER_REGISTRY must be a dict, got {type(mod.T0_HANDLER_REGISTRY)}"


def test_wb3_register_t0_handler_decorator_adds_to_registry():
    """WB3: @register_t0_handler decorator correctly adds handler to T0_HANDLER_REGISTRY."""
    unique_type = "__wb3_test_type__"

    # Clean up before test in case previous run left an entry
    T0_HANDLER_REGISTRY.pop(unique_type, None)
    assert unique_type not in T0_HANDLER_REGISTRY, "Pre-condition: key must not exist"

    @register_t0_handler(unique_type)
    def _my_wb3_handler(payload: dict) -> dict:
        return {"wb3": True}

    assert unique_type in T0_HANDLER_REGISTRY, \
        f"'{unique_type}' must be in T0_HANDLER_REGISTRY after decoration"
    assert T0_HANDLER_REGISTRY[unique_type] is _my_wb3_handler

    # Clean up
    del T0_HANDLER_REGISTRY[unique_type]


@pytest.mark.asyncio
async def test_wb4_event_logged_to_events_jsonl_on_t0_execution(tmp_path, monkeypatch):
    """WB4: Execution event is logged to events.jsonl after a successful T0 dispatch."""
    log_dir = tmp_path / "observability"

    # Patch EVENTS_DIR inside the module so logs go to tmp_path
    import core.routing.tier_executor as mod
    monkeypatch.setattr(mod, "EVENTS_DIR", log_dir)

    executor = TierExecutor(dispatch_fn=None)
    task_payload = {"type": "health_check", "task_id": "wb4-task"}
    decision = _make_decision("T0")

    await executor.execute(task_payload, decision)

    log_file = log_dir / "events.jsonl"
    assert log_file.exists(), "events.jsonl should be created after execution"

    lines = log_file.read_text().strip().split("\n")
    assert len(lines) >= 1, "At least one log entry should be written"

    event = json.loads(lines[-1])
    assert event["event_type"] == "tier_execution"
    assert event["task_id"] == "wb4-task"
    assert event["tier"] == "T0"
    assert event["model"] == "python_function"
    assert "timestamp" in event


@pytest.mark.asyncio
async def test_wb5_t0_handler_receives_original_task_payload():
    """WB5: T0 handler is called with the original task_payload dict (not a modified copy)."""
    received_payloads = []

    unique_type = "__wb5_handler_test__"
    T0_HANDLER_REGISTRY.pop(unique_type, None)

    @register_t0_handler(unique_type)
    def _wb5_handler(payload: dict) -> dict:
        received_payloads.append(payload)
        return {"status": "ok", "tier": "T0", "model": "python_function", "output": {}}

    executor = TierExecutor(dispatch_fn=None)
    task_payload = {"type": unique_type, "task_id": "wb5-task", "extra": "data"}
    decision = _make_decision("T0")

    await executor.execute(task_payload, decision)

    assert len(received_payloads) == 1
    # The handler received the exact same object (identity or at least same content)
    assert received_payloads[0]["type"] == unique_type
    assert received_payloads[0]["task_id"] == "wb5-task"
    assert received_payloads[0]["extra"] == "data"

    # Clean up
    del T0_HANDLER_REGISTRY[unique_type]


@pytest.mark.asyncio
async def test_wb6_t1_enriched_payload_contains_tier_and_model():
    """WB6: T1 enriched payload forwarded to dispatch_fn contains tier='T1' and model='gemini-flash'."""
    captured = []

    async def capture(payload: dict) -> dict:
        captured.append(payload)
        return {"status": "ok"}

    executor = TierExecutor(dispatch_fn=capture)
    task_payload = {"type": "email_draft", "task_id": "wb6-task", "body": "hello"}
    decision = _make_decision("T1")

    await executor.execute(task_payload, decision)

    assert len(captured) == 1
    enriched = captured[0]
    assert enriched["tier"] == "T1", "T1 enriched payload must contain tier='T1'"
    assert enriched["model"] == "gemini-flash", "T1 enriched payload must contain model='gemini-flash'"
    # Original keys must be preserved
    assert enriched["body"] == "hello"
    assert enriched["task_id"] == "wb6-task"


# ===========================================================================
# Import / export sanity
# ===========================================================================

def test_tier_executor_importable_from_package():
    """TierExecutor is exported from core.routing.__init__."""
    from core.routing import TierExecutor as TE
    assert TE is TierExecutor


def test_t0_handler_registry_importable_from_package():
    """T0_HANDLER_REGISTRY is exported from core.routing.__init__."""
    from core.routing import T0_HANDLER_REGISTRY as reg
    assert isinstance(reg, dict)


def test_register_t0_handler_importable_from_package():
    """register_t0_handler is exported from core.routing.__init__."""
    from core.routing import register_t0_handler as rth
    assert callable(rth)


# ===========================================================================
# Standalone runner
# ===========================================================================

if __name__ == "__main__":
    import asyncio

    async def _smoke():
        print("Running manual smoke tests for Story 4.03...")

        # BB1: T0 health_check
        executor = TierExecutor()
        result = await executor.execute(
            {"type": "health_check", "task_id": "smoke-1"},
            _make_decision("T0"),
        )
        assert result["tier"] == "T0" and result["output"]["healthy"] is True
        print("[PASS] BB1: T0 health_check → python_function")

        # BB5: No dispatch_fn → T1 stub
        result = await executor.execute(
            {"type": "email_draft", "task_id": "smoke-2"},
            _make_decision("T1"),
        )
        assert result["model"] == "gemini-flash" and result["output"] is None
        print("[PASS] BB5: No dispatch_fn → T1 stub result")

        # BB5: No dispatch_fn → T2 stub
        result = await executor.execute(
            {"type": "novel_task", "task_id": "smoke-3"},
            _make_decision("T2"),
        )
        assert result["model"] == "claude-opus-4-6" and result["output"] is None
        print("[PASS] BB5: No dispatch_fn → T2 stub result")

        # WB3: register_t0_handler
        test_type = "__smoke_handler__"
        T0_HANDLER_REGISTRY.pop(test_type, None)

        @register_t0_handler(test_type)
        def _smoke_handler(p): return {"status": "ok", "tier": "T0", "model": "python_function", "output": {}}

        assert test_type in T0_HANDLER_REGISTRY
        del T0_HANDLER_REGISTRY[test_type]
        print("[PASS] WB3: register_t0_handler adds to registry")

        print("\nAll manual smoke tests passed — Story 4.03 (Track B)")

    asyncio.run(_smoke())