"""Tests for AgentAsTool behavior, metadata, and result strategies."""

import sys
from collections.abc import AsyncGenerator
from types import ModuleType
from typing import Any, Dict, List, Optional, Union

import pytest

from agentbyte.agents import Agent, AgentAsTool
from agentbyte.agents.types import AgentEvent, AgentResponse
from agentbyte.context import AgentContext
from agentbyte.llm.base import BaseChatCompletionClient, BaseChatCompletionClientConfig
from agentbyte.llm.types import ChatCompletionChunk, ChatCompletionResult
from agentbyte.messages import (
    AssistantMessage,
    Message,
    ToolCallRequest,
    Usage,
    UserMessage,
)
from agentbyte.middleware.otel import OTelMiddleware
from agentbyte.types import ToolResult


class FakeClientConfig(BaseChatCompletionClientConfig):
    pass


class FakeModelClient(BaseChatCompletionClient):
    def __init__(self, response: ChatCompletionResult) -> None:
        super().__init__(model="fake", client=object())
        self._response = response

    async def create(
        self,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        **kwargs: Any,
    ) -> ChatCompletionResult:
        return self._response

    async def create_stream(
        self,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        **kwargs: Any,
    ) -> AsyncGenerator[ChatCompletionChunk, None]:
        yield ChatCompletionChunk(content="x", model="fake", is_complete=False)
        yield ChatCompletionChunk(content="", model="fake", is_complete=True)

    def _to_config(self) -> BaseChatCompletionClientConfig:
        return FakeClientConfig(model=self.model, config=self.config)

    @classmethod
    def _from_config(cls, config: BaseChatCompletionClientConfig) -> "FakeModelClient":
        response = ChatCompletionResult(
            message=AssistantMessage(content="ok", source="fake"),
            usage=Usage(),
            model="fake",
        )
        return cls(response=response)


class SequencedFakeModelClient(BaseChatCompletionClient):
    def __init__(self, responses: List[ChatCompletionResult]) -> None:
        super().__init__(model="fake", client=object())
        self._responses = responses
        self._index = 0

    async def create(
        self,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        **kwargs: Any,
    ) -> ChatCompletionResult:
        response = self._responses[self._index]
        self._index += 1
        return response

    async def create_stream(
        self,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        **kwargs: Any,
    ) -> AsyncGenerator[ChatCompletionChunk, None]:
        yield ChatCompletionChunk(content="x", model="fake", is_complete=False)
        yield ChatCompletionChunk(content="", model="fake", is_complete=True)

    def _to_config(self) -> BaseChatCompletionClientConfig:
        return FakeClientConfig(model=self.model, config=self.config)

    @classmethod
    def _from_config(
        cls, config: BaseChatCompletionClientConfig
    ) -> "SequencedFakeModelClient":
        return cls(responses=[])


class _FakeStatusCode:
    OK = "OK"
    ERROR = "ERROR"


class _FakeStatus:
    def __init__(self, code: str, description: str | None = None) -> None:
        self.code = code
        self.description = description


class _FakeSpan:
    def __init__(self, name: str) -> None:
        self.name = name
        self.attributes: dict[str, Any] = {}
        self.status: _FakeStatus | None = None
        self.recorded_exceptions: list[Exception] = []
        self.parent_name: str | None = None
        self.ended = False

    def set_attribute(self, key: str, value: Any) -> None:
        self.attributes[key] = value

    def set_status(self, status: _FakeStatus) -> None:
        self.status = status

    def record_exception(self, exc: Exception) -> None:
        self.recorded_exceptions.append(exc)

    def end(self) -> None:
        self.ended = True


class _FakeSpanContextManager:
    def __init__(self, span: _FakeSpan, active_stack: list[_FakeSpan]) -> None:
        self.span = span
        self._active_stack = active_stack

    def __enter__(self) -> _FakeSpan:
        if self._active_stack:
            self.span.parent_name = self._active_stack[-1].name
        self._active_stack.append(self.span)
        return self.span

    def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> bool:
        self._active_stack.pop()
        self.span.end()
        return False


class _FakeTracer:
    def __init__(self, spans: list[_FakeSpan], active_stack: list[_FakeSpan]) -> None:
        self._spans = spans
        self._active_stack = active_stack

    def start_as_current_span(self, name: str) -> _FakeSpanContextManager:
        span = _FakeSpan(name)
        self._spans.append(span)
        return _FakeSpanContextManager(span, self._active_stack)

    def start_span(self, name: str) -> _FakeSpan:
        span = _FakeSpan(name)
        self._spans.append(span)
        return span


