"""The support assistant: one customer message in, one reply out, with tool calls in between."""

from __future__ import annotations

import asyncio
import json
import logging
import time
from dataclasses import dataclass, field
from typing import Any, Protocol

import openai

from .language import language_instruction
from .memory import ConversationMemory
from .tools import TOOL_SCHEMAS, ToolContext, ToolExecutor

logger = logging.getLogger(__name__)

FALLBACK_ERROR = "سمحلي، صرا مشكل تقني. عاود جرب بعد شوية.\nDésolé, un problème technique est survenu. Réessayez dans un instant."
FALLBACK_BUSY = "كاين ضغط كبير دروك. عاود ابعث رسالتك بعد دقيقة.\nLe service est très sollicité. Renvoyez votre message dans une minute."
FALLBACK_EMPTY = "سمحلي، ما فهمتش مليح. تقدر تعاود تكتب سؤالك؟\nDésolé, je n'ai pas bien compris. Pouvez-vous reformuler ?"


class ChatCompletionsClient(Protocol):
    """The part of the OpenAI-compatible client this module uses."""

    @property
    def chat(self) -> Any: ...


@dataclass
class Reply:
    text: str
    tools_used: list[str] = field(default_factory=list)
    ticket_ids: list[int] = field(default_factory=list)
    failed: bool = False


class StoreAssistant:
    def __init__(
        self,
        client: ChatCompletionsClient,
        model: str,
        executor: ToolExecutor,
        memory: ConversationMemory,
        system_prompt: str,
        *,
        max_tool_rounds: int = 4,
        max_output_tokens: int = 8192,
    ) -> None:
        self._client = client
        self._model = model
        self._executor = executor
        self._memory = memory
        self._system_prompt = system_prompt
        self._max_tool_rounds = max_tool_rounds
        self._max_output_tokens = max_output_tokens

    def reset(self, chat_id: int) -> None:
        self._memory.clear(chat_id)

    async def reply(self, chat_id: int, user_text: str, username: str | None = None) -> Reply:
        ctx = ToolContext(chat_id=chat_id, username=username)
        messages: list[dict[str, Any]] = [
            {"role": "system", "content": self._system_prompt},
            *(dict(turn) for turn in self._memory.history(chat_id)),
        ]
        hint = language_instruction(user_text)
        if hint:
            messages.append({"role": "system", "content": hint})
        messages.append({"role": "user", "content": user_text})
        tools_used: list[str] = []
        started = time.monotonic()

        try:
            text = await self._run_loop(messages, ctx, tools_used)
        except openai.RateLimitError:
            logger.warning("chat=%s rate limited by the model API", chat_id)
            return Reply(FALLBACK_BUSY, tools_used, ctx.tickets, failed=True)
        except (openai.APIConnectionError, openai.APITimeoutError):
            logger.warning("chat=%s model API unreachable or timed out", chat_id)
            return Reply(FALLBACK_ERROR, tools_used, ctx.tickets, failed=True)
        except openai.APIStatusError as exc:
            logger.error("chat=%s model API returned HTTP %s", chat_id, exc.status_code)
            return Reply(FALLBACK_ERROR, tools_used, ctx.tickets, failed=True)

        if not text:
            text = FALLBACK_EMPTY
        else:
            self._memory.add_exchange(chat_id, user_text, text)

        logger.info(
            "chat=%s replied in %.1fs tools=%s tickets=%s",
            chat_id,
            time.monotonic() - started,
            ",".join(tools_used) or "-",
            ctx.tickets or "-",
        )
        return Reply(text, tools_used, ctx.tickets)

    async def _run_loop(self, messages: list[dict[str, Any]], ctx: ToolContext, tools_used: list[str]) -> str:
        cache: dict[tuple[str, str], str] = {}
        for round_no in range(self._max_tool_rounds + 1):
            # On the last round the model must answer with text.
            final_round = round_no == self._max_tool_rounds
            response = await self._client.chat.completions.create(
                model=self._model,
                messages=messages,
                tools=TOOL_SCHEMAS,
                tool_choice="none" if final_round else "auto",
                max_tokens=self._max_output_tokens,
            )
            if not response.choices:
                return ""
            message = response.choices[0].message
            tool_calls = message.tool_calls or []
            if not tool_calls or final_round:
                return (message.content or "").strip()

            if response.choices[0].finish_reason == "length":
                logger.warning("model hit the output limit while calling tools; consider raising MAX_OUTPUT_TOKENS")

            messages.append(_assistant_turn(message))
            for call in tool_calls:
                name = call.function.name
                key = (name, _canonical_arguments(call.function.arguments))
                if key in cache:
                    # Models sometimes repeat the same call; answer it from the first result
                    # (and never open a second ticket for the same request).
                    result = cache[key]
                else:
                    tools_used.append(name)
                    result = await asyncio.to_thread(self._executor.run, name, call.function.arguments, ctx)
                    cache[key] = result
                messages.append({"role": "tool", "tool_call_id": call.id, "content": result})
        return ""


def _canonical_arguments(raw: str | None) -> str:
    try:
        return json.dumps(json.loads(raw or "{}"), sort_keys=True, ensure_ascii=False)
    except json.JSONDecodeError:
        return raw or ""


def _assistant_turn(message: Any) -> dict[str, Any]:
    """Rebuild the assistant message that requested tools, to send it back with the results."""
    turn: dict[str, Any] = {
        "role": "assistant",
        "content": message.content or "",
        "tool_calls": [
            {
                "id": call.id,
                "type": "function",
                "function": {"name": call.function.name, "arguments": call.function.arguments or "{}"},
            }
            for call in message.tool_calls
        ],
    }
    # Reasoning models may return their reasoning next to the tool calls; some providers
    # expect it back in the same turn.
    reasoning = getattr(message, "reasoning_content", None)
    if reasoning:
        turn["reasoning_content"] = reasoning
    return turn