from __future__ import annotations

"""Runtime agent implementation and loop strategies."""

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

from agentrspy.core import (
    AgentError,
    AgentOutput,
    AgentTrait,
    CompletionRequest,
    CompletionResponse,
    DoneEvent,
    LlmProvider,
    Memory,
    Message,
    ThinkingEvent,
    TokenEvent,
    ToolOutput,
    Usage,
)
from agentrspy.tools import ToolRegistry


class LoopStrategy:
    """Built-in loop strategies."""

    def __init__(
        self, kind: str, *, max_steps: int | None = None, instruction: str | None = None
    ) -> None:
        self.kind = kind
        self.max_steps = max_steps
        self.instruction = instruction

    @classmethod
    def react(cls, max_steps: int = 8) -> "LoopStrategy":
        return cls("react", max_steps=max_steps)

    @classmethod
    def cot(cls) -> "LoopStrategy":
        return cls("cot")

    @classmethod
    def plan_and_execute(cls, max_steps: int = 8) -> "LoopStrategy":
        return cls("plan_and_execute", max_steps=max_steps)

    @classmethod
    def custom(cls, instruction: str) -> "LoopStrategy":
        return cls("custom", instruction=instruction)

    def max_steps_hint(self, fallback: int) -> int:
        return (
            self.max_steps
            if self.kind in {"react", "plan_and_execute"} and self.max_steps is not None
            else fallback
        )


LoopStrategy.ReAct = LoopStrategy.react  # type: ignore[attr-defined]
LoopStrategy.CoT = LoopStrategy.cot()  # type: ignore[attr-defined]
LoopStrategy.PlanAndExecute = LoopStrategy.plan_and_execute  # type: ignore[attr-defined]
LoopStrategy.Custom = LoopStrategy.custom  # type: ignore[attr-defined]


@dataclass(slots=True)
class AgentConfig:
    """Runtime agent configuration."""

    model: str = ""
    temperature: float | None = 0.2
    max_tokens: int | None = 4096
    loop_strategy: LoopStrategy = field(default_factory=lambda: LoopStrategy.react(8))
    max_steps: int = 8


@dataclass(slots=True)
class AgentRunner(AgentTrait):
    """Runnable agent implementation."""

    llm: LlmProvider
    memory: Memory
    tools: ToolRegistry
    system_prompt: str | None
    config: AgentConfig

    async def run(self, input_text: str) -> AgentOutput:
        strategy = self.config.loop_strategy
        if strategy.kind == "react":
            return await self._run_react(input_text)
        if strategy.kind == "cot":
            return await self._run_cot(input_text)
        if strategy.kind == "plan_and_execute":
            return await self._run_plan_execute(
                input_text, strategy.max_steps or self.config.max_steps
            )
        if strategy.kind == "custom":
            return await self._run_cot(
                f"{strategy.instruction}\n\nUser task: {input_text}"
            )
        raise AgentError.invalid_configuration(
            f"unknown loop strategy: {strategy.kind}"
        )

    async def stream_run(self, input_text: str) -> AsyncIterator[Any]:
        output = await self.run(input_text)

        async def iterator() -> AsyncIterator[Any]:
            yield ThinkingEvent("completed")
            for token in output.text.split():
                yield TokenEvent(f"{token} ")
            yield DoneEvent(output)

        return iterator()

    async def _run_react(self, input_text: str) -> AgentOutput:
        await self.memory.store("user", Message.user(input_text))
        max_steps = self.config.loop_strategy.max_steps or self.config.max_steps
        for step in range(1, max_steps + 1):
            history = await self.memory.history()
            response = await self.llm.complete(
                self._build_request(history, include_tools=not self.tools.is_empty())
            )
            assistant_message = response.message
            await self.memory.store("assistant", assistant_message)
            if assistant_message.tool_calls:
                for message in await self._execute_tool_calls(
                    assistant_message.tool_calls
                ):
                    await self.memory.store("tool", message)
                continue
            return await self._finish_output(response, step)
        raise AgentError.max_steps_reached(max_steps)

    async def _run_cot(self, input_text: str) -> AgentOutput:
        await self.memory.store("user", Message.user(input_text))
        history = await self.memory.history()
        response = await self.llm.complete(
            self._build_request(history, include_tools=False)
        )
        await self.memory.store("assistant", response.message)
        return await self._finish_output(response, 1)

    async def _run_plan_execute(self, input_text: str, max_steps: int) -> AgentOutput:
        planner_prompt = (
            f"Create a concise numbered plan to solve the user task. Task: {input_text}"
        )
        plan_response = await self.llm.complete(
            CompletionRequest(
                messages=[Message.user(planner_prompt)],
                tools=None,
                model=self.config.model,
                temperature=0.1,
                max_tokens=self.config.max_tokens,
                stream=False,
                system=self.system_prompt,
                extra={},
            )
        )
        await self.memory.store(
            "plan", Message.assistant(plan_response.message.text_content())
        )
        execution_prompt = f"Use this plan to solve the task.\nPlan:\n{plan_response.message.text_content()}\n\nTask: {input_text}"
        await self.memory.store("user", Message.user(execution_prompt))
        output = await self._run_react(input_text)
        output.steps = max(output.steps, min(max_steps, max(output.steps, 1)))
        return output

    def _build_request(
        self, history: list[Message], include_tools: bool
    ) -> CompletionRequest:
        return CompletionRequest(
            messages=history,
            tools=self.tools.to_definitions() if include_tools else None,
            model=self.config.model,
            temperature=self.config.temperature,
            max_tokens=self.config.max_tokens,
            stream=False,
            system=self.system_prompt,
            extra={},
        )

    async def _execute_tool_calls(self, tool_calls: list[Any]) -> list[Message]:
        async def invoke(tool_call: Any) -> Message:
            try:
                output = await self.tools.call(tool_call.name, tool_call.arguments)
            except Exception as error:
                output = ToolOutput.error(str(error))
            return Message.tool_result(tool_call.id, tool_call.name, output)

        return await asyncio.gather(*(invoke(tool_call) for tool_call in tool_calls))

    async def _finish_output(
        self, response: CompletionResponse, steps: int
    ) -> AgentOutput:
        history = await self.memory.history()
        return AgentOutput(
            text=response.message.text_content(),
            steps=steps,
            usage=response.usage or Usage(),
            messages=history,
            metadata={},
        )
