from __future__ import annotations

import asyncio
from typing import Any, Awaitable, Callable

from ..providers.base import ProviderError, is_transient_provider_error
from ..schemas import ChatMessage, TELEGRAM_ACTIVE_USER_ID_METADATA_KEY, ToolCall, ToolExecutionContext
from .prompting import build_system_prompt

ProgressCallback = Callable[[dict[str, Any]], Awaitable[None] | None]


class AgentLoop:
    def __init__(self, runtime):
        self.runtime = runtime

    def _provider_retry_attempts(self) -> int:
        raw = self.runtime.config.get("provider_retry_attempts", 3)
        try:
            return max(int(raw), 1)
        except (TypeError, ValueError):
            return 3

    def _provider_retry_base_delay_sec(self) -> float:
        raw = self.runtime.config.get("provider_retry_base_delay_sec", 1.0)
        try:
            return max(float(raw), 0.0)
        except (TypeError, ValueError):
            return 1.0

    async def _emit(self, callback: ProgressCallback | None, payload: dict[str, Any]) -> None:
        if callback is None:
            return
        result = callback(payload)
        if result is not None and hasattr(result, "__await__"):
            await result

    def _resolve_owner(self, *, transport: str, user_id: int | None, chat_id: str) -> bool:
        if transport != "telegram":
            return True
        if hasattr(self.runtime, "user_is_admin"):
            return bool(self.runtime.user_is_admin(user_id=user_id, chat_id=chat_id))
        if hasattr(self.runtime, "is_authorized"):
            return bool(self.runtime.is_authorized(user_id, chat_id))
        if hasattr(self.runtime, "auth_is_allowed"):
            return bool(self.runtime.auth_is_allowed(user_id=user_id, chat_id=chat_id))
        if hasattr(self.runtime, "is_owner"):
            return bool(self.runtime.is_owner(user_id, chat_id))
        return False

    def _update_transport_session_state(
        self,
        *,
        session,
        transport: str,
        user_id: int | None,
        active: bool,
    ) -> None:
        if transport != "telegram":
            return
        if active:
            session.metadata[TELEGRAM_ACTIVE_USER_ID_METADATA_KEY] = user_id
            return
        session.metadata.pop(TELEGRAM_ACTIVE_USER_ID_METADATA_KEY, None)

    async def run_turn(
        self,
        *,
        session_id: str,
        chat_id: str,
        user_id: int | None,
        text: str,
        transport: str = "local",
        transport_state: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        session = self.runtime.get_session(session_id=session_id, chat_id=chat_id, user_id=user_id)
        if session.closed:
            session.closed = False
            self.runtime.session_store.save_state(session)

        self.runtime.sync_skill_tools()
        system_prompt = build_system_prompt(
            self.runtime.config,
            session,
            provider_names=self.runtime.list_provider_names(),
            skill_fragments=self.runtime.skill_manager.get_system_prompt_fragments(),
            root_agents_prompt=(self.runtime.read_root_agents_md() if hasattr(self.runtime, "read_root_agents_md") else ""),
        )
        tool_defs = self.runtime.tool_registry.list()
        tool_schemas = [tool.openai_schema() for tool in tool_defs]
        conversation = [message.to_dict() for message in session.history]
        conversation.append({"role": "user", "content": text})

        final_text = ""
        last_result: dict[str, Any] = {}
        progress_callback = transport_state.get("progress_callback") if transport_state else None
        max_steps = self.runtime.max_agent_steps()
        iteration = 0

        await self._emit(
            progress_callback,
            {
                "event": "agent_started",
                "session_id": session.session_id,
                "chat_id": chat_id,
                "user_id": user_id,
                "text": text,
                "max_steps": max_steps,
                "security_mode": session.security_mode,
            },
        )

        async def generate_with_retry() -> dict[str, Any]:
            attempts = self._provider_retry_attempts()
            base_delay_sec = self._provider_retry_base_delay_sec()
            last_error: ProviderError | None = None
            for attempt in range(1, attempts + 1):
                try:
                    provider = self.runtime.provider_registry.get(session.provider)
                    return await provider.generate(
                        messages=conversation,
                        tools=tool_schemas,
                        model=session.model,
                        system_prompt=system_prompt,
                    )
                except ProviderError as exc:
                    last_error = exc
                    can_retry = attempt < attempts and is_transient_provider_error(exc)
                    await self._emit(
                        progress_callback,
                        {
                            "event": "provider_generate_error",
                            "iteration": iteration,
                            "attempt": attempt,
                            "max_attempts": attempts,
                            "provider": session.provider,
                            "model": session.model,
                            "error": str(exc),
                            "will_retry": can_retry,
                        },
                    )
                    if not can_retry:
                        raise
                    delay_sec = base_delay_sec * (2 ** (attempt - 1))
                    await self._emit(
                        progress_callback,
                        {
                            "event": "provider_retry_scheduled",
                            "iteration": iteration,
                            "attempt": attempt,
                            "max_attempts": attempts,
                            "provider": session.provider,
                            "model": session.model,
                            "error": str(exc),
                            "delay_sec": delay_sec,
                        },
                    )
                    if delay_sec > 0:
                        await asyncio.sleep(delay_sec)
            if last_error is not None:
                raise last_error
            raise ProviderError("provider generate failed without an explicit error")

        while True:
            iteration += 1
            if max_steps is not None and iteration > max_steps:
                final_text = f"Agent step limit reached: {max_steps}"
                await self._emit(
                    progress_callback,
                    {
                        "event": "agent_step_limit_reached",
                        "iteration": iteration,
                        "limit": max_steps,
                    },
                )
                break

            await self._emit(
                progress_callback,
                {
                    "event": "agent_iteration_started",
                    "iteration": iteration,
                    "session_id": session.session_id,
                    "provider": session.provider,
                    "model": session.model,
                },
            )
            try:
                result = await generate_with_retry()
            except ProviderError as exc:
                final_text = (
                    f"Провайдер `{session.provider}` не готов: {exc}. "
                    "Настрой ключ/модель через Telegram `/settings`, `/apikey <provider> <key>`, "
                    "`/provider <name>`, `/model <name>` или inline-кнопки настроек."
                )
                last_result = {
                    "text": final_text,
                    "tool_calls": [],
                    "usage": {},
                    "provider": session.provider,
                    "model": session.model,
                    "raw": {"error": str(exc)},
                }
                await self._emit(
                    progress_callback,
                    {
                        "event": "assistant_response",
                        "iteration": iteration,
                        "text": final_text,
                        "tool_calls": [],
                        "usage": {},
                    },
                )
                break
            last_result = result
            await self._emit(
                progress_callback,
                {
                    "event": "assistant_response",
                    "iteration": iteration,
                    "text": result.get("text", ""),
                    "tool_calls": result.get("tool_calls") or [],
                    "usage": result.get("usage") or {},
                },
            )

            tool_calls = [
                ToolCall(
                    id=str(item.get("id") or f"call-{index + 1}"),
                    name=str(item.get("name") or "unknown"),
                    arguments=item.get("arguments") or {},
                )
                for index, item in enumerate(result.get("tool_calls") or [])
            ]
            assistant_entry: dict[str, Any] = {"role": "assistant", "content": result.get("text", "")}
            if tool_calls:
                assistant_entry["tool_calls"] = [
                    {"id": call.id, "name": call.name, "arguments": call.arguments} for call in tool_calls
                ]
            conversation.append(assistant_entry)

            if not tool_calls:
                final_text = str(result.get("text") or "").strip() or "Готово."
                break

            for call in tool_calls:
                await self._emit(
                    progress_callback,
                    {
                        "event": "tool_call_started",
                        "iteration": iteration,
                        "tool": call.name,
                        "arguments": call.arguments,
                    },
                )
                tool_context = ToolExecutionContext(
                    runtime=self.runtime,
                    session=session,
                    config=self.runtime.config,
                    chat_id=chat_id,
                    user_id=user_id,
                    is_owner=self._resolve_owner(transport=transport, user_id=user_id, chat_id=chat_id),
                    transport=transport,
                    transport_state=transport_state or {},
                )
                try:
                    tool_result = await self.runtime.tool_registry.execute(call.name, call.arguments, tool_context)
                except Exception as exc:
                    tool_result = {
                        "error": str(exc),
                        "tool": call.name,
                        "recoverable": True,
                    }
                    await self._emit(
                        progress_callback,
                        {
                            "event": "tool_call_error",
                            "iteration": iteration,
                            "tool": call.name,
                            "error": str(exc),
                        },
                    )
                await self._emit(
                    progress_callback,
                    {
                        "event": "tool_call_finished",
                        "iteration": iteration,
                        "tool": call.name,
                        "result": tool_result,
                    },
                )
                conversation.append(
                    {
                        "role": "tool",
                        "tool_call_id": call.id,
                        "name": call.name,
                        "content": tool_result,
                    }
                )
                if call.name == "session_end":
                    session.closed = True
                    self.runtime.session_store.save_state(session)

            if session.closed:
                final_text = str(last_result.get("text") or "Сессия закрыта агентом.")
                break

        self._update_transport_session_state(
            session=session,
            transport=transport,
            user_id=user_id,
            active=not session.closed,
        )
        self.runtime.session_store.append_messages(
            session,
            [ChatMessage(role="user", content=text), ChatMessage(role="assistant", content=final_text)],
        )
        await self._emit(
            progress_callback,
            {
                "event": "agent_finished",
                "iteration": iteration,
                "text": final_text,
                "session_closed": session.closed,
                "usage": last_result.get("usage") or {},
            },
        )
        return {
            "text": final_text,
            "usage": last_result.get("usage") or {},
            "provider": session.provider,
            "model": session.model,
            "security_mode": session.security_mode,
            "tool_count": len(last_result.get("tool_calls") or []),
            "session_closed": session.closed,
            "iterations": iteration,
        }
