from __future__ import annotations

import inspect
import mimetypes
import re
import tempfile
from collections import OrderedDict
from collections.abc import Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from micro_agents import Aggregator, Media, OpenAIBasedLLM, UserMessage
from micro_agents_tools import BraveSearchTools, FilesystemTools, PythonTools

from app.agent_tools import AgentToolManager, TelegramTools
from app.config import AiConfig, AiProviderConfig
from app.llm import AnthropicMessagesLLM, ChatCompletionsLLM, OpenRouterChatLLM
from app.web_tools import WebPageTools


class AgentError(RuntimeError):
    pass


class AgentSyntaxError(AgentError):
    pass


@dataclass(frozen=True, slots=True)
class AgentInvocation:
    prompt: str
    context_messages: int = 0
    use_deepseek: bool = False
    use_claude: bool = False
    use_openrouter: bool = False
    thinking: bool = False


@dataclass(frozen=True, slots=True)
class PromptTemplates:
    master: str
    user: str


@dataclass(slots=True, eq=False)
class AgentSessionHandle:
    session: Any
    workspace: tempfile.TemporaryDirectory[str]
    response_keys: set[tuple[Any, Any]]

    def close(self) -> None:
        self.workspace.cleanup()


@dataclass(frozen=True, slots=True)
class AgentTurnResult:
    answer: str
    handle: AgentSessionHandle


class AgentSessionStore:
    def __init__(self, max_sessions: int = 100) -> None:
        if max_sessions < 1:
            raise ValueError("max_sessions must be positive")
        self._max_sessions = max_sessions
        self._by_response: dict[tuple[Any, Any], AgentSessionHandle] = {}
        self._handles: OrderedDict[AgentSessionHandle, None] = OrderedDict()

    def find(self, chat_id: Any, message_id: Any) -> AgentSessionHandle | None:
        handle = self._by_response.get((chat_id, message_id))
        if handle is not None and handle in self._handles:
            self._handles.move_to_end(handle)
        return handle

    def register_response(self, handle: AgentSessionHandle, message: Any) -> None:
        key = _message_key(message)
        if key[0] is None or key[1] is None:
            return
        self._handles[handle] = None
        self._handles.move_to_end(handle)
        handle.response_keys.add(key)
        self._by_response[key] = handle
        self._evict()

    def discard(self, handle: AgentSessionHandle) -> None:
        self._handles.pop(handle, None)
        for key in handle.response_keys:
            if self._by_response.get(key) is handle:
                self._by_response.pop(key, None)
        handle.response_keys.clear()
        handle.close()

    def clear(self) -> None:
        for handle in list(self._handles):
            self.discard(handle)

    def _evict(self) -> None:
        while len(self._handles) > self._max_sessions:
            oldest = next(iter(self._handles))
            self.discard(oldest)


_PLACEHOLDER = re.compile(
    r"\{\{\s*([A-Za-z_]\w*(?:\.\w+)*)\s*\}\}|\{([A-Za-z_]\w*(?:\.\w+)*)\}"
)
_BLOCK_COMMENT = re.compile(r"/\*.*?\*/", re.DOTALL)
_LINE_COMMENT = re.compile(r"(?m)(?<!\S)//.*$")


def strip_prompt_comments(value: str) -> str:
    without_blocks = _BLOCK_COMMENT.sub(
        lambda match: "\n" * match.group(0).count("\n"),
        value,
    )
    return _LINE_COMMENT.sub("", without_blocks)


def load_prompt_templates(prompt_dir: Path) -> PromptTemplates:
    try:
        master = strip_prompt_comments(
            (prompt_dir / "MASTER.txt").read_text(encoding="utf-8")
        )
        user = strip_prompt_comments(
            (prompt_dir / "USER.txt").read_text(encoding="utf-8")
        )
    except OSError as exc:
        raise AgentError(
            f"Could not load {prompt_dir / 'MASTER.txt'} and {prompt_dir / 'USER.txt'}"
        ) from exc
    if not master.strip() or not user.strip():
        raise AgentError("Prompt templates must not be empty")
    return PromptTemplates(master=master, user=user)


def _template_value(context: Mapping[str, Any], path: str) -> Any:
    value: Any = context
    for part in path.split("."):
        if not isinstance(value, Mapping) or part not in value:
            raise AgentError(f"Unknown prompt placeholder: {path}")
        value = value[part]
    return value


