"""
Unit tests for the AutoGen integration.
Fully mocked — no autogen-agentchat, autogen-ext, or network required.
"""
import pytest
from dataclasses import dataclass, field
from typing import List


# ------------------------------------------------------------------
# Mock AutoGen types (simulating v0.7.5 API)
# ------------------------------------------------------------------

@dataclass
class RequestUsage:
    prompt_tokens: int = 0
    completion_tokens: int = 0


@dataclass
class FunctionCall:
    id: str = ""
    name: str = ""
    arguments: str = ""


@dataclass
class FunctionExecutionResult:
    content: str = ""
    name: str = ""
    call_id: str = ""
    is_error: bool = False


@dataclass
class CreateResult:
    content: object = ""
    usage: RequestUsage = None
    finish_reason: str = "stop"
    cached: bool = False


@dataclass
class UserMessage:
    content: str = ""
    type: str = "UserMessage"
    source: str = "user"


@dataclass
class AssistantMessage:
    content: str = ""
    type: str = "AssistantMessage"
    source: str = "assistant"


@dataclass
class FunctionExecutionResultMessage:
    content: list = field(default_factory=list)
    type: str = "FunctionExecutionResultMessage"
    source: str = "tool"


class MockModelClient:
    """Simulates OpenAIChatCompletionClient for testing."""

    def __init__(self, responses: list):
        self._responses = responses
        self._call_index = 0
        self._raw_config = {"model": "gpt-4o-mini"}
        self.model_info = {"family": "gpt-4o-mini"}
        self._total_prompt = 0
        self._total_completion = 0

    async def create(self, messages, *args, **kwargs):
        if self._call_index >= len(self._responses):
            raise RuntimeError("No more mock responses")
        result = self._responses[self._call_index]
        self._call_index += 1
        if result.usage:
            self._total_prompt += result.usage.prompt_tokens
            self._total_completion += result.usage.completion_tokens
        return result

    def total_usage(self):
        return RequestUsage(self._total_prompt, self._total_completion)

    def actual_usage(self):
        return self.total_usage()


class MockVisibe:
    """Mock Visibe client that captures create_trace, queue_span, complete_trace."""

    def __init__(self):
        self.session_id = "test-session-123"
        self._trace_header = None
        self._spans = []
        self._summary = None

    def create_trace(self, trace_data):
        self._trace_header = trace_data
        return True

    def queue_span(self, trace_id, span_data):
        self._spans.append(span_data)
        return None

    def complete_trace(self, trace_id, summary):
        self._summary = summary
        return True

    @property
    def traces(self):
        """Reconstruct a trace dict for test assertions."""
        if not self._trace_header or not self._summary:
            return []

        # agent_cost_breakdown from llm_call spans
        agent_costs = {}
        for span in self._spans:
            if span.get("type") == "llm_call":
                agent = span.get("agent_name", "autogen")
                if agent not in agent_costs:
                    agent_costs[agent] = {
                        "agent_name": agent,
                        "model": span.get("model", "unknown"),
                        "input_tokens": 0,
                        "output_tokens": 0,
                        "cost": 0.0,
                    }
                agent_costs[agent]["input_tokens"] += span.get("input_tokens", 0)
                agent_costs[agent]["output_tokens"] += span.get("output_tokens", 0)
                agent_costs[agent]["cost"] += span.get("cost", 0)
        agent_cost_breakdown = list(agent_costs.values())

        # tools_used from tool_call spans
        tools_used = [
            {
                "tool_name": s.get("tool_name", "unknown"),
                "agent_name": s.get("agent_name", "autogen"),
                "execution_time_ms": s.get("duration_ms", 0),
                "status": s.get("status", "success"),
                "input": s.get("input_text", ""),
                "output": s.get("output_text", ""),
            }
            for s in self._spans if s.get("type") == "tool_call"
        ]

        # errors from error spans
        errors = [
            {
                "timestamp": s.get("timestamp", ""),
                "agent_name": s.get("agent_name", "unknown"),
                "tool_name": s.get("tool_name"),
                "error_type": s.get("error_type", "unknown"),
                "message": s.get("error_message", ""),
            }
            for s in self._spans if s.get("type") == "error"
        ]

        # steps
        steps = [
            {**s, "step_id": s.get("span_id", s.get("step_id", ""))}
            for s in self._spans
        ]

        agents_executed = list(dict.fromkeys(
            s.get("agent_name") for s in self._spans
            if s.get("type") == "agent_start" and s.get("agent_name")
        )) or (["autogen"] if agent_cost_breakdown else [])

        return [{
            "trace_id": self._trace_header.get("trace_id"),
            "prompt": self._summary.get("prompt", ""),
            "total_input_tokens": self._summary.get("total_input_tokens", 0),
            "total_output_tokens": self._summary.get("total_output_tokens", 0),
            "total_cost": self._summary.get("total_cost", 0),
            "status": self._summary.get("status", "completed"),
            "agent_cost_breakdown": agent_cost_breakdown,
            "agents_executed": agents_executed,
            "task_cost_breakdown": agent_cost_breakdown[:1] if agent_cost_breakdown else [],
            "tools_used": tools_used,
            "errors": errors,
            "metadata": {"steps": steps},
        }]


