"""Async agent implementation for concurrent operations."""

from __future__ import annotations

import asyncio
import uuid
from dataclasses import dataclass, field
from typing import Any, AsyncIterator

from agenai.agent.config import AgentConfig
from agenai.agent.context import AgentContext
from agenai.agent.core import AgentResult, StepResult, StepType
from agenai.agent.hooks import HookData, HookEvent, HookRegistry
from agenai.agent.llm import LLMProvider, LLMResponse
from agenai.agent.memory import ConversationMemory
from agenai.agent.retry import RetryConfig
from agenai.agent.tools import ToolRegistry, ToolResult


@dataclass
class AsyncRetryConfig:
    """Configuration for async retry behavior."""

    max_retries: int = 3
    base_delay: float = 1.0
    max_delay: float = 60.0
    exponential_base: float = 2.0


class AsyncLLMProvider:
    """Async wrapper for LLM providers."""

    def __init__(self, sync_provider: LLMProvider):
        self._sync = sync_provider

    async def chat(
        self,
        messages: list[dict[str, Any]],
        system: str | None = None,
        tools: list[dict[str, Any]] | None = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
    ) -> LLMResponse:
        """Async chat completion using thread pool."""
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            None,
            lambda: self._sync.chat(
                messages=messages,
                system=system,
                tools=tools,
                max_tokens=max_tokens,
                temperature=temperature,
            ),
        )


@dataclass
class StreamChunk:
    """A chunk of streamed output."""

    type: str  # "text", "tool_start", "tool_end", "done"
    content: str = ""
    tool_name: str | None = None
    tool_result: ToolResult | None = None