def render_prompt_template(template: str, context: Mapping[str, Any]) -> str:
    values = dict(context)
    if "source_text" in values:
        values["text"] = values["source_text"]

    def replace(match: re.Match[str]) -> str:
        path = match.group(1) or match.group(2)
        value = _template_value(values, path)
        if value is None:
            return ""
        if isinstance(value, bool):
            return str(value).lower()
        return str(value)

    rendered = _PLACEHOLDER.sub(replace, template).strip()
    if not rendered:
        raise AgentError("Rendered prompt must not be empty")
    return rendered


def parse_agent_invocation(value: str, config: AiConfig) -> AgentInvocation:
    tokens = value.split()
    if not tokens:
        raise AgentSyntaxError("AI prompt is required")

    index = 0
    context_messages = 0
    if tokens[0].lstrip("+-").isdigit():
        try:
            context_messages = int(tokens[0])
        except ValueError as exc:  # pragma: no cover - guarded by isdigit
            raise AgentSyntaxError("Invalid context message count") from exc
        if not 0 <= context_messages <= config.max_context_messages:
            raise AgentSyntaxError("Context message count is outside the configured range")
        index = 1

    provider = "openrouter"
    thinking = False
    while index < len(tokens) and tokens[index].startswith("."):
        flag = tokens[index].lower()
        if flag == ".glm":
            provider = "glm"
        elif flag == ".ds":
            provider = "deepseek"
        elif flag in {".cl", ".claude"}:
            provider = "claude"
        elif flag in {".or", ".openrouter"}:
            provider = "openrouter"
        elif flag == ".th":
            thinking = True
        index += 1

    prompt = " ".join(tokens[index:]).strip()
    if not prompt:
        raise AgentSyntaxError("AI prompt is required")

    return AgentInvocation(
        prompt=prompt,
        context_messages=context_messages,
        use_deepseek=provider == "deepseek",
        use_claude=provider == "claude",
        use_openrouter=provider == "openrouter",
        thinking=thinking,
    )


def _message_key(message: Any) -> tuple[Any, Any]:
    chat = getattr(message, "chat", None)
    return getattr(chat, "id", None), getattr(message, "id", None)


def _user_context(user: Any | None) -> dict[str, Any]:
    if user is None:
        return {"id": None, "name": "unknown", "username": None}
    return {
        "id": getattr(user, "id", None),
        "name": (
            getattr(user, "full_name", None)
            or getattr(user, "first_name", None)
            or getattr(user, "username", None)
            or "unknown"
        ),
        "username": getattr(user, "username", None),
    }


def _chat_context(chat: Any | None) -> dict[str, Any]:
    chat_type = getattr(chat, "type", None)
    return {
        "id": getattr(chat, "id", None),
        "title": getattr(chat, "title", None),
        "type": getattr(chat_type, "value", chat_type),
        "username": getattr(chat, "username", None),
    }


def _agent_context(command_message: Any, account: Any, invocation: AgentInvocation) -> dict[str, Any]:
    caller = getattr(command_message, "from_user", None)
    owner = getattr(account, "me", None)
    reply = getattr(command_message, "reply_to_message", None)
    account_config = getattr(account, "config", None)
    caller_id = getattr(caller, "id", None)
    owner_id = getattr(owner, "id", None)
    return {
        "caller_user": _user_context(caller),
        "owner": _user_context(owner),
        "is_owner": owner_id is not None and caller_id == owner_id,
        "text": invocation.prompt,
        "chat": _chat_context(getattr(command_message, "chat", None)),
        "message": _message_context(command_message, is_reply=False),
        "reply": {
            "id": getattr(reply, "id", None),
            "text": (
                getattr(reply, "text", None) or getattr(reply, "caption", None) or ""
                if reply is not None
                else ""
            ),
            "user": _user_context(getattr(reply, "from_user", None)),
        },
        "account": {
            "name": getattr(account_config, "name", None),
            "session": getattr(account_config, "session", None),
        },
        "provider": _provider_name(invocation),
        "thinking": invocation.thinking,
        "context_messages": invocation.context_messages,
    }


