"""Implement infrastructure support for openai compatible answer generator."""

from __future__ import annotations

import time
from typing import Generator

from app.config.settings import settings
from app.domain.conversation import AnswerGenerator, AnswerResult, AnswerSource
from app.domain.retrieval import RetrievedChunk
from app.services.llm.llm_factory import get_llm_client
from app.services.rag.prompt_templates import PromptTemplates
# Keep adapter behavior explicit so integration details remain easy to audit.


# Fallback system prompts used when no rich template matches.
_FALLBACK_PROMPTS = {
    "default": "你是法规知识问答助手。请仅依据提供的上下文回答；如果上下文不足，明确说明。",
    "compliance_qa": "你是法规合规问答助手。优先引用给定法规原文，回答要准确、克制，并注明依据来源。",
}


class OpenAICompatibleAnswerGenerator(AnswerGenerator):
    """Represent the Open A I Compatible Answer Generator type."""

    @staticmethod
    def _estimate_tokens(text: str) -> int:
        """Estimate token count for mixed Chinese/English text.

        Chinese chars are ~1.5 chars/token; ASCII is ~4 chars/token.
        """
        chinese = sum(1 for c in text if "一" <= c <= "鿿")
        other = len(text) - chinese
        return int(chinese / 1.5 + other / 4) + 1

    def _build_messages(
        self,
        *,
        query: str,
        retrieved_chunks: list[RetrievedChunk],
        history: list[dict[str, str]] | None,
        prompt_template: str | None,
        context_text: str | None = None,
        context_filename: str | None = None,
    ) -> tuple[list[dict[str, str]], int]:
        """Build the message list to send to the LLM.

        When context_text is provided the user's document is injected as a
        dedicated section BEFORE the retrieved regulation chunks so the LLM
        can reason about the document directly while still referencing regulations.
        The retrieval step uses only the user's question, not the document text,
        so embedding quality is preserved.

        System prompt selection priority:
          1. Rich template from PromptTemplates (compliance_qa / comparison /
             compliance_check / clause_interpretation / …)
          2. Fallback hardcoded prompt when no rich template matches.
        """
        # Look up the rich template first; fall back to simple hardcoded prompts.
        tpl_name = prompt_template or "compliance_qa"
        rich_tpl = PromptTemplates.get_template(tpl_name)
        if rich_tpl:
            system_prompt = rich_tpl.system_prompt
        else:
            system_prompt = _FALLBACK_PROMPTS.get(tpl_name, _FALLBACK_PROMPTS["default"])
        context_blocks = []
        context_tokens = 0

        # ── User document context (if attached) ───────────────────────────────
        if context_text and context_text.strip():
            doc_label = f"附件文档：{context_filename}" if context_filename else "附件文档"
            doc_block = f"[{doc_label}]\n{context_text.strip()}"
            doc_tokens = self._estimate_tokens(doc_block)
            # Reserve at most half the context budget for the user document
            half_budget = settings.rag_max_context_tokens // 2
            if doc_tokens > half_budget:
                # Truncate document to fit half the budget
                ratio = half_budget / doc_tokens
                doc_block = doc_block[: int(len(doc_block) * ratio)] + "\n…（文档已截断）"
                doc_tokens = half_budget
            context_blocks.append(doc_block)
            context_tokens += doc_tokens

        # ── Retrieved regulation chunks ────────────────────────────────────────
        remaining_budget = settings.rag_max_context_tokens - context_tokens
        for idx, chunk in enumerate(retrieved_chunks, start=1):
            block = (
                f"[法规{idx}] 文档: {chunk.doc_title}\n"
                f"章节: {chunk.section_title or '未标注'}\n"
                f"页码: {chunk.page_start}" + (f"-{chunk.page_end}" if chunk.page_end and chunk.page_end != chunk.page_start else "") + "\n"
                f"内容: {chunk.text}"
            )
            block_tokens = self._estimate_tokens(block)
            if block_tokens > remaining_budget:
                break
            remaining_budget -= block_tokens
            context_tokens += block_tokens
            context_blocks.append(block)

        context = "\n\n".join(context_blocks)
        messages = [{"role": "system", "content": system_prompt}]
        for item in history or []:
            messages.append({"role": item["role"], "content": item["content"]})

        # Craft the user turn differently when a document is attached
        if context_text and context_text.strip():
            user_content = (
                f"问题：{query}\n\n"
                f"请先基于上方附件文档内容进行分析，再结合法规参考上下文给出合规评估。"
                f"\n\n参考上下文：\n{context}\n\n"
                f"请在回答中注明引用来源编号（如适用）。"
            )
        else:
            user_content = f"问题：{query}\n\n参考上下文：\n{context}\n\n请在回答后给出简要引用编号。"

        messages.append({"role": "user", "content": user_content})
        return messages, context_tokens

    def _is_context_truncated(self, *, retrieved_chunks: list[RetrievedChunk], context_tokens: int) -> bool:
        """Return whether the prompt context had to omit retrieved chunks to fit the token budget."""
        if not retrieved_chunks:
            return False
        estimated_total_tokens = sum(
            self._estimate_tokens(
                f"[{idx}] 文档: {chunk.doc_title}\n"
                f"章节: {chunk.section_title or '未标注'}\n"
                f"页码: {chunk.page_start}" + (f"-{chunk.page_end}" if chunk.page_end and chunk.page_end != chunk.page_start else "") + "\n"
                f"内容: {chunk.text}"
            )
            for idx, chunk in enumerate(retrieved_chunks, start=1)
        )
        return estimated_total_tokens > context_tokens

    def _sources(self, chunks: list[RetrievedChunk]) -> list[AnswerSource]:
        """Handle sources for this module for the Open A I Compatible Answer Generator instance."""
        return [
            AnswerSource(
                doc_id=chunk.doc_id,
                doc_title=chunk.doc_title,
                chunk_id=chunk.chunk_id,
                chunk_type=chunk.chunk_type,
                section_title=chunk.section_title,
                page_start=chunk.page_start,
                page_end=chunk.page_end,
                section_level=chunk.section_level,
                chunk_index=chunk.chunk_index,
                piece_index=chunk.piece_index,
                score=chunk.score,
                text=chunk.text,
                metadata=chunk.metadata,
            )
            for chunk in chunks
        ]

    def generate(
        self,
        *,
        query: str,
        retrieved_chunks: list[RetrievedChunk],
        history: list[dict[str, str]] | None = None,
        provider: str | None = None,
        model: str | None = None,
        prompt_template: str | None = None,
        context_text: str | None = None,
        context_filename: str | None = None,
    ) -> AnswerResult:
        """Handle generate for the Open A I Compatible Answer Generator instance."""
        start = time.time()
        messages, context_tokens = self._build_messages(
            query=query,
            retrieved_chunks=retrieved_chunks,
            history=history,
            prompt_template=prompt_template,
            context_text=context_text,
            context_filename=context_filename,
        )
        client = get_llm_client(provider=provider or settings.llm_provider, model=model or settings.llm_model)
        response = client.chat(messages)
        latency_ms = int((time.time() - start) * 1000)
        return AnswerResult(
            answer=response.content if response.is_success else "",
            sources=self._sources(retrieved_chunks),
            model=response.model or (model or settings.llm_model),
            latency_ms=latency_ms,
            retrieved_count=len(retrieved_chunks),
            context_tokens=context_tokens,
            truncated=self._is_context_truncated(
                retrieved_chunks=retrieved_chunks,
                context_tokens=context_tokens,
            ),
            error=response.error,
        )

    def stream_generate(
        self,
        *,
        query: str,
        retrieved_chunks: list[RetrievedChunk],
        history: list[dict[str, str]] | None = None,
        provider: str | None = None,
        model: str | None = None,
        prompt_template: str | None = None,
        context_text: str | None = None,
        context_filename: str | None = None,
    ) -> Generator[dict, None, AnswerResult]:
        """Stream generate for the Open A I Compatible Answer Generator instance."""
        start = time.time()
        messages, context_tokens = self._build_messages(
            query=query,
            retrieved_chunks=retrieved_chunks,
            history=history,
            prompt_template=prompt_template,
            context_text=context_text,
            context_filename=context_filename,
        )
        sources = [source.__dict__ for source in self._sources(retrieved_chunks)]
        yield {"event": "sources", "data": sources}
        client = get_llm_client(provider=provider or settings.llm_provider, model=model or settings.llm_model)
        answer_parts: list[str] = []
        try:
            if hasattr(client, "stream_chat"):
                for chunk in client.stream_chat(messages):
                    answer_parts.append(chunk)
                    yield {"event": "content", "data": chunk}
            else:
                response = client.chat(messages)
                answer_parts.append(response.content)
                yield {"event": "content", "data": response.content}
        except Exception as exc:
            yield {"event": "error", "data": str(exc)}
            return
        yield {
            "event": "done",
            "data": {
                "latency_ms": int((time.time() - start) * 1000),
                "retrieved_count": len(retrieved_chunks),
                "context_tokens": context_tokens,
                "model": model or settings.llm_model,
            },
        }
