from __future__ import annotations

import contextlib
from dataclasses import dataclass
from types import SimpleNamespace
from typing import Any, Callable, Dict, List

import pytest


@dataclass
class _FakeSpanContext:
    trace_id: int = 123
    span_id: int = 456


class _FakeRootSpan:
    def __init__(self, trace_id: int = 123):
        self._ctx = _FakeSpanContext(trace_id=trace_id, span_id=456)

    def get_span_context(self):
        return self._ctx


@contextlib.contextmanager
def _dummy_trace_run(_name: str, metadata: Dict[str, Any] | None = None):
    yield _FakeRootSpan()


@contextlib.contextmanager
def _dummy_trace_agent(_name: str, metadata: Dict[str, Any] | None = None):
    yield None


class _ToolUseBlock:
    def __init__(self, *, name: str, input: Dict[str, Any], tool_use_id: str = "tu_1"):
        self.type = "tool_use"
        self.name = name
        self.input = input
        self.id = tool_use_id


class _TextBlock:
    def __init__(self, text: str):
        self.type = "text"
        self.text = text


class _FakeAnthropicResponse:
    def __init__(self, *, stop_reason: str, content: List[Any]):
        self.stop_reason = stop_reason
        self.content = content


class _FakeMessages:
    def __init__(self, scripted: List[_FakeAnthropicResponse]):
        self._scripted = list(scripted)
        self.calls: List[Dict[str, Any]] = []

    def create(self, **kwargs):
        self.calls.append(kwargs)
        if not self._scripted:
            # Default end_turn to stop loops cleanly.
            return _FakeAnthropicResponse(stop_reason="end_turn", content=[_TextBlock("ok")])
        return self._scripted.pop(0)


class _FakeAnthropicClient:
    def __init__(self, scripted: List[_FakeAnthropicResponse]):
        self.messages = _FakeMessages(scripted)


def _assert_tools_schema_matches_dispatch(mod):
    schema_names = {t["name"] for t in getattr(mod, "TOOLS_SCHEMA")}
    dispatch_names = set(getattr(mod, "TOOL_DISPATCH").keys())
    assert schema_names == dispatch_names


def test_custom_scorer_demo_non_live_tool_loop(monkeypatch):
    import tests.test_custom_scorer as demo

    # Prevent trace initialization requirements.
    monkeypatch.setattr(demo, "trace_run", _dummy_trace_run)
    monkeypatch.setattr(demo, "trace_agent", _dummy_trace_agent)
    _assert_tools_schema_matches_dispatch(demo)

    called: List[Dict[str, Any]] = []
    orig = demo.TOOL_DISPATCH["analyze_code"]

    def wrapped_analyze_code(**kwargs):
        called.append(kwargs)
        return orig(**kwargs)

    demo.TOOL_DISPATCH["analyze_code"] = wrapped_analyze_code

    client = _FakeAnthropicClient(
        scripted=[
            _FakeAnthropicResponse(
                stop_reason="tool_use",
                content=[
                    _ToolUseBlock(
                        name="analyze_code",
                        input={"code": "def f():\n  return 1\n", "language": "python"},
                        tool_use_id="tu_1",
                    )
                ],
            ),
            _FakeAnthropicResponse(stop_reason="end_turn", content=[_TextBlock("analysis done")]),
            _FakeAnthropicResponse(stop_reason="end_turn", content=[_TextBlock("review done")]),
        ]
    )

    trace_id, _root = demo.run_review_agent(client, "def f():\n  return 1\n", "python")
    assert isinstance(trace_id, str) and len(trace_id) == 32
    assert called, "Expected analyze_code to be executed via TOOL_DISPATCH"