# ------------------------------------------------------------------
# Tests
# ------------------------------------------------------------------

@pytest.mark.asyncio
async def test_basic_llm_call_tracking():
    from visibe.integrations.autogen import _AutoGenTracker

    mock_visibe = MockVisibe()
    model_client = MockModelClient([
        CreateResult(
            content="The answer is 100.",
            usage=RequestUsage(prompt_tokens=39, completion_tokens=250),
        )
    ])

    tracker = _AutoGenTracker(mock_visibe, model_client, "test-basic", agents=[])
    tracker.start()

    await model_client.create([UserMessage(content="Calculate 25 * 4")])
    tracker.stop()

    trace = mock_visibe.traces[0]
    assert trace["total_input_tokens"] == 39
    assert trace["total_output_tokens"] == 250
    assert trace["total_cost"] > 0
    assert trace["status"] == "completed"
    assert len(trace["metadata"]["steps"]) >= 1
    assert len(trace["task_cost_breakdown"]) == 1
    assert trace["task_cost_breakdown"][0]["input_tokens"] == 39
    assert trace["task_cost_breakdown"][0]["output_tokens"] == 250


@pytest.mark.asyncio
async def test_tool_call_tracking():
    from visibe.integrations.autogen import _AutoGenTracker

    mock_visibe = MockVisibe()
    model_client = MockModelClient([
        CreateResult(
            content=[FunctionCall(id="call_abc", name="calculate", arguments='{"expression": "25 * 4"}')],
            usage=RequestUsage(prompt_tokens=50, completion_tokens=20),
        ),
        CreateResult(
            content="Based on the calculation, 25 * 4 = 100.",
            usage=RequestUsage(prompt_tokens=80, completion_tokens=30),
        ),
    ])

    tracker = _AutoGenTracker(mock_visibe, model_client, "test-tools", agents=[])
    tracker.start()

    await model_client.create([UserMessage(content="Calculate 25 * 4")])
    await model_client.create([
        UserMessage(content="Calculate 25 * 4"),
        FunctionExecutionResultMessage(content=[
            FunctionExecutionResult(content="100", name="calculate", call_id="call_abc", is_error=False)
        ]),
    ])

    tracker.stop()

    trace = mock_visibe.traces[0]
    assert len(trace["tools_used"]) == 1
    assert trace["tools_used"][0]["tool_name"] == "calculate"
    assert trace["tools_used"][0]["status"] == "success"
    assert trace["tools_used"][0]["output"] == "100"
    assert "25 * 4" in trace["tools_used"][0]["input"]
    assert len(trace["agent_cost_breakdown"]) == 1
    assert trace["agent_cost_breakdown"][0]["input_tokens"] == 130  # 50 + 80
    assert trace["agent_cost_breakdown"][0]["output_tokens"] == 50   # 20 + 30


@pytest.mark.asyncio
async def test_tool_error_tracking():
    from visibe.integrations.autogen import _AutoGenTracker

    mock_visibe = MockVisibe()
    model_client = MockModelClient([
        CreateResult(
            content=[FunctionCall(id="call_err", name="web_search", arguments='{"query": "news"}')],
            usage=RequestUsage(prompt_tokens=40, completion_tokens=15),
        ),
        CreateResult(
            content="Sorry, the search failed.",
            usage=RequestUsage(prompt_tokens=90, completion_tokens=40),
        ),
    ])

    tracker = _AutoGenTracker(mock_visibe, model_client, "test-error", agents=[])
    tracker.start()

    await model_client.create([UserMessage(content="Search for news")])
    await model_client.create([
        FunctionExecutionResultMessage(content=[
            FunctionExecutionResult(
                content="ConnectionError: Failed",
                name="web_search",
                call_id="call_err",
                is_error=True,
            )
        ])
    ])

    tracker.stop()

    trace = mock_visibe.traces[0]
    assert len(trace["tools_used"]) == 1
    assert trace["tools_used"][0]["status"] == "failed"
    assert len(trace["errors"]) == 1
    assert trace["errors"][0]["error_type"] == "TOOL_ERROR"
    assert "ConnectionError" in trace["errors"][0]["message"]
    assert trace["status"] == "failed"


