"""Translation team agent: MsgHub-based multi-agent collaboration."""

from __future__ import annotations

import asyncio
import logging
import re
import threading
from typing import Any

from agentscope.agent import ReActAgent
from agentscope.formatter import OpenAIChatFormatter
from agentscope.message import Msg

from ..config import get_model
from ..prompts import battle, mystery, daily

logger = logging.getLogger(__name__)

# Per-thread persistent event loop to avoid the overhead of asyncio.run()
# which creates/destroys a new loop on every call.
_thread_local = threading.local()


def _get_or_create_loop() -> asyncio.AbstractEventLoop:
    """Get or create a persistent event loop for the current thread."""
    loop = getattr(_thread_local, "loop", None)
    if loop is None or loop.is_closed():
        loop = asyncio.new_event_loop()
        _thread_local.loop = loop
    return loop

# ---------------------------------------------------------------------------
# Style -> prompt builder mapping
# ---------------------------------------------------------------------------
_STYLE_PROMPT_MAP: dict[str, Any] = {
    "battle": battle.build_prompt,
    "mystery": mystery.build_prompt,
    "daily": daily.build_prompt,
}


def _select_prompt_builder(novel_type: str):
    """Return the appropriate prompt builder based on novel type."""
    key = novel_type.lower().strip()
    for style_key in ("battle", "mystery"):
        if style_key in key:
            return _STYLE_PROMPT_MAP[style_key]
    return _STYLE_PROMPT_MAP["daily"]


# ---------------------------------------------------------------------------
# LLM output cleanup helpers
# ---------------------------------------------------------------------------
# Patterns commonly generated by Chinese LLMs that wrap the actual output
_LLM_CLEAN_PATTERNS: list[re.Pattern[str]] = [
    # Strip markdown code blocks (``` ... ```)
    re.compile(r"```[\w]*\n([\s\S]*?)```", re.MULTILINE),
    # Strip common preamble phrases followed by a colon or newline
    re.compile(r"^(?:好的[，,。]?|以下是[^\n]*[：:]\s*\n?)", re.MULTILINE),
    # Strip trailing sign-off lines like "希望对你有帮助！"
    re.compile(r"\n\s*希望[^\n]*[！!]\s*$"),
]


def _extract_text(content) -> str:
    """Extract text from agent response content.

    agentscope 1.x may return content as a string or as a list of
    content blocks (e.g., thinking + text).  Handle both cases.
    """
    if isinstance(content, list):
        # Extract 'text' fields from content blocks, skip 'thinking'
        parts = []
        for block in content:
            if isinstance(block, dict):
                if block.get("type") == "text":
                    parts.append(block.get("text", ""))
                elif "content" in block:
                    parts.append(block["content"])
                elif "text" in block:
                    parts.append(block["text"])
            elif isinstance(block, str):
                parts.append(block)
        return "\n".join(parts)
    return str(content) if content else ""


def _clean_llm_output(raw: str) -> str:
    """Strip common LLM preamble/postamble and markdown wrappers.

    Preserves internal newlines and paragraph structure of the translation.
    Only removes outer preamble phrases and markdown code block fences.
    """
    text = raw.strip()

    # Try to extract content from markdown code blocks first
    for pattern in _LLM_CLEAN_PATTERNS:
        match = pattern.search(text)
        if match:
            text = match.group(1) if match.lastindex else text
            break

    # Strip residual `---` separator lines that LLM may have copied from prompt
    lines = text.split("\n")
    cleaned = [line for line in lines if re.fullmatch(r"-{3,}", line.strip()) is None]
    text = "\n".join(cleaned)

    # Only strip outer whitespace, preserve internal newlines
    return text.strip()