class _FakeOtelContext:
    def __init__(self, active_stack: list[_FakeSpan]) -> None:
        self._active_stack = active_stack

    def attach(self, span: _FakeSpan) -> _FakeSpan:
        if self._active_stack:
            span.parent_name = self._active_stack[-1].name
        self._active_stack.append(span)
        return span

    def detach(self, token: _FakeSpan) -> None:
        popped = self._active_stack.pop()
        assert popped is token


def _install_fake_opentelemetry(
    monkeypatch: pytest.MonkeyPatch, spans: list[_FakeSpan]
) -> _FakeTracer:
    active_stack: list[_FakeSpan] = []
    tracer = _FakeTracer(spans, active_stack)
    fake_context = _FakeOtelContext(active_stack)

    trace_module = ModuleType("opentelemetry.trace")
    trace_module.get_tracer = lambda name: tracer  # type: ignore[attr-defined]
    trace_module.set_span_in_context = lambda span: span  # type: ignore[attr-defined]
    trace_module.Status = _FakeStatus  # type: ignore[attr-defined]
    trace_module.StatusCode = _FakeStatusCode  # type: ignore[attr-defined]

    context_module = ModuleType("opentelemetry.context")
    context_module.attach = fake_context.attach  # type: ignore[attr-defined]
    context_module.detach = fake_context.detach  # type: ignore[attr-defined]

    otel_module = ModuleType("opentelemetry")
    otel_module.trace = trace_module  # type: ignore[attr-defined]
    otel_module.context = context_module  # type: ignore[attr-defined]

    monkeypatch.setitem(sys.modules, "opentelemetry", otel_module)
    monkeypatch.setitem(sys.modules, "opentelemetry.trace", trace_module)
    monkeypatch.setitem(sys.modules, "opentelemetry.context", context_module)

    import agentbyte.middleware.otel as otel_module_ref

    monkeypatch.setattr(otel_module_ref, "trace", trace_module)
    monkeypatch.setattr(otel_module_ref, "otel_context", context_module)
    monkeypatch.setattr(otel_module_ref, "Status", _FakeStatus)
    monkeypatch.setattr(otel_module_ref, "StatusCode", _FakeStatusCode)

    return tracer


def _make_fake_otel_middleware(tracer: _FakeTracer) -> OTelMiddleware:
    middleware = OTelMiddleware.__new__(OTelMiddleware)
    middleware._enabled = True
    middleware._capture_content = False
    middleware._tracer = tracer
    middleware._meter = None
    middleware._token_histogram = None
    middleware._duration_histogram = None
    return middleware


class MockAgent:
    """Simple protocol-compatible mock for AgentAsTool strategy tests."""

    def __init__(
        self, name: str, messages: List[str], finish_reason: str = "stop"
    ) -> None:
        self.name = name
        self.description = f"{name} description"
        self._messages = messages
        self._finish_reason = finish_reason

    async def run(
        self,
        task: Optional[Union[UserMessage, List[Message]]] = None,
        context: Optional[AgentContext] = None,
        cancellation_token: Optional[Any] = None,
    ) -> AgentResponse:
        ctx = context or AgentContext()
        # For strategy tests, we only include the assistant responses
        for content in self._messages:
            ctx.add_message(AssistantMessage(content=content, source=self.name))

        return AgentResponse(
            context=ctx,
            source=self.name,
            finish_reason=self._finish_reason,
            usage=Usage(tokens_input=11, tokens_output=7, duration_ms=5),
        )

    async def run_stream(
        self,
        task: Optional[Union[UserMessage, List[Message]]] = None,
        context: Optional[AgentContext] = None,
        cancellation_token: Optional[Any] = None,
        verbose: bool = False,
        stream_tokens: bool = False,
    ) -> AsyncGenerator[Union[Message, AgentEvent, AgentResponse], None]:
        for content in self._messages:
            yield AssistantMessage(content=content, source=self.name)

        yield await self.run(
            task=task, context=context, cancellation_token=cancellation_token
        )


