"""Assistant-visible text sanitization (think tags, CoT, tool markup, topics)."""
from __future__ import annotations

import json
import logging
import re
from typing import Any

from agent_loop_stream import normalize_stream_text
from agent_loop_tool_parse import _strip_xml_tool_calls, _TOOL_CALL_TAG_RE
from utils.assistant_visibility_placeholders import PLACEHOLDER_NO_VISIBLE_REPLY, PLACEHOLDER_TOOL_FORMAT_GLITCH
from utils.llm_output_sanitize import (
    strip_reasoning_preamble_leak as _strip_reasoning_preamble_leak,
    strip_untagged_cot_with_meta as _strip_untagged_cot_with_meta,
)
from utils.transcript_roles import is_real_user_message

_log = logging.getLogger(__name__)

_THINK_TAG_RE = re.compile(
    r"<think>[\s\S]*?</think>"
    r"|<\|channel>thought[\s\S]*?<channel\|>",
    re.IGNORECASE,
)
_THINK_UNCLOSED_RE = re.compile(
    r"<think>[\s\S]*$"
    r"|<\|channel>thought[\s\S]*$",
    re.IGNORECASE,
)
_THINK_OPEN_TAGS = ("<think>", "<|channel>thought")
_THINK_CLOSE_TAGS = ("</think>", "<channel|>")
_ORPHAN_THINK_CLOSE_RE = re.compile(r"^[\s\S]*?(?:</think>|<channel\|>)", re.IGNORECASE)

_TOPIC_TOKEN_RE = re.compile(r"[a-zA-Z0-9_]{3,}")
_TOPIC_STOPWORDS = frozenset(
    {
        "about",
        "after",
        "again",
        "agent",
        "also",
        "and",
        "assistant",
        "because",
        "been",
        "before",
        "being",
        "between",
        "both",
        "build",
        "can",
        "could",
        "does",
        "done",
        "from",
        "have",
        "help",
        "into",
        "just",
        "like",
        "make",
        "more",
        "most",
        "need",
        "only",
        "other",
        "over",
        "please",
        "really",
        "should",
        "some",
        "than",
        "that",
        "their",
        "them",
        "then",
        "there",
        "these",
        "they",
        "this",
        "those",
        "through",
        "topic",
        "topics",
        # Tool/JSON noise that should not become curiosity topics.
        "tool",
        "json",
        "file",
        "files",
        "error",
        "errors",
        "content",
        "message",
        "messages",
        "response",
        "result",
        "results",
        "data",
        "null",
        "true",
        "false",
        "none",
        "init",
        "unknown",
        "very",
        "want",
        "what",
        "when",
        "where",
        "which",
        "with",
        "would",
        "your",
        "you",
    }
)
_HEDGING_PHRASES = (
    "i think",
    "i might be wrong",
    "i'm not sure",
    "i am not sure",
    "might be",
    "possibly",
    "probably",
    "unclear",
    "not certain",
    "could be",
)

_CJK_CHAR_RE = re.compile(r"[\u3400-\u4DBF\u4E00-\u9FFF\uF900-\uFAFF]")
_LATIN_EMBEDDED_CJK_RE = re.compile(
    r"(?<=[A-Za-z])[\u3400-\u4DBF\u4E00-\u9FFF\uF900-\uFAFF]{1,8}(?=[A-Za-z])"
)
_NON_ENGLISH_LANG_HINTS = (
    "chinese",
    "mandarin",
    "cantonese",
    "japanese",
    "korean",
    "spanish",
    "french",
    "german",
    "italian",
    "portuguese",
    "russian",
    "arabic",
    "hindi",
    "urdu",
)
_NON_ENGLISH_REQUEST_CUES = (
    "reply in",
    "respond in",
    "answer in",
    "write in",
    "speak in",
    "translate",
)

_BARE_CHANNEL_PREAMBLE_RE = re.compile(
    r"^\s*(?:thought|thinking|analysis)\s*[:\n\r]",
    re.IGNORECASE,
)