def _message_context(message: Any, *, is_reply: bool) -> dict[str, Any]:
    date = getattr(message, "date", None)
    return {
        "id": getattr(message, "id", None),
        "text": getattr(message, "text", None) or getattr(message, "caption", None) or "",
        "date": date.isoformat() if hasattr(date, "isoformat") else None,
        "user": _user_context(getattr(message, "from_user", None)),
        "is_reply": is_reply,
    }


def _media_object(message: Any) -> Any | None:
    for field_name in (
        "photo",
        "document",
        "audio",
        "voice",
        "video",
        "video_note",
        "animation",
        "sticker",
    ):
        value = getattr(message, field_name, None)
        if value is not None:
            return value
    return None


def _media_mime_type(media_object: Any, path: Path) -> str:
    configured = getattr(media_object, "mime_type", None)
    if configured:
        return str(configured)
    guessed = mimetypes.guess_type(path.name)[0]
    return guessed or "application/octet-stream"


def _media_kind(mime_type: str) -> str:
    if mime_type.startswith("image/"):
        return "image"
    if mime_type in {"audio/mpeg", "audio/wav", "audio/x-wav"}:
        return "audio"
    return "file"


async def _download_message_media(
    client: Any,
    message: Any,
    directory: Path,
    max_bytes: int,
) -> list[Media]:
    media_object = _media_object(message)
    if media_object is None:
        return []

    declared_size = getattr(media_object, "file_size", None)
    if not isinstance(declared_size, int) or declared_size > max_bytes:
        return []

    message_directory = directory / str(getattr(message, "id", "unknown"))
    message_directory.mkdir(parents=True, exist_ok=True)
    downloaded = await client.download_media(message, file_name=str(message_directory) + "/")
    if not isinstance(downloaded, str):
        raise AgentError("Telegram media download returned no file")

    path = Path(downloaded)
    if path.stat().st_size > max_bytes:
        return []
    mime_type = _media_mime_type(media_object, path)
    return [
        Media(
            kind=_media_kind(mime_type),
            data=path.read_bytes(),
            filename=getattr(media_object, "file_name", None) or path.name,
            mime_type=mime_type,
        )
    ]


async def _get_history(client: Any, command_message: Any, limit: int) -> list[Any]:
    if limit == 0:
        return []
    chat_id = getattr(getattr(command_message, "chat", None), "id", None)
    if chat_id is None:
        raise AgentError("Command message has no chat")

    history = client.get_chat_history(
        chat_id,
        limit=limit,
        offset_id=getattr(command_message, "id", 0),
    )
    if inspect.isawaitable(history):
        history = await history
    if hasattr(history, "__aiter__"):
        messages = [message async for message in history]
    else:
        messages = list(history)
    return list(reversed(messages))


async def _build_context_messages(
    client: Any,
    command_message: Any,
    invocation: AgentInvocation,
    config: AiConfig,
    directory: Path,
    *,
    user_template: str,
    template_context: Mapping[str, Any],
) -> list[UserMessage]:
    telegram_messages = await _get_history(client, command_message, invocation.context_messages)
    seen = {_message_key(message) for message in telegram_messages}
    reply = getattr(command_message, "reply_to_message", None)
    reply_key = _message_key(reply) if reply is not None else None
    if reply is not None and reply_key not in seen:
        telegram_messages.append(reply)

    messages: list[UserMessage] = []
    for telegram_message in telegram_messages:
        media = await _download_message_media(
            client,
            telegram_message,
            directory,
            config.max_media_bytes,
        )
        values = dict(template_context)
        values["text"] = (
            getattr(telegram_message, "text", None)
            or getattr(telegram_message, "caption", None)
            or "[media without a caption]"
        )
        values["message"] = _message_context(
            telegram_message,
            is_reply=_message_key(telegram_message) == reply_key,
        )
        messages.append(
            UserMessage(render_prompt_template(user_template, values), media)
        )
    return messages


def _provider_name(invocation: AgentInvocation) -> str:
    if invocation.use_openrouter:
        return "openrouter"
    if invocation.use_claude:
        return "claude"
    if invocation.use_deepseek:
        return "deepseek"
    return "glm"