@pytest.mark.asyncio
async def test_agent_as_tool_execute_last_strategy() -> None:
    response = ChatCompletionResult(
        message=AssistantMessage(content="final answer", source="fake"),
        usage=Usage(tokens_input=2, tokens_output=2),
        model="fake",
    )
    client = FakeModelClient(response)
    agent = Agent(
        name="researcher",
        description="research",
        instructions="do research",
        model_client=client,
    )

    tool_wrapper = AgentAsTool(agent=agent)
    result = await tool_wrapper.execute({"task": "find X"})

    assert result.success is True
    assert result.result == "final answer"
    assert result.metadata["agent_name"] == "researcher"
    assert result.metadata["finish_reason"] == "stop"
    assert result.metadata["usage"]["tokens_input"] == 2
    assert result.metadata["usage"]["tokens_output"] == 2


@pytest.mark.asyncio
async def test_agent_as_tool_last_n_strategy() -> None:
    response = ChatCompletionResult(
        message=AssistantMessage(content="step 2", source="fake"),
        usage=Usage(tokens_input=2, tokens_output=2),
        model="fake",
    )
    client = FakeModelClient(response)
    agent = Agent(
        name="planner",
        description="plan",
        instructions="plan",
        model_client=client,
    )

    # Seed extra history to validate concatenation behavior.
    ctx = AgentContext()
    ctx.add_message(AssistantMessage(content="step 0", source="planner"))
    ctx.add_message(AssistantMessage(content="step 1", source="planner"))

    wrapper = AgentAsTool(agent=agent, result_strategy="last:2")
    result = await wrapper.execute({"task": "plan"}, context=ctx)

    assert result.success is True
    assert "plan" in result.result
    assert "step 2" in result.result


@pytest.mark.asyncio
async def test_agent_as_tool_all_strategy_uses_all_messages() -> None:
    tool_wrapper = AgentAsTool(
        agent=MockAgent(name="specialist", messages=["one", "two", "three"]),
        result_strategy="all",
    )

    result = await tool_wrapper.execute({"task": "aggregate"})

    assert result.success is True
    assert result.result == "one\ntwo\nthree"


@pytest.mark.asyncio
async def test_agent_as_tool_custom_callable_strategy() -> None:
    def extract_middle(messages: List[Message]) -> str:
        # Use messages[1:] assuming first might be user message or skip first assistant
        return " | ".join(message.content.upper() for message in messages[1:])

    tool_wrapper = AgentAsTool(
        agent=MockAgent(name="specialist", messages=["alpha", "beta", "gamma"]),
        result_strategy=extract_middle,
    )

    result = await tool_wrapper.execute({"task": "transform"})

    assert result.success is True
    assert result.result == "BETA | GAMMA"


@pytest.mark.asyncio
async def test_agent_as_tool_supports_custom_task_parameter_name() -> None:
    tool_wrapper = AgentAsTool(
        agent=MockAgent(name="worker", messages=["done"]),
        task_parameter_name="job",
    )

    assert tool_wrapper.parameters["required"] == ["job"]
    result = await tool_wrapper.execute({"job": "do this"})

    assert result.success is True
    assert result.result == "done"


@pytest.mark.asyncio
async def test_base_agent_as_tool_passes_task_parameter_name() -> None:
    response = ChatCompletionResult(
        message=AssistantMessage(content="completed", source="fake"),
        usage=Usage(tokens_input=1, tokens_output=1),
        model="fake",
    )
    agent = Agent(
        name="delegate",
        description="delegates",
        instructions="delegate",
        model_client=FakeModelClient(response),
    )

    wrapped = agent.as_tool(task_parameter_name="job")
    result = await wrapped.execute({"job": "perform"})

    assert wrapped.parameters["required"] == ["job"]
    assert result.success is True
    assert result.result == "completed"


@pytest.mark.asyncio
async def test_agent_as_tool_stream_emits_usage_and_finish_reason_metadata() -> None:
    tool_wrapper = AgentAsTool(
        agent=MockAgent(name="streamer", messages=["first", "second"]),
        result_strategy="last",
    )

    final_result: ToolResult | None = None
    async for item in tool_wrapper.execute_stream({"task": "stream"}):
        if isinstance(item, ToolResult):
            final_result = item

    assert final_result is not None
    assert final_result.success is True
    assert final_result.result == "second"
    assert final_result.metadata["finish_reason"] == "stop"
    assert final_result.metadata["usage"]["tokens_input"] == 11
    assert final_result.metadata["usage"]["tokens_output"] == 7