def test_prebuilt_evals_demo_non_live_tool_loop(monkeypatch):
    import tests.test_prebuilt_evals as demo

    monkeypatch.setattr(demo, "trace_run", _dummy_trace_run)
    monkeypatch.setattr(demo, "trace_agent", _dummy_trace_agent)
    _assert_tools_schema_matches_dispatch(demo)

    called: List[Dict[str, Any]] = []
    orig = demo.TOOL_DISPATCH["search_orders"]

    def wrapped_search_orders(**kwargs):
        called.append(kwargs)
        return orig(**kwargs)

    demo.TOOL_DISPATCH["search_orders"] = wrapped_search_orders

    client = _FakeAnthropicClient(
        scripted=[
            _FakeAnthropicResponse(
                stop_reason="tool_use",
                content=[_ToolUseBlock(name="search_orders", input={"customer_email": "alice@example.com"}, tool_use_id="tu_1")],
            ),
            _FakeAnthropicResponse(stop_reason="end_turn", content=[_TextBlock("triage done")]),
            _FakeAnthropicResponse(stop_reason="end_turn", content=[_TextBlock("resolution done")]),
        ]
    )

    trace_id, _root = demo.run_support_agent(client, "I want to return my headphones.")
    assert isinstance(trace_id, str) and len(trace_id) == 32
    assert called, "Expected search_orders to be executed via TOOL_DISPATCH"


def test_realtime_custom_evals_demo_non_live_tool_loop(monkeypatch):
    import tests.test_realtime_custom_evals as demo

    monkeypatch.setattr(demo, "trace_run", _dummy_trace_run)
    monkeypatch.setattr(demo, "trace_agent", _dummy_trace_agent)
    _assert_tools_schema_matches_dispatch(demo)

    called: List[Dict[str, Any]] = []
    orig = demo.TOOL_DISPATCH["lookup_error"]

    def wrapped_lookup_error(**kwargs):
        called.append(kwargs)
        return orig(**kwargs)

    demo.TOOL_DISPATCH["lookup_error"] = wrapped_lookup_error

    client = _FakeAnthropicClient(
        scripted=[
            _FakeAnthropicResponse(
                stop_reason="tool_use",
                content=[_ToolUseBlock(name="lookup_error", input={"error_code": "DB-WARN-4201"}, tool_use_id="tu_1")],
            ),
            _FakeAnthropicResponse(stop_reason="end_turn", content=[_TextBlock("diagnosis done")]),
            _FakeAnthropicResponse(stop_reason="end_turn", content=[_TextBlock("resolution done")]),
        ]
    )

    trace_id, _root = demo.run_support_agent(client, "Seeing DB-WARN-4201, what do I do?")
    assert isinstance(trace_id, str) and len(trace_id) == 32
    assert called, "Expected lookup_error to be executed via TOOL_DISPATCH"


def test_span_evals_demo_non_live_tool_loop(monkeypatch):
    import tests.test_span_evals as demo

    monkeypatch.setattr(demo, "trace_run", _dummy_trace_run)
    monkeypatch.setattr(demo, "trace_agent", _dummy_trace_agent)
    _assert_tools_schema_matches_dispatch(demo)

    called: List[Dict[str, Any]] = []
    orig = demo.TOOL_DISPATCH["search_recipes"]

    def wrapped_search_recipes(**kwargs):
        called.append(kwargs)
        return orig(**kwargs)

    demo.TOOL_DISPATCH["search_recipes"] = wrapped_search_recipes

    client = _FakeAnthropicClient(
        scripted=[
            _FakeAnthropicResponse(
                stop_reason="tool_use",
                content=[_ToolUseBlock(name="search_recipes", input={"query": "pasta", "dietary_preference": "none"}, tool_use_id="tu_1")],
            ),
            _FakeAnthropicResponse(stop_reason="end_turn", content=[_TextBlock("research done")]),
            _FakeAnthropicResponse(stop_reason="end_turn", content=[_TextBlock("planner done")]),
        ]
    )

    trace_id, _root = demo.run_recipe_agent(client, "Find a pasta recipe.")
    assert isinstance(trace_id, str) and len(trace_id) == 32
    assert called, "Expected search_recipes to be executed via TOOL_DISPATCH"


def test_session_chatbot_helpers_deterministic():
    # The interactive chatbot is inherently non-deterministic.
    # We still validate that the pure helper functions are deterministic.
    import tests.test_session_chatbot as demo

    assert demo.calculate_compound_growth(1000, 0.05, 2)["final_balance"] == 1102.5
    assert demo.format_currency(1234.5) == "$1,234.50"
    r = demo.assess_risk_score(age=30, income=100_000, debt=0)
    assert 1 <= r["score"] <= 10