class AsyncAgent:
    """Async version of the Agent for concurrent operations.

    Provides:
    - Async run() method
    - Streaming via run_stream()
    - Concurrent tool execution
    """

    def __init__(
        self,
        config: AgentConfig | None = None,
        llm: LLMProvider | None = None,
        tools: ToolRegistry | None = None,
        memory: ConversationMemory | None = None,
        hooks: HookRegistry | None = None,
        retry_config: AsyncRetryConfig | None = None,
    ):
        self.config = config or AgentConfig()
        self.tools = tools or ToolRegistry()
        self.memory = memory or ConversationMemory()
        self.hooks = hooks or HookRegistry()
        self.retry_config = retry_config or AsyncRetryConfig()

        # Wrap sync LLM provider
        if llm:
            self.llm = AsyncLLMProvider(llm)
        else:
            from agenai.agent.llm import AnthropicProvider

            self.llm = AsyncLLMProvider(AnthropicProvider(model=self.config.model))

        self._context: AgentContext | None = None

    async def run(
        self,
        user_input: str,
        context: AgentContext | None = None,
    ) -> AgentResult:
        """Run the agent asynchronously."""
        self._context = context or AgentContext(
            run_id=str(uuid.uuid4()),
            user_input=user_input,
            metadata={"model": self.config.model},
        )

        await self._trigger_hook_async(HookEvent.AGENT_START)
        self.memory.add_user(user_input)

        steps: list[StepResult] = []
        total_tokens = 0

        try:
            for step_num in range(self.config.max_steps):
                self._context.step = step_num + 1
                await self._trigger_hook_async(HookEvent.STEP_START, step_number=step_num + 1)

                step_result, response = await self._step()
                steps.append(step_result)

                if response:
                    total_tokens += response.usage.get("input_tokens", 0)
                    total_tokens += response.usage.get("output_tokens", 0)

                await self._trigger_hook_async(HookEvent.STEP_END, step_number=step_num + 1)

                if step_result.step_type == StepType.FINAL_ANSWER:
                    await self._trigger_hook_async(HookEvent.AGENT_END)
                    return AgentResult(
                        success=True,
                        output=step_result.content,
                        steps=steps,
                        total_tokens=total_tokens,
                        context=self._context,
                    )

                if step_result.step_type == StepType.ERROR:
                    await self._trigger_hook_async(
                        HookEvent.AGENT_ERROR,
                        error=Exception(step_result.content),
                    )
                    return AgentResult(
                        success=False,
                        output="",
                        steps=steps,
                        error=step_result.content,
                        total_tokens=total_tokens,
                        context=self._context,
                    )

            # Max steps reached
            await self._trigger_hook_async(HookEvent.AGENT_END)
            return AgentResult(
                success=False,
                output="",
                steps=steps,
                error=f"Max steps ({self.config.max_steps}) reached",
                total_tokens=total_tokens,
                context=self._context,
            )

        except Exception as e:
            await self._trigger_hook_async(HookEvent.AGENT_ERROR, error=e)
            return AgentResult(
                success=False,
                output="",
                steps=steps,
                error=str(e),
                total_tokens=total_tokens,
                context=self._context,
            )

    async def run_stream(
        self,
        user_input: str,
        context: AgentContext | None = None,
    ) -> AsyncIterator[StreamChunk]:
        """Run the agent with streaming output."""
        self._context = context or AgentContext(
            run_id=str(uuid.uuid4()),
            user_input=user_input,
        )

        self.memory.add_user(user_input)

        for step_num in range(self.config.max_steps):
            self._context.step = step_num + 1

            try:
                step_result, response = await self._step()

                if step_result.thinking:
                    yield StreamChunk(type="text", content=step_result.thinking)

                if step_result.step_type == StepType.TOOL_CALL:
                    yield StreamChunk(
                        type="tool_end",
                        tool_name=step_result.tool_name,
                        tool_result=step_result.tool_result,
                    )

                if step_result.step_type == StepType.FINAL_ANSWER:
                    yield StreamChunk(type="text", content=step_result.content)
                    yield StreamChunk(type="done")
                    return

                if step_result.step_type == StepType.ERROR:
                    yield StreamChunk(type="error", content=step_result.content)
                    return

            except Exception as e:
                yield StreamChunk(type="error", content=str(e))
                return

        yield StreamChunk(type="error", content="Max steps reached")

    async def _step(self) -> tuple[StepResult, LLMResponse | None]:
        """Execute a single async step."""
        try:
            messages = self.memory.get_anthropic_messages()
            tools = self.tools.list() if len(self.tools) > 0 else None

            await self._trigger_hook_async(HookEvent.LLM_START, messages=messages)

            response = await self._call_llm_with_retry(messages, tools)

            await self._trigger_hook_async(HookEvent.LLM_END, response=response)

            if response.is_final:
                text = response.text or ""
                self.memory.add_assistant(text)
                return StepResult(step_type=StepType.FINAL_ANSWER, content=text), response

            return await self._handle_tool_calls(response), response

        except Exception as e:
            return StepResult(step_type=StepType.ERROR, content=str(e)), None

    async def _call_llm_with_retry(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None,
    ) -> LLMResponse:
        """Call LLM with async retry logic."""
        last_error: Exception | None = None

        for attempt in range(self.retry_config.max_retries + 1):
            try:
                return await self.llm.chat(
                    messages=messages,
                    system=self.config.system_prompt,
                    tools=tools,
                    max_tokens=self.config.max_tokens,
                    temperature=self.config.temperature,
                )
            except Exception as e:
                last_error = e
                if attempt < self.retry_config.max_retries:
                    delay = min(
                        self.retry_config.base_delay
                        * (self.retry_config.exponential_base**attempt),
                        self.retry_config.max_delay,
                    )
                    await asyncio.sleep(delay)

        raise last_error  # type: ignore

    async def _handle_tool_calls(self, response: LLMResponse) -> StepResult:
        """Handle tool calls with concurrent execution."""
        thinking_text = response.text

        # Build assistant message
        tool_call_blocks = []
        if thinking_text:
            tool_call_blocks.append({"type": "text", "text": thinking_text})

        for tc in response.tool_calls:
            tool_call_blocks.append({
                "type": "tool_use",
                "id": tc.id,
                "name": tc.name,
                "input": tc.arguments,
            })

        self.memory.add_assistant_tool_calls(tool_call_blocks)

        # Execute tools concurrently
        async def execute_tool(tc) -> tuple[str, dict, ToolResult]:
            await self._trigger_hook_async(
                HookEvent.TOOL_START,
                tool_name=tc.name,
                tool_args=tc.arguments,
            )

            if tc.name not in self.tools:
                result = ToolResult.fail(f"Unknown tool: {tc.name}")
            else:
                # Run sync tool in thread pool
                loop = asyncio.get_event_loop()
                try:
                    result = await loop.run_in_executor(
                        None,
                        lambda: self.tools.execute(tc.name, **tc.arguments),
                    )
                except Exception as e:
                    result = ToolResult.fail(f"Tool error: {e}")

            await self._trigger_hook_async(
                HookEvent.TOOL_END,
                tool_name=tc.name,
                tool_result=result,
            )

            # Add to memory
            result_text = result.output if result.success else f"Error: {result.error}"
            self.memory.add_tool_result(tc.id, tc.name, result_text)

            return tc.name, tc.arguments, result

        # Run all tools concurrently
        tasks = [execute_tool(tc) for tc in response.tool_calls]
        results = await asyncio.gather(*tasks)

        # Build step result
        result_texts = [f"{name}: {r.output if r.success else r.error}" for name, _, r in results]
        last_name, last_args, last_result = results[-1] if results else (None, None, None)

        return StepResult(
            step_type=StepType.TOOL_CALL,
            content="\n".join(result_texts),
            tool_name=last_name,
            tool_args=last_args,
            tool_result=last_result,
            thinking=thinking_text,
        )

    async def _trigger_hook_async(self, event: HookEvent, **kwargs: Any) -> None:
        """Trigger hook (runs sync hooks in thread pool)."""
        data = HookData(event=event, context=self._context, **kwargs)
        loop = asyncio.get_event_loop()
        await loop.run_in_executor(None, lambda: self.hooks.trigger(data))

    def reset(self) -> None:
        """Reset agent state."""
        self.memory.clear()
        self._context = None

    def register_tool(self, tool: Any) -> None:
        """Register a tool."""
        self.tools.register(tool)

    def register_hook(self, hook: Any) -> None:
        """Register a hook."""
        self.hooks.register(hook)