@pytest.mark.asyncio
async def test_agent_as_tool_streaming_execution_preserves_tool_span_hierarchy(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    spans: list[_FakeSpan] = []
    tracer = _install_fake_opentelemetry(monkeypatch, spans)
    middleware = _make_fake_otel_middleware(tracer)

    monkeypatch.setenv("AGENTBYTE_ENABLE_OTEL", "true")

    specialist = Agent(
        name="specialist",
        description="Handles delegated research",
        instructions="Answer the delegated task clearly.",
        model_client=SequencedFakeModelClient(
            [
                ChatCompletionResult(
                    message=AssistantMessage(
                        content="specialist result",
                        source="fake",
                    ),
                    usage=Usage(
                        tokens_input=7,
                        tokens_output=3,
                        cost_estimate=0.0015,
                    ),
                    model="fake",
                )
            ]
        ),
        middlewares=[middleware],
    )

    coordinator = Agent(
        name="coordinator",
        description="Coordinates delegated work",
        instructions="Delegate to the specialist when needed.",
        model_client=SequencedFakeModelClient(
            [
                ChatCompletionResult(
                    message=AssistantMessage(
                        content="delegating",
                        source="fake",
                        tool_calls=[
                            ToolCallRequest(
                                tool_name="specialist",
                                parameters={"task": "Investigate the request"},
                                call_id="call-1",
                            )
                        ],
                    ),
                    usage=Usage(
                        tokens_input=12,
                        tokens_output=4,
                        cost_estimate=0.0025,
                    ),
                    model="fake",
                ),
                ChatCompletionResult(
                    message=AssistantMessage(
                        content="final answer",
                        source="fake",
                    ),
                    usage=Usage(
                        tokens_input=9,
                        tokens_output=5,
                        cost_estimate=0.0035,
                    ),
                    model="fake",
                ),
            ]
        ),
        tools=[specialist.as_tool()],
        middlewares=[middleware],
    )

    response = await coordinator.run("Please investigate and summarize.")

    assert response.messages[-1].content == "final answer"
    assert [span.name for span in spans] == [
        "agent coordinator",
        "chat fake",
        "tool specialist",
        "agent specialist",
        "chat fake",
        "chat fake",
    ]

    first_chat_span = spans[1]
    tool_span = spans[2]
    specialist_agent_span = spans[3]
    specialist_chat_span = spans[4]
    final_chat_span = spans[5]

    assert first_chat_span.parent_name == "agent coordinator"
    assert tool_span.parent_name == "agent coordinator"
    assert specialist_agent_span.parent_name == "tool specialist"
    assert specialist_chat_span.parent_name == "agent specialist"
    assert final_chat_span.parent_name == "agent coordinator"

    assert tool_span.attributes["gen_ai.tool.name"] == "specialist"
    assert tool_span.attributes["gen_ai.tool.success"] is True

    assert first_chat_span.attributes["gen_ai.usage.input_tokens"] == 12
    assert first_chat_span.attributes["gen_ai.usage.output_tokens"] == 4
    assert first_chat_span.attributes["gen_ai.usage.total_tokens"] == 16
    assert first_chat_span.attributes["gen_ai.usage.cost_estimate_usd"] == 0.0025

    assert specialist_chat_span.attributes["gen_ai.usage.input_tokens"] == 7
    assert specialist_chat_span.attributes["gen_ai.usage.output_tokens"] == 3
    assert specialist_chat_span.attributes["gen_ai.usage.total_tokens"] == 10
    assert specialist_chat_span.attributes["gen_ai.usage.cost_estimate_usd"] == 0.0015

    assert final_chat_span.attributes["gen_ai.usage.input_tokens"] == 9
    assert final_chat_span.attributes["gen_ai.usage.output_tokens"] == 5
    assert final_chat_span.attributes["gen_ai.usage.total_tokens"] == 14
    assert final_chat_span.attributes["gen_ai.usage.cost_estimate_usd"] == 0.0035


def test_agent_as_tool_invalid_strategy_raises() -> None:
    response = ChatCompletionResult(
        message=AssistantMessage(content="ok", source="fake"),
        usage=Usage(),
        model="fake",
    )
    agent = Agent(
        name="bad",
        description="bad",
        instructions="bad",
        model_client=FakeModelClient(response),
    )

    with pytest.raises(ValueError):
        AgentAsTool(agent=agent, result_strategy="unknown")


def test_agent_as_tool_invalid_last_n_raises() -> None:
    wrapper_agent = MockAgent(name="mock", messages=["x"])

    with pytest.raises(ValueError):
        AgentAsTool(agent=wrapper_agent, result_strategy="last:0")