@pytest.mark.asyncio
async def test_multiple_tool_calls():
    from visibe.integrations.autogen import _AutoGenTracker

    mock_visibe = MockVisibe()
    model_client = MockModelClient([
        CreateResult(
            content=[
                FunctionCall(id="call_a", name="get_weather", arguments='{"city": "NYC"}'),
                FunctionCall(id="call_b", name="get_time", arguments='{"timezone": "EST"}'),
            ],
            usage=RequestUsage(prompt_tokens=60, completion_tokens=35),
        ),
        CreateResult(
            content="It's 72F in NYC at 3:00 PM.",
            usage=RequestUsage(prompt_tokens=120, completion_tokens=25),
        ),
    ])

    tracker = _AutoGenTracker(mock_visibe, model_client, "test-multi-tools", agents=[])
    tracker.start()

    await model_client.create([UserMessage(content="Weather and time in NYC?")])
    await model_client.create([
        FunctionExecutionResultMessage(content=[
            FunctionExecutionResult(content="72F, sunny", name="get_weather", call_id="call_a"),
            FunctionExecutionResult(content="3:00 PM", name="get_time", call_id="call_b"),
        ])
    ])

    tracker.stop()

    trace = mock_visibe.traces[0]
    assert len(trace["tools_used"]) == 2
    tool_names = {t["tool_name"] for t in trace["tools_used"]}
    assert tool_names == {"get_weather", "get_time"}
    assert all(t["status"] == "success" for t in trace["tools_used"])


@pytest.mark.asyncio
async def test_llm_error_tracking():
    from visibe.integrations.autogen import _AutoGenTracker

    mock_visibe = MockVisibe()

    class ErrorModelClient(MockModelClient):
        async def create(self, messages, *args, **kwargs):
            raise RuntimeError("API rate limit exceeded")

    model_client = ErrorModelClient([])
    tracker = _AutoGenTracker(mock_visibe, model_client, "test-llm-error", agents=[])
    tracker.start()

    try:
        await model_client.create([UserMessage(content="Hello")])
    except RuntimeError:
        pass

    tracker.stop()

    trace = mock_visibe.traces[0]
    assert trace["status"] == "failed"
    assert len(trace["errors"]) == 1
    assert trace["errors"][0]["error_type"] == "RuntimeError"
    assert "rate limit" in trace["errors"][0]["message"]


@pytest.mark.asyncio
async def test_fallback_usage_tracking():
    from visibe.integrations.autogen import _AutoGenTracker

    mock_visibe = MockVisibe()
    model_client = MockModelClient([])

    tracker = _AutoGenTracker(mock_visibe, model_client, "test-fallback", agents=[])
    tracker.start()

    # Simulate calls happening outside the wrapper (edge case)
    model_client._total_prompt = 100
    model_client._total_completion = 200

    tracker.stop()

    trace = mock_visibe.traces[0]
    assert trace["total_input_tokens"] == 100
    assert trace["total_output_tokens"] == 200
    assert trace["total_cost"] > 0


@pytest.mark.asyncio
async def test_prompt_text_captured():
    from visibe.integrations.autogen import _AutoGenTracker

    mock_visibe = MockVisibe()
    model_client = MockModelClient([
        CreateResult(
            content="Paris is the capital of France.",
            usage=RequestUsage(prompt_tokens=20, completion_tokens=10),
        )
    ])

    tracker = _AutoGenTracker(mock_visibe, model_client, "test-text", agents=[])
    tracker.start()
    await model_client.create([UserMessage(content="What is the capital of France?")])
    tracker.stop()

    trace = mock_visibe.traces[0]
    assert "France" in trace["prompt"]

    last_llm_step = [s for s in trace["metadata"]["steps"] if s["type"] == "llm_call"][-1]
    assert "Paris" in last_llm_step["output_text"]


@pytest.mark.asyncio
async def test_dynamic_agent_names():
    from visibe.integrations.autogen import _AutoGenTracker

    class FakeAgent:
        name = "research_assistant"

    mock_visibe = MockVisibe()
    model_client = MockModelClient([
        CreateResult(
            content="Done!",
            usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
        )
    ])

    tracker = _AutoGenTracker(
        mock_visibe, model_client, "test-agent-names",
        agents=[FakeAgent()]
    )
    tracker.start()
    await model_client.create([UserMessage(content="Hello")])
    tracker.stop()

    trace = mock_visibe.traces[0]
    assert trace["agents_executed"] == ["research_assistant"]
    assert trace["agent_cost_breakdown"][0]["agent_name"] == "research_assistant"
    assert trace["metadata"]["steps"][0]["agent_name"] == "research_assistant"
    assert trace["task_cost_breakdown"][0]["agent_name"] == "research_assistant"