def _provider(config: AiConfig, invocation: AgentInvocation) -> AiProviderConfig:
    if invocation.use_openrouter:
        return config.openrouter
    if invocation.use_claude:
        return config.claude
    if invocation.use_deepseek:
        return config.deepseek
    return config.glm


def _create_llm(provider: AiProviderConfig, invocation: AgentInvocation) -> Any:
    if invocation.use_openrouter:
        llm_class: type[Any] = OpenRouterChatLLM
    elif invocation.use_claude:
        llm_class = AnthropicMessagesLLM
    elif invocation.use_deepseek:
        llm_class = OpenAIBasedLLM
    else:
        llm_class = ChatCompletionsLLM
    return llm_class(
        api_url=provider.api_url,
        api_key=provider.api_key,
        model=provider.model,
    )


async def run_agent_turn(
    client: Any,
    command_message: Any,
    account: Any,
    session_store: AgentSessionStore,
    invocation: AgentInvocation,
    config: AiConfig,
) -> AgentTurnResult:
    chat_id = getattr(getattr(command_message, "chat", None), "id", None)
    if chat_id is None:
        raise AgentError("Command message has no chat")

    reply = getattr(command_message, "reply_to_message", None)
    existing = (
        session_store.find(chat_id, getattr(reply, "id", None))
        if reply is not None
        else None
    )
    if existing is not None:
        context = _agent_context(command_message, account, invocation)
        context["base_dir"] = config.base_dir
        existing.session.context.clear()
        existing.session.context.update(context)
        media = await _download_message_media(
            client,
            command_message,
            Path(existing.workspace.name),
            config.max_media_bytes,
        )
        return AgentTurnResult(
            answer=await _run_session(existing.session, invocation.prompt, media=media),
            handle=existing,
        )

    provider = _provider(config, invocation)
    if not provider.api_key:
        raise AgentError("The selected AI provider has no API key")

    llm = _create_llm(provider, invocation)
    templates = load_prompt_templates(config.prompt_dir)
    context = _agent_context(command_message, account, invocation)
    context["base_dir"] = config.base_dir

    workspace = tempfile.TemporaryDirectory(prefix="multi-bot-ai-")
    try:
        temp_path = Path(workspace.name)
        messages = await _build_context_messages(
            client,
            command_message,
            invocation,
            config,
            temp_path,
            user_template=templates.user,
            template_context=context,
        )
        media = await _download_message_media(
            client,
            command_message,
            temp_path,
            config.max_media_bytes,
        )
        tools = [
            *WebPageTools(
                timeout_seconds=config.web_page_timeout_seconds,
                max_chars=config.web_page_max_chars,
            ).tools(),
            *BraveSearchTools(
                api_key=config.brave_api_key,
                timeout_seconds=config.brave_timeout_seconds,
            ).tools(),
            *FilesystemTools(max_attachment_bytes=config.max_media_bytes).tools(),
            *PythonTools().tools(),
            *TelegramTools(
                client,
                chat_id,
                base_dir=config.base_dir,
                guest_base_dir=temp_path,
                reply_to_message_id=getattr(command_message, "id", None),
            ).tools(),
        ]
        agent = Aggregator(
            llm,
            AgentToolManager(tools),
            master_prompt=lambda values: render_prompt_template(
                load_prompt_templates(config.prompt_dir).master,
                values,
            ),
            user_prompt=lambda values: render_prompt_template(
                load_prompt_templates(config.prompt_dir).user,
                values,
            ),
            max_tool_rounds=config.max_tool_rounds,
            thinking=True if invocation.thinking else None,
        )
        session = agent.create_session(
            context=context,
            session_name=f"telegram:{chat_id}",
            messages=messages,
        )
        answer = await _run_session(session, invocation.prompt, media=media)
        return AgentTurnResult(
            answer=answer,
            handle=AgentSessionHandle(session, workspace, set()),
        )
    except Exception:
        workspace.cleanup()
        raise


async def _run_session(session: Any, prompt: str, *, media: list[Media] | None = None) -> str:
    if media:
        result = await session.run(prompt, media=media)
    else:
        result = await session.run(prompt)
    if not result.succeeded or result.final_message is None:
        raise AgentError("Agent loop failed") from result.error
    answer = result.final_message.message().strip()
    if not answer:
        raise AgentError("Agent returned an empty answer")
    return answer