class TranslationTeam:
    """Multi-agent translation team orchestrated via agentscope MsgHub.

    The team consists of three roles collaborating in a shared message hub:
      - **Translator**: performs the primary translation
      - **Terminologist**: checks glossary consistency and terminology
      - **Polisher**: refines tone, flow, and style

    After the group discussion, the final translated text is extracted.

    **Usage**: Create one instance via ``TranslationTeam.create(...)`` and
    call ``.translate(text)`` multiple times to reuse the same agents.
    This avoids the overhead of rebuilding agents for each chunk.
    """

    translator: ReActAgent | None
    terminologist: ReActAgent | None
    polisher: ReActAgent | None

    def __init__(
        self,
        source_lang: str = "Japanese",
        target_lang: str = "Chinese",
        novel_type: str = "daily",
        glossary_terms: list[dict[str, str]] | None = None,
        placeholders: dict[str, str] | None = None,
    ) -> None:
        self.source_lang = source_lang
        self.target_lang = target_lang
        self.novel_type = novel_type
        self.glossary_terms = glossary_terms or []
        self.placeholders = placeholders or {}
        self.model = get_model()
        self._agents: list[ReActAgent] = []

        # Declare agent attributes so type checkers know they exist
        self.translator = None
        self.terminologist = None
        self.polisher = None

        self._build_agents()

    # ------------------------------------------------------------------
    # Factory: create a reusable team instance
    # ------------------------------------------------------------------
    @classmethod
    def create(
        cls,
        source_lang: str = "Japanese",
        target_lang: str = "Chinese",
        novel_type: str = "daily",
        glossary_terms: list[dict[str, str]] | None = None,
        placeholders: dict[str, str] | None = None,
    ) -> "TranslationTeam":
        """Create a reusable TranslationTeam instance.

        The caller should create one instance and call ``.translate()``
        multiple times for different chunks. Reusing the instance avoids
        the overhead of rebuilding agents and re-initializing the model
        connection for each chunk.

        Args:
            source_lang: Source language name (e.g. "Japanese").
            target_lang: Target language name (e.g. "Chinese").
            novel_type: Novel style/type (e.g. "battle", "mystery", "daily").
            glossary_terms: Optional list of {"source": ..., "target": ...} dicts.
            placeholders: Optional dict of placeholder mappings (e.g. {"[{1}]": "翻译后的名词"}).

        Returns:
            A fully initialized TranslationTeam ready for repeated use.
        """
        return cls(
            source_lang=source_lang,
            target_lang=target_lang,
            novel_type=novel_type,
            glossary_terms=glossary_terms,
            placeholders=placeholders,
        )

    # ------------------------------------------------------------------
    # Agent construction
    # ------------------------------------------------------------------
    def _build_agents(self) -> None:
        prompt_builder = _select_prompt_builder(self.novel_type)
        sys_prompt = prompt_builder(
            source_lang=self.source_lang,
            target_lang=self.target_lang,
            glossary_terms=self.glossary_terms,
        )

        # Add placeholder protection instruction if placeholders exist
        if self.placeholders:
            placeholder_keys = ", ".join(self.placeholders.keys())
            placeholder_instruction = (
                f"\n\n重要：文本中包含占位符如 {placeholder_keys} 等，"
                "这些是已翻译的专有名词占位符，请勿翻译、修改或删除它们，"
                "在译文中保持原样输出。"
            )
            sys_prompt += placeholder_instruction

        # --- Translator ---
        self.translator = ReActAgent(
            name="Translator",
            sys_prompt=sys_prompt,
            model=self.model,
            formatter=OpenAIChatFormatter(),
        )

        # --- Terminologist ---
        term_prompt = self._terminologist_prompt()
        self.terminologist = ReActAgent(
            name="Terminologist",
            sys_prompt=term_prompt,
            model=self.model,
            formatter=OpenAIChatFormatter(),
        )

        # --- Polisher ---
        polish_prompt = self._polisher_prompt()
        self.polisher = ReActAgent(
            name="Polisher",
            sys_prompt=polish_prompt,
            model=self.model,
            formatter=OpenAIChatFormatter(),
        )

        self._agents = [self.translator, self.terminologist, self.polisher]

    # ------------------------------------------------------------------
    # Role-specific system prompts
    # ------------------------------------------------------------------
    def _terminologist_prompt(self) -> str:
        terms_text = ""
        if self.glossary_terms:
            lines = [
                f"  - {t.get('source', '')} -> {t.get('target', '')}"
                for t in self.glossary_terms
            ]
            terms_text = "\n".join(lines)

        placeholder_instruction = ""
        if self.placeholders:
            placeholder_keys = ", ".join(self.placeholders.keys())
            placeholder_instruction = (
                f"\n注意：文本中包含占位符如 {placeholder_keys} 等，"
                "这些是已翻译的专有名词占位符，请勿翻译或修改它们。"
            )

        return (
            f"你是术语审校专家。你的职责是检查译文中的专业术语和专有名词"
            f"是否与以下术语表一致：\n{terms_text}\n\n"
            "如果发现不一致的术语，请指出原文中的具体位置并给出修正建议。\n"
            "如果没有术语表，则检查通用译名的一致性。\n"
            "请只输出需要修改的部分，不要重复全文。"
            f"{placeholder_instruction}"
        )

    def _polisher_prompt(self) -> str:
        placeholder_instruction = ""
        if self.placeholders:
            placeholder_keys = ", ".join(self.placeholders.keys())
            placeholder_instruction = (
                f"\n注意：文本中包含占位符如 {placeholder_keys} 等，"
                "这些是已翻译的专有名词占位符，请勿翻译或修改它们。"
            )
        return (
            "你是文字润色专家。你的职责是在保持原意的前提下，"
            "对译文进行语言润色，使其更加流畅自然、符合目标语言的表达习惯。\n"
            "重点关注：句式变化、节奏感、情感传递、口语/书面语风格的匹配。\n"
            "请输出润色后的完整译文。"
            f"{placeholder_instruction}"
        )

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------
    def translate(self, text: str) -> str:
        """Run the full team translation pipeline (sync wrapper).

        Uses a persistent per-thread event loop to avoid the CPU overhead
        of creating/destroying a new loop on every request (asyncio.run).
        """
        if not text.strip():
            return ""
        loop = _get_or_create_loop()
        return loop.run_until_complete(self._async_translate(text))

    async def _async_translate(self, text: str) -> str:
        """Async implementation of the translation pipeline.

        1. Translator produces initial translation.
        2. Terminologist reviews terminology consistency (graceful fallback).
        3. Polisher produces the final polished version (graceful fallback).

        This method is safe to call multiple times on the same instance.

        Note: agents are called directly without msghub since agentscope 1.x
        MsgHub is async and cannot be nested within a running event loop.
        """
        logger.info(
            "TranslationTeam: starting translation (%d chars, src=%s, tgt=%s)",
            len(text),
            self.source_lang,
            self.target_lang,
        )

        # Ensure agents are initialized (defensive check)
        if not self.translator or not self.terminologist or not self.polisher:
            raise RuntimeError("Agents not initialized. Call _build_agents() first.")

        # ---------- Step 1: Translator produces initial translation -----
        translation_text = ""
        try:
            msg_translate = Msg(
                name="User",
                role="user",
                content=(
                    f"请将以下{self.source_lang}文本翻译为{self.target_lang}：\n\n"
                    f"<user_input>\n{text}\n</user_input>"
                ),
            )
            # Broadcast to all agents' memory
            for agent in self._agents:
                await agent.observe(msg_translate)
            resp_translator = await self.translator(msg_translate)
            translation_text = _extract_text(resp_translator.content)
            logger.debug(
                "Translator output length: %d", len(str(translation_text))
            )
        except Exception:
            logger.exception("Step 1 (Translator) failed")
            # Fall through: attempt remaining steps with empty translation

        # ---------- Step 2: Terminologist reviews terminology -----------
        term_feedback = "无需修改"
        try:
            msg_review = Msg(
                name="User",
                role="user",
                content=f"以下是初步译文，请检查术语一致性：\n\n<user_input>\n{translation_text}\n</user_input>",
            )
            for agent in self._agents:
                await agent.observe(msg_review)
            resp_terminologist = await self.terminologist(msg_review)
            term_feedback = _extract_text(resp_terminologist.content) or "无需修改"
            logger.debug("Terminologist feedback: %s", term_feedback[:200])
        except Exception:
            logger.warning(
                "Step 2 (Terminologist) failed, continuing with '无需修改'"
            )
            term_feedback = "无需修改"

        # ---------- Step 3: Polisher produces final version -------------
        try:
            msg_polish = Msg(
                name="User",
                role="user",
                content=(
                    f"以下是初步译文和术语审校意见：\n\n"
                    f"<user_input>\n{translation_text}\n</user_input>\n\n"
                    f"<terminologist_feedback>\n{term_feedback}\n</terminologist_feedback>\n\n"
                    f"请结合以上意见，输出最终润色后的完整译文。"
                ),
            )
            for agent in self._agents:
                await agent.observe(msg_polish)
            resp_polisher = await self.polisher(msg_polish)
            final_text = _extract_text(resp_polisher.content) or translation_text
        except Exception:
            logger.warning(
                "Step 3 (Polisher) failed, returning translator output"
            )
            final_text = translation_text

        # Clean up LLM preamble/postamble from the final output
        final_text = _clean_llm_output(str(final_text))

        logger.info(
            "TranslationTeam: finished translation (%d chars)", len(final_text)
        )
        return final_text
