"""
Concrete Agent implementation with core execution loop.
"""

import asyncio
import os
import time
import uuid
import warnings
from collections.abc import AsyncGenerator
from typing import List, Optional, Union

from agentbyte.cancellation_token import CancellationToken
from agentbyte.context import AgentContext
from agentbyte.llm.types import ChatCompletionChunk
from agentbyte.messages import (
    AssistantMessage,
    Message,
    ToolCallRequest,
    ToolMessage,
    Usage,
    UserMessage,
)
from agentbyte.types import ToolResult
from agentbyte.middleware.base import OperationType
from agentbyte.tools import ApprovalMode

from .base import AgentConfigurationError, AgentExecutionError, BaseAgent
from .types import (
    AgentEvent,
    AgentResponse,
    BaseEvent,
    ErrorEvent,
    ModelCallEvent,
    ModelResponseEvent,
    ModelStreamChunkEvent,
    TaskCompleteEvent,
    TaskStartEvent,
    ToolApprovalEvent,
    ToolCallEvent,
    ToolCallResponseEvent,
)


class Agent(BaseAgent):
    """Standard single-agent execution loop with tools and cancellation."""

    async def _run_internal(
        self,
        task: Optional[Union[str, UserMessage, List[Message]]],
        working_context: AgentContext,
        cancellation_token: Optional[CancellationToken] = None,
    ) -> AgentResponse:
        final_response: Optional[AgentResponse] = None

        async for item in self.run_stream(
            task=task,
            context=working_context,
            cancellation_token=cancellation_token,
            verbose=False,
            stream_tokens=False,
        ):
            if isinstance(item, AgentResponse):
                final_response = item

        if final_response is None:
            raise AgentExecutionError("run_stream did not yield AgentResponse")

        return final_response

    async def run(
        self,
        task: Optional[Union[str, UserMessage, List[Message]]] = None,
        context: Optional[AgentContext] = None,
        cancellation_token: Optional[CancellationToken] = None,
    ) -> AgentResponse:
        working_context = context if context is not None else AgentContext()
        if "run_id" not in working_context.metadata:
            working_context.metadata["run_id"] = str(uuid.uuid4())
        if not task and working_context.is_empty:
            raise AgentConfigurationError(
                "Either 'task' must be provided or 'context' must be non-empty."
            )
        return await self._run_internal(task, working_context, cancellation_token)

    async def _execute_tool_call(
        self,
        tool_call: ToolCallRequest,
        working_context: AgentContext,
        cancellation_token: Optional[CancellationToken],
        verbose: bool,
        skip_approval_check: bool = False,
    ) -> AsyncGenerator[Union[Message, AgentEvent], None]:
        if verbose:
            yield ToolCallEvent(
                source=self.name,
                tool_name=tool_call.tool_name,
                parameters=tool_call.parameters,
                call_id=tool_call.call_id,
            )

        tool = self._find_tool(tool_call.tool_name)
        if tool is None:
            message = ToolMessage(
                content=f"Tool '{tool_call.tool_name}' not found",
                source=self.name,
                tool_call_id=tool_call.call_id,
                tool_name=tool_call.tool_name,
                success=False,
                error="Tool not found",
            )
            working_context.add_message(message)
            if verbose:
                yield ToolCallResponseEvent(
                    source=self.name,
                    call_id=tool_call.call_id,
                    result=None,
                )
            yield message
            return

        if cancellation_token and cancellation_token.is_cancelled():
            raise asyncio.CancelledError()

        if tool.approval_mode == ApprovalMode.ALWAYS and not skip_approval_check:
            existing_approval = working_context.get_approval_response(tool_call.call_id)

            if existing_approval is None:
                approval_request = working_context.add_approval_request(
                    tool_call=tool_call,
                    tool_name=tool_call.tool_name,
                )
                yield ToolApprovalEvent(
                    source=self.name,
                    approval_request=approval_request,
                )
                return

            consumed_approval = working_context.consume_approval_response(
                tool_call.call_id
            )
            if consumed_approval is not None and not consumed_approval.approved:
                message = ToolMessage(
                    content=(
                        "Tool execution denied: "
                        f"{consumed_approval.reason or 'User declined approval'}"
                    ),
                    source=self.name,
                    tool_call_id=tool_call.call_id,
                    tool_name=tool_call.tool_name,
                    success=False,
                    error="Approval denied",
                )

                working_context.add_message(message)

                if verbose:
                    yield ToolCallResponseEvent(
                        source=self.name,
                        call_id=tool_call.call_id,
                        result=None,
                    )

                yield message
                return

        if tool.supports_streaming():
            try:

                async def _execute_streaming_tool(
                    payload: dict[str, object],
                ) -> AsyncGenerator[object, None]:
                    tool_parameters = payload.get("parameters", {})
                    if not isinstance(tool_parameters, dict):
                        raise ValueError(
                            "tool_call payload parameters must be a dictionary"
                        )

                    async for item in tool.execute_stream(
                        tool_parameters,
                        cancellation_token,
                    ):
                        yield item

                streamed_result: object | None = None
                async for item in self.middleware_chain.execute_tool_stream(
                    operation=OperationType.TOOL_CALL,
                    agent_name=self.name,
                    agent_context=working_context,
                    data={
                        "tool_name": tool_call.tool_name,
                        "parameters": tool_call.parameters,
                        "call_id": tool_call.call_id,
                    },
                    func=_execute_streaming_tool,
                ):
                    if isinstance(item, BaseEvent):
                        if verbose or isinstance(item, ToolApprovalEvent):
                            yield item

                        if isinstance(item, ToolApprovalEvent):
                            return

                        continue

                    if isinstance(item, ToolResult):
                        streamed_result = item
                    else:
                        yield item

                if streamed_result is None:
                    raise AgentExecutionError(
                        "Streaming tool middleware execution did not yield a final result"
                    )

                tool_message = ToolMessage(
                    content=str(streamed_result.result)
                    if streamed_result.success
                    else (streamed_result.error or ""),
                    source=self.name,
                    tool_call_id=tool_call.call_id,
                    tool_name=tool_call.tool_name,
                    success=streamed_result.success,
                    error=streamed_result.error,
                )
                working_context.add_message(tool_message)
                yield tool_message

                if verbose:
                    yield ToolCallResponseEvent(
                        source=self.name,
                        call_id=tool_call.call_id,
                        result=streamed_result,
                    )

                return
            except Exception as exc:
                error_message = str(exc)
                if error_message == "Streaming tool returned no final result":
                    error_message = "Streaming tool middleware execution did not yield a final result"

                message = ToolMessage(
                    content=f"Tool execution failed: {error_message}",
                    source=self.name,
                    tool_call_id=tool_call.call_id,
                    tool_name=tool_call.tool_name,
                    success=False,
                    error=error_message,
                )
                working_context.add_message(message)

                if verbose:
                    yield ToolCallResponseEvent(
                        source=self.name,
                        call_id=tool_call.call_id,
                        result=None,
                    )

                yield message
                return

        try:

            async def _execute_tool(payload: dict[str, object]) -> object:
                tool_parameters = payload.get("parameters", {})
                if not isinstance(tool_parameters, dict):
                    raise ValueError(
                        "tool_call payload parameters must be a dictionary"
                    )
                return await tool.execute(tool_parameters)

            result: object | None = None
            async for item in self.middleware_chain.execute_stream(
                operation=OperationType.TOOL_CALL,
                agent_name=self.name,
                agent_context=working_context,
                data={
                    "tool_name": tool_call.tool_name,
                    "parameters": tool_call.parameters,
                    "call_id": tool_call.call_id,
                },
                func=_execute_tool,
            ):
                if isinstance(item, BaseEvent):
                    if verbose or isinstance(item, ToolApprovalEvent):
                        yield item

                    if isinstance(item, ToolApprovalEvent):
                        return

                    continue

                result = item

            if result is None:
                raise AgentExecutionError(
                    "Tool middleware execution did not yield a final result"
                )

            content = str(result.result) if result.success else (result.error or "")
            message = ToolMessage(
                content=content,
                source=self.name,
                tool_call_id=tool_call.call_id,
                tool_name=tool_call.tool_name,
                success=result.success,
                error=result.error,
            )
        except Exception as exc:
            message = ToolMessage(
                content=f"Tool execution failed: {exc}",
                source=self.name,
                tool_call_id=tool_call.call_id,
                tool_name=tool_call.tool_name,
                success=False,
                error=str(exc),
            )

        working_context.add_message(message)

        if verbose:
            yield ToolCallResponseEvent(
                source=self.name,
                call_id=tool_call.call_id,
                result=result if "result" in locals() else None,
            )

        yield message

    async def run_stream(
        self,
        task: Optional[Union[str, UserMessage, List[Message]]] = None,
        context: Optional[AgentContext] = None,
        cancellation_token: Optional[CancellationToken] = None,
        verbose: bool = False,
        stream_tokens: bool = False,
    ) -> AsyncGenerator[
        Union[Message, AgentEvent, AgentResponse, ChatCompletionChunk], None
    ]:
        """Stream agent execution events, ending with a final AgentResponse.

        Creates an OTel root span (``agent <name>``) when ``AGENTBYTE_ENABLE_OTEL=true``
        so that all per-call ``chat`` and ``tool`` child spans are grouped under one trace.
        """
        working_context = context if context is not None else AgentContext()
        if not task and working_context.is_empty:
            raise AgentConfigurationError(
                "Either 'task' must be provided or 'context' must be non-empty."
            )
        if os.getenv("AGENTBYTE_ENABLE_OTEL", "false").lower() in ("true", "1", "yes"):
            try:
                from opentelemetry import trace

                tracer = trace.get_tracer("agentbyte")
                with tracer.start_as_current_span(f"agent {self.name}"):
                    async for item in self._run_stream_impl(
                        task=task,
                        working_context=working_context,
                        cancellation_token=cancellation_token,
                        verbose=verbose,
                        stream_tokens=stream_tokens,
                    ):
                        yield item
                return
            except Exception:
                pass
        async for item in self._run_stream_impl(
            task=task,
            working_context=working_context,
            cancellation_token=cancellation_token,
            verbose=verbose,
            stream_tokens=stream_tokens,
        ):
            yield item

    async def _run_stream_impl(
        self,
        task: Optional[Union[str, UserMessage, List[Message]]],
        working_context: AgentContext,
        cancellation_token: Optional[CancellationToken] = None,
        verbose: bool = False,
        stream_tokens: bool = False,
    ) -> AsyncGenerator[
        Union[Message, AgentEvent, AgentResponse, ChatCompletionChunk], None
    ]:
        """Core streaming execution loop."""
        start_time = time.time()
        llm_calls = 0
        tokens_input = 0
        tokens_output = 0
        tool_calls = 0
        memory_operations = 0
        cost_estimate = 0.0
        finish_reason = "stop"

        effective_stream_tokens = stream_tokens
        if stream_tokens and self.middleware_chain.middlewares:
            warnings.warn(
                "stream_tokens=True is disabled when middleware is configured; "
                "falling back to non-token middleware-compatible execution.",
                stacklevel=2,
            )
            effective_stream_tokens = False

        try:
            task_messages = self._convert_task_to_messages(task) if task else []
            for message in task_messages:
                working_context.add_message(message)

            if task_messages:
                yield task_messages[0]
                if verbose:
                    yield TaskStartEvent(
                        source=self.name, task=task_messages[0].content
                    )

            iteration = 0
            while iteration < self.max_iterations:
                if cancellation_token and cancellation_token.is_cancelled():
                    raise asyncio.CancelledError()

                approved_calls = working_context.get_approved_tool_calls()
                rejected_calls = working_context.get_rejected_tool_calls()

                for approved_call in approved_calls:
                    tool_calls += 1
                    async for item in self._execute_tool_call(
                        approved_call,
                        working_context,
                        cancellation_token,
                        verbose,
                        skip_approval_check=True,
                    ):
                        yield item

                for rejected_call_id, rejected_call in rejected_calls:
                    denial_message = ToolMessage(
                        content="Tool execution denied: User declined approval",
                        source=self.name,
                        tool_call_id=rejected_call_id,
                        tool_name=rejected_call.tool_name,
                        success=False,
                        error="Approval denied",
                    )
                    working_context.add_message(denial_message)
                    yield denial_message

                llm_messages = await self._prepare_llm_messages([], working_context)
                if self.memory is not None:
                    memory_operations += 1
                tools = self._get_tools_for_llm() if self.tools else None

                if verbose:
                    yield ModelCallEvent(
                        source=self.name,
                        model=self.model_client.model,
                        input_messages=llm_messages,
                    )

                if effective_stream_tokens and not tools:
                    accumulated_content = ""
                    last_usage: Optional[Usage] = None

                    async for chunk in self.model_client.create_stream(
                        llm_messages,
                        tools=None,
                        output_format=self.output_format,
                    ):
                        if cancellation_token and cancellation_token.is_cancelled():
                            raise asyncio.CancelledError()

                        if chunk.usage is not None:
                            last_usage = chunk.usage

                        if not chunk.is_complete:
                            accumulated_content += chunk.content
                            if verbose:
                                yield ModelStreamChunkEvent(
                                    source=self.name,
                                    chunk=chunk.content,
                                    is_final=False,
                                )
                            yield chunk

                    if verbose:
                        yield ModelStreamChunkEvent(
                            source=self.name,
                            chunk="",
                            is_final=True,
                        )

                    assistant_message = AssistantMessage(
                        content=accumulated_content,
                        source=self.name,
                        tool_calls=None,
                        usage=last_usage,
                    )

                    usage_for_count = last_usage or Usage()
                    llm_calls += 1
                    tokens_input += usage_for_count.tokens_input
                    tokens_output += usage_for_count.tokens_output
                    if usage_for_count.cost_estimate is not None:
                        cost_estimate += usage_for_count.cost_estimate

                    working_context.add_message(assistant_message)
                    yield assistant_message
                    finish_reason = "stop"
                    break

                async def _execute_model_call(payload: dict[str, object]):
                    payload_messages = payload.get("messages", llm_messages)
                    if not isinstance(payload_messages, list):
                        raise ValueError("model_call payload messages must be a list")

                    payload_tools = payload.get("tools", tools)
                    if payload_tools is not None and not isinstance(
                        payload_tools, list
                    ):
                        raise ValueError(
                            "model_call payload tools must be a list or None"
                        )

                    payload_output_format = payload.get(
                        "output_format", self.output_format
                    )
                    if payload_output_format is not None and not isinstance(
                        payload_output_format, type
                    ):
                        raise ValueError(
                            "model_call payload output_format must be a type or None"
                        )

                    return await self.model_client.create(
                        payload_messages,
                        tools=payload_tools,
                        output_format=payload_output_format,
                    )

                completion_result = None
                middleware_paused = False

                async for middleware_item in self.middleware_chain.execute_stream(
                    operation=OperationType.MODEL_CALL,
                    agent_name=self.name,
                    agent_context=working_context,
                    data={
                        "messages": llm_messages,
                        "tools": tools,
                        "output_format": self.output_format,
                    },
                    metadata={
                        "messages_count": len(llm_messages),
                        "has_tools": bool(tools),
                        "output_format": self.output_format.__name__
                        if self.output_format is not None
                        else None,
                        "model": getattr(self.model_client, "model", "unknown"),
                    },
                    func=_execute_model_call,
                ):
                    if isinstance(middleware_item, BaseEvent):
                        if verbose or isinstance(middleware_item, ToolApprovalEvent):
                            yield middleware_item

                        if isinstance(middleware_item, ToolApprovalEvent):
                            finish_reason = "approval_needed"
                            middleware_paused = True
                            break

                        continue

                    completion_result = middleware_item

                if middleware_paused:
                    break

                if completion_result is None:
                    raise AgentExecutionError(
                        "Model middleware execution did not yield a completion result",
                    )

                llm_calls += 1
                tokens_input += completion_result.usage.tokens_input
                tokens_output += completion_result.usage.tokens_output
                if completion_result.usage.cost_estimate is not None:
                    cost_estimate += completion_result.usage.cost_estimate

                assistant_message = AssistantMessage(
                    content=completion_result.message.content,
                    source=self.name,
                    tool_calls=completion_result.message.tool_calls,
                    structured_content=completion_result.message.structured_content,
                    usage=completion_result.usage,
                )

                if verbose:
                    yield ModelResponseEvent(
                        source=self.name,
                        model=completion_result.model,
                        response=completion_result.message.content,
                        has_tool_calls=bool(completion_result.message.tool_calls),
                    )

                working_context.add_message(assistant_message)
                yield assistant_message

                if not assistant_message.tool_calls:
                    finish_reason = completion_result.finish_reason or "stop"
                    break

                approval_pause = False
                for tool_call in assistant_message.tool_calls:
                    tool_calls += 1
                    async for item in self._execute_tool_call(
                        tool_call,
                        working_context,
                        cancellation_token,
                        verbose,
                    ):
                        yield item

                        if isinstance(item, ToolApprovalEvent):
                            finish_reason = "approval_needed"
                            approval_pause = True
                            break

                    if approval_pause:
                        break

                if approval_pause:
                    break

                if not self.summarize_tool_result:
                    finish_reason = "tool_executed"
                    break

                iteration += 1

            if iteration >= self.max_iterations and finish_reason == "stop":
                finish_reason = "max_iterations"

        except asyncio.CancelledError:
            finish_reason = "cancelled"
            if verbose:
                yield ErrorEvent(
                    source=self.name,
                    error_message="Execution cancelled",
                    error_type="CancelledError",
                    is_recoverable=True,
                )
        except Exception as exc:
            finish_reason = "error"
            error_message = AssistantMessage(
                content=f"Error: {exc}",
                source=self.name,
            )
            working_context.add_message(error_message)
            yield error_message
            if verbose:
                yield ErrorEvent(
                    source=self.name,
                    error_message=str(exc),
                    error_type=type(exc).__name__,
                    is_recoverable=False,
                )

        duration_ms = int((time.time() - start_time) * 1000)
        usage = Usage(
            duration_ms=duration_ms,
            llm_calls=llm_calls,
            tokens_input=tokens_input,
            tokens_output=tokens_output,
            tool_calls=tool_calls,
            memory_operations=memory_operations,
            cost_estimate=cost_estimate if cost_estimate > 0 else None,
        )

        response = AgentResponse(
            context=working_context,
            source=self.name,
            finish_reason=finish_reason,
            usage=usage,
        )

        if verbose:
            yield TaskCompleteEvent(
                source=self.name,
                result=response.final_content,
                finish_reason=finish_reason,
            )

        yield response


__all__ = ["Agent"]
