from __future__ import annotations

import re
from typing import Any

_THINK_RE = re.compile(
    r"<think(?:\s[^>]*)?>.*?</think\s*>"
    r"|<\|channel>thought[\s\S]*?<channel\|>",
    re.IGNORECASE | re.DOTALL,
)
# Unclosed open tags: <think> or <|channel>thought with no matching close tag.
_THINK_UNCLOSED_RE = re.compile(
    r"<think(?:\s[^>]*)?>[\s\S]*$"
    r"|<\|channel>thought[\s\S]*$",
    re.IGNORECASE,
)
# Malformed / non-thought channel tags (e.g. ``<|channel>s`` glued to the visible
# answer) leak as-is — strip only the tag token, keeping the content visible.
# Case-sensitive lowercase run so ``<|channel>sYou are…`` drops just ``s``.
_CHANNEL_TAG_RESIDUE_RE = re.compile(r"<\|channel>[a-z0-9_]{0,24}|<channel\|>")

# Line-start pseudo workspace paths: models use these as CoT delimiters (often invalid names
# like <proactive-queue.md>). Restrict to path-like tokens; avoids matching <tool_call>.
_MD_DELIM_LINE_RE = re.compile(
    r"(?:^|\n)(\s*)<([A-Za-z0-9][A-Za-z0-9_./-]{0,119})\.md>\s*",
    re.IGNORECASE | re.MULTILINE,
)
_TRAILING_MD_LINE_RE = re.compile(
    r"^\s*<[A-Za-z0-9][A-Za-z0-9_./-]{0,119}\.md>\s*$",
    re.IGNORECASE,
)
# Single-line planning leaks when the model omits a real reply after a delimiter line.
_COT_SINGLE_LINE_PREFIXES = (
    "i need to ",
    "i'll need to ",
    "let me ",
    "first, ",
    "looking at ",
    "i'm going to ",
)
# Do not treat these as planning-only (common user-facing openers).
_LET_ME_KEEP_PREFIXES = (
    "let me know",
    "let me help",
    "let me get you",
    "let me grab",
    "let me pull",
)
_SAFETY_KEEP_TOKENS = (
    "can't",
    "cannot",
    "won't",
    "will not",
    "shouldn't",
    "should not",
    "must not",
    "unable",
    "not help",
    "not assist",
    "decline",
    "refuse",
    "sorry",
)

# Reasoning preambles (plain prose, not <think> tags). One or more newlines after ":".
# Qwen-style models often emit "Here's a thinking process:" / "Here is a thinking process:" instead of "Thinking process:".
_REASONING_HEADER_RE = re.compile(
    r"^(?:(?:here(?:['\u2019]s| is)\s+)(?:a|the)\s+)?"
    r"(?:thinking\s+process|thought\s+process|my\s+thinking|internal\s+reasoning|"
    r"chain[- ]of[- ]thought)\s*:\s*\n+"
    r"|^(?:reasoning|analysis)\s*:\s*\n+",
    re.IGNORECASE,
)

# Lines inside a numbered CoT step that are not * / - bullets but still meta (draft quotes, echoes).
_COT_LINE_ECHO_HINTS = (
    "system prompt",
    "do not emit hidden",
    "do not use tools",
    "synthetic startup",
    "user-visible answer",
    "user input:",
    "constraints:",
)


def _looks_like_numbered_cot_continuation(raw_line: str, stripped: str) -> bool:
    if not stripped:
        return True
    if stripped.startswith("*") or re.match(r"^[-•]\s", stripped):
        return True
    if re.match(r"^\*\*[^*]+\*\*:", stripped):
        return True
    # Indented continuation under a step
    if raw_line[:1] in (" ", "\t") and raw_line.strip():
        return True
    # Draft / meta fragments
    if re.match(r'^["\u201c\u2018\\]', stripped) or re.match(r"^\\+\s*[\"'\u201c]", stripped):
        return True
    low = stripped.lower()
    if re.match(r"^(?:let'?s|wait,|wait:|->)\b", low):
        return True
    if stripped.startswith("(") and len(stripped) < 160:
        return True
    if any(h in low for h in _COT_LINE_ECHO_HINTS):
        return True
    # Field-heavy analysis lines (role / peer / context dumps)
    if stripped.count("**") >= 2 and ":" in stripped:
        return True
    return False


def strip_think_tags(text: str) -> str:
    if not isinstance(text, str):
        return ""
    low = text.lower()
    if "<think" not in low and "<|channel>" not in low and "<channel|>" not in low:
        return text
    result = _THINK_RE.sub("", text)
    result = _THINK_UNCLOSED_RE.sub("", result)
    result = _CHANNEL_TAG_RESIDUE_RE.sub("", result)
    return result.strip()


def _strip_trailing_md_delimiter_lines(text: str) -> str:
    """Drop trailing lines that are only a pseudo <something.md> tag."""
    s = text.rstrip("\n\r ")
    while True:
        if not s:
            return s
        last_nl = s.rfind("\n")
        last_line = s[last_nl + 1 :] if last_nl != -1 else s
        if not _TRAILING_MD_LINE_RE.match(last_line):
            return s
        s = s[:last_nl].rstrip("\n\r ") if last_nl != -1 else ""


def _strip_planning_only_single_line(text: str) -> str:
    """Remove a lone planning sentence left after delimiter stripping."""
    t = text.strip()
    if not t or "\n" in t:
        return text
    low = t.lower()
    for p in _COT_SINGLE_LINE_PREFIXES:
        if low.startswith(p):
            if p == "let me " and any(low.startswith(k) for k in _LET_ME_KEEP_PREFIXES):
                continue
            if any(token in low for token in _SAFETY_KEEP_TOKENS):
                # Keep likely refusals/safety responses even when they look "planning-like".
                continue
            return ""
    return text


def strip_untagged_cot_with_meta(text: str) -> tuple[str, dict[str, Any]]:
    """Strip untagged CoT and return (clean_text, metadata).

    Metadata keys:
      - delimiter_found: bool
      - planning_stripped: bool
      - trailing_tag_stripped: bool
      - removed_chars: int
    """
    if not isinstance(text, str) or not text:
        out = text or ""
        return out, {
            "delimiter_found": False,
            "planning_stripped": False,
            "trailing_tag_stripped": False,
            "removed_chars": 0,
        }
    original = text
    s = text.replace("\r\n", "\n")
    delimiter_found = False
    planning_stripped = False
    first_pass = _strip_trailing_md_delimiter_lines(s)
    trailing_tag_stripped = first_pass != s
    s = first_pass
    while True:
        m = _MD_DELIM_LINE_RE.search(s)
        if not m:
            break
        delimiter_found = True
        tail = s[m.end() :].lstrip("\n\r ")
        if tail:
            s = tail
            break
        s = s[: m.start()].rstrip("\n\r ")
        next_pass = _strip_trailing_md_delimiter_lines(s)
        trailing_tag_stripped = trailing_tag_stripped or (next_pass != s)
        s = next_pass
    if delimiter_found:
        stripped = _strip_planning_only_single_line(s)
        planning_stripped = stripped != s
        s = stripped
    out = s.strip()
    removed_chars = max(0, len(original) - len(out))
    return out, {
        "delimiter_found": delimiter_found,
        "planning_stripped": planning_stripped,
        "trailing_tag_stripped": trailing_tag_stripped,
        "removed_chars": removed_chars,
    }


def _strip_leading_numbered_analysis_block(s: str) -> str:
    """Remove a leading '1. **...** / * bullet' outline often emitted after 'Thinking Process:'."""
    lines = s.split("\n")
    idx = 0
    n = len(lines)
    while idx < n and not lines[idx].strip():
        idx += 1
    if idx >= n:
        return ""
    if not re.match(r"^\d+\.\s+", lines[idx].strip()):
        return s
    while idx < n:
        st = lines[idx].strip()
        if not st:
            idx += 1
            continue
        if not re.match(r"^\d+\.\s+", st):
            break
        idx += 1
        while idx < n:
            raw2 = lines[idx]
            st2 = raw2.strip()
            if not st2:
                idx += 1
                continue
            if re.match(r"^\d+\.\s+", st2):
                break
            if _looks_like_numbered_cot_continuation(raw2, st2):
                idx += 1
                continue
            break
    return "\n".join(lines[idx:]).strip()


def strip_reasoning_preamble_leak(text: str) -> str:
    """Drop plain-text reasoning headers ('Thinking Process:', etc.) and following outline bullets.

    Used for chat bubbles and history. Only strips the numbered/bullet *outline* when at least
    one reasoning header was removed, so legitimate assistant replies that start with '1. …' are kept.
    """
    if not isinstance(text, str) or not text.strip():
        return (text or "").strip()
    s = text.replace("\r\n", "\n")
    preamble_removed = False
    while True:
        m = _REASONING_HEADER_RE.match(s)
        if not m:
            break
        s = s[m.end() :].lstrip("\n\r \t")
        preamble_removed = True
        if not s:
            return ""
    if preamble_removed:
        s = _strip_leading_numbered_analysis_block(s)
    return s.strip()


def strip_untagged_cot(text: str) -> str:
    """Strip untagged chain-of-thought that precedes a pseudo workspace <*.md> delimiter line.

    Local models emit plain-text reasoning then a line like <IDENTITY.md> or
    <proactive-queue.md> before the user-visible reply. Trailing orphan tags
    (e.g. only <SESSION-STATE.md>) are removed so bubbles are not tag-only.
    """
    cleaned, _meta = strip_untagged_cot_with_meta(text)
    return cleaned


def strip_markdown_fences(text: str) -> str:
    stripped = (text or "").strip()
    if not stripped.startswith("```"):
        return stripped
    stripped = re.sub(r"^\s*```(?:[a-zA-Z0-9_-]+)?\s*\n?", "", stripped)
    stripped = re.sub(r"\n?\s*```\s*$", "", stripped)
    return stripped.strip()


def extract_structured_tail(text: str) -> str:
    raw = (text or "").strip()
    if not raw:
        return ""
    if "---OUTPUT---" in raw:
        return raw.split("---OUTPUT---", 1)[-1].strip()
    parts = [p.strip() for p in re.split(r"\n\s*\n", raw) if p.strip()]
    return parts[-1] if parts else raw


def sanitize_summary_text(text: str) -> str:
    """Strip think tags, unclosed tags, and reasoning preambles from LLM summary output."""
    if not isinstance(text, str) or not text.strip():
        return (text or "").strip()
    s = strip_think_tags(text)
    s = strip_untagged_cot(s)
    s = strip_reasoning_preamble_leak(s)
    return s.strip()
