"""Provider-agnostic chat service wrapping LangChain."""

from __future__ import annotations

import re
from typing import Any

from langchain_core.language_models import BaseChatModel

# Regex to strip <think>...</think> blocks from reasoning models (e.g. Qwen3)
# Also handles unclosed <think> blocks (when output is truncated by max_tokens)
_THINK_TAG_RE = re.compile(r"<think>.*?</think>\s*|<think>.*", re.DOTALL)


def strip_thinking(content: str) -> str:
    """Strip <think>...</think> blocks from reasoning models."""
    return _THINK_TAG_RE.sub("", content).strip()


class ChatService:
    """Provider-agnostic chat service wrapping a LangChain chat model.

    Accepts any LangChain BaseChatModel (AzureChatOpenAI, ChatOpenAI, etc.).
    Use ``get_chat_service()`` from ``src.llm.factory`` to construct an instance
    with the correct provider backend.
    """

    def __init__(self, llm: BaseChatModel):
        self.llm = llm

    def invoke(self, messages: list[dict[str, str]], **kwargs: Any) -> str:
        response = self.llm.invoke(messages, **kwargs)
        return strip_thinking(response.content)

    async def ainvoke(self, messages: list[dict[str, str]], **kwargs: Any) -> str:
        response = await self.llm.ainvoke(messages, **kwargs)
        return strip_thinking(response.content)

    def stream(self, messages: list[dict[str, str]], **kwargs: Any):
        buffer = ""
        for chunk in self.llm.stream(messages, **kwargs):
            if chunk.content:
                buffer += chunk.content
        yield strip_thinking(buffer)

    async def astream(self, messages: list[dict[str, str]], **kwargs: Any):
        buffer = ""
        async for chunk in self.llm.astream(messages, **kwargs):
            if chunk.content:
                buffer += chunk.content
        yield strip_thinking(buffer)

    def get_llm(self) -> BaseChatModel:
        return self.llm

    def update_config(
        self,
        temperature: float | None = None,
        max_tokens: int | None = None,
        **kwargs: Any,
    ) -> None:
        if temperature is not None:
            self.llm.temperature = temperature
        if max_tokens is not None:
            self.llm.max_tokens = max_tokens
        for key, value in kwargs.items():
            if hasattr(self.llm, key):
                setattr(self.llm, key, value)