def strip_leading_reasoning_channel_marker(text: str) -> str:
    """Drop a bare leading reasoning-channel token from assistant content.

    Some local models (e.g. Gemma served via MLX/LM Studio with tools in
    context) emit the channel *name* as plain text — ``"thought\\nHello"`` —
    once the surrounding ``<|channel>``/``<channel|>`` delimiters are stripped
    upstream. The tagged forms are handled by ``_strip_think_tags``; this catches
    the bare residue. Conservative: only fires when the very first token of the
    message is exactly one of the known channel names followed by ``:`` or a
    newline, so ordinary prose like "Thoughts on this:" is never touched
    (the trailing 's' breaks the anchor).
    """
    if not isinstance(text, str) or not text:
        return text
    m = _BARE_CHANNEL_PREAMBLE_RE.match(text)
    if not m:
        return text
    return text[m.end():].lstrip()


def _strip_think_tags(text: str) -> str:
    """Remove <think>...</think> blocks from text for storage/display.
    Also strips orphaned close tags (no opener) and unclosed open tags (no closer)."""
    result = _THINK_TAG_RE.sub("", text)
    result = _ORPHAN_THINK_CLOSE_RE.sub("", result)
    result = _THINK_UNCLOSED_RE.sub("", result)
    return result.strip()


def _user_requested_non_english(messages: list[dict] | None, lookback: int = 6) -> bool:
    if not isinstance(messages, list) or not messages:
        return False
    recent = [m for m in messages if isinstance(m, dict) and is_real_user_message(m)][-max(1, int(lookback)) :]
    for m in recent:
        content = m.get("content")
        if not isinstance(content, str) or not content.strip():
            continue
        text = content.strip()
        low = text.lower()
        if _CJK_CHAR_RE.search(text):
            return True
        if any(lang in low for lang in _NON_ENGLISH_LANG_HINTS) and any(cue in low for cue in _NON_ENGLISH_REQUEST_CUES):
            return True
    return False


def _strip_mixed_script_glitches(
    text: str,
    messages: list[dict] | None,
    stage: str = "final",
) -> str:
    """Drop tiny CJK spans embedded in Latin words for English threads.

    This targets generation glitches like 'still验证ing' while preserving
    responses when the user explicitly asked for non-English output.
    """
    if not isinstance(text, str) or not text:
        return text
    if _user_requested_non_english(messages):
        return text
    if not _CJK_CHAR_RE.search(text):
        return text

    ascii_letters = len(re.findall(r"[A-Za-z]", text))
    cjk_chars = len(_CJK_CHAR_RE.findall(text))
    # Keep potentially intentional non-English output; only scrub sparse
    # CJK insertions in otherwise English text.
    if ascii_letters < 20 or cjk_chars > max(8, int(ascii_letters * 0.25)):
        return text

    cleaned = _LATIN_EMBEDDED_CJK_RE.sub("", text)
    if cleaned != text:
        _log.debug(
            "Assistant sanitize (%s): removed mixed-script embedded CJK chars=%d; len %d -> %d",
            stage,
            cjk_chars,
            len(text),
            len(cleaned),
        )
    return cleaned


def _finalize_assistant_text_for_user(
    text: str,
    stage: str = "final",
    messages: list[dict] | None = None,
) -> str:
    """Canonical sanitizer for assistant-visible text (persist + UI).

    Order matters:
      normalize_stream_text -> strip think tags -> strip untagged CoT
      -> strip plain-text reasoning preambles (e.g. Qwen 'Thinking Process:')
      -> strip XML tool-calls -> mixed-script glitch cleanup
    """
    raw = text if isinstance(text, str) else (str(text) if text is not None else "")
    try:
        normalized = normalize_stream_text(raw)
    except Exception:
        _log.debug("Failed to normalize assistant text (%s)", stage, exc_info=True)
        normalized = raw

    no_think = _strip_think_tags(normalized)
    cot_clean, cot_meta = _strip_untagged_cot_with_meta(no_think)
    no_reasoning = _strip_reasoning_preamble_leak(cot_clean)
    no_xml = _strip_xml_tool_calls(no_reasoning)
    mixed_clean = _strip_mixed_script_glitches(no_xml, messages, stage=stage)
    final = mixed_clean.strip()

    removed_think = max(0, len(normalized) - len(no_think))
    removed_cot = int(cot_meta.get("removed_chars", 0)) if isinstance(cot_meta, dict) else 0
    removed_reasoning = max(0, len(cot_clean) - len(no_reasoning))
    removed_xml = max(0, len(no_reasoning) - len(no_xml))
    if removed_think or removed_cot or removed_reasoning or removed_xml:
        _log.debug(
            "Assistant sanitize (%s): removed think=%d cot=%d reasoning=%d xml=%d; len %d -> %d; delimiter=%s planning=%s",
            stage,
            removed_think,
            removed_cot,
            removed_reasoning,
            removed_xml,
            len(raw),
            len(final),
            bool(cot_meta.get("delimiter_found")) if isinstance(cot_meta, dict) else False,
            bool(cot_meta.get("planning_stripped")) if isinstance(cot_meta, dict) else False,
        )
    return final


def _assistant_history_content_after_tool_extraction(
    text: str,
    messages: list[dict] | None = None,
) -> str:
    """Clean assistant content before appending a tool-call round to history.

    Tool extraction may rely on raw XML/Gemma markup, but the next LLM request
    should not see that markup as ordinary assistant prose. Keeping it in
    history makes prompts grow and can teach the model to repeat malformed
    fallback syntax.
    """
    cleaned = _finalize_assistant_text_for_user(
        text or "",
        stage="tool_round_history",
        messages=messages,
    )
    if _assistant_final_text_looks_like_tool_dump(cleaned):
        return ""
    return cleaned


def _assistant_final_text_looks_like_tool_dump(text: str) -> bool:
    """True when user-visible text still looks like raw tool-call payload (should not ship to UI)."""
    if not isinstance(text, str):
        return False
    s = text.strip()
    if not s:
        return False
    if _TOOL_CALL_TAG_RE.search(s):
        return True
    if re.search(r"</tool_call>|<tool_call\|>", s, re.IGNORECASE):
        return True
    if s.startswith("{") and len(s) <= 1800:
        low = s.lower()
        if '"name"' in low and '"arguments"' in low:
            try:
                obj = json.loads(s)
            except (json.JSONDecodeError, TypeError, ValueError):
                obj = None
            if isinstance(obj, dict) and isinstance(obj.get("name"), str) and "arguments" in obj:
                return True
    return False


def _ensure_nonempty_assistant_reply(cleaned: str, raw: str) -> str:
    """Avoid blank assistant bubbles when the model produced only internal markup."""
    if (cleaned or "").strip():
        if _assistant_final_text_looks_like_tool_dump(cleaned):
            _log.debug("Assistant sanitize: replacing reply that still resembles tool-call markup")
            return PLACEHOLDER_TOOL_FORMAT_GLITCH
        return cleaned
    if (raw or "").strip():
        return PLACEHOLDER_NO_VISIBLE_REPLY
    return ""


def _extract_conversation_topics(messages: list[dict], limit: int = 5) -> list[str]:
    """Extract lightweight topic keywords from the latest user message."""
    if not isinstance(messages, list):
        return []
    latest_user = ""
    for msg in reversed(messages):
        if not is_real_user_message(msg):
            continue
        content = msg.get("content")
        if isinstance(content, str) and content.strip():
            latest_user = content.strip()
            break
    if not latest_user:
        return []
    counts: dict[str, int] = {}
    first_seen: dict[str, int] = {}
    idx = 0
    for token in _TOPIC_TOKEN_RE.findall(latest_user.lower()):
        if "_" in token or len(token) < 4:
            continue
        if token in _TOPIC_STOPWORDS:
            continue
        counts[token] = counts.get(token, 0) + 1
        if token not in first_seen:
            first_seen[token] = idx
            idx += 1
    ordered = sorted(counts.items(), key=lambda kv: (-kv[1], first_seen.get(kv[0], 0), kv[0]))
    return [token for token, _ in ordered[: max(0, int(limit))]]


def _estimate_response_confidence(response_text: str) -> float:
    """Estimate confidence from response hedging language."""
    text = (response_text or "").strip().lower()
    if not text:
        return 0.5
    score = 0.6
    for phrase in _HEDGING_PHRASES:
        if phrase in text:
            score -= 0.08
    if "?" in text:
        score -= 0.03
    if "definitely" in text or "certainly" in text:
        score += 0.04
    return max(0.1, min(0.95, score))
