"""Stream token filtering, markdown sanitization, and history compression utilities."""

from __future__ import annotations

import json
import re
from typing import Iterator

from services.observability import setup_logger

logger = setup_logger(__name__)

MAX_HISTORY_CHARS = 400

_SUPPRESS_TAGS = [
    ("<think>", "</think>"),
    ("<tool_call>", "</tool_call>"),
    ("<|constrain|>", None),
]


def _find_start_tag(buf: str) -> tuple[bool, str, str, str | None]:
    """Finds the earliest occurrence of an internal start tag in the buffer."""
    for start_tag, end_tag in _SUPPRESS_TAGS:
        idx = buf.find(start_tag)
        if idx != -1:
            before = buf[:idx]
            after = buf[idx + len(start_tag):]
            return True, before, after, end_tag
    return False, "", buf, None


def _consume_suppressed(buf: str, suppress_end: str | None) -> tuple[str, bool]:
    """Discards suppressed content until the closing delimiter is encountered."""
    if suppress_end is None:
        return "", True
    end_idx = buf.find(suppress_end)
    if end_idx != -1:
        return buf[end_idx + len(suppress_end):].lstrip("\n\r "), False
    return "", True


def _handle_unsuppressed_buf(buf: str) -> tuple[str, bool, str | None, list[str]]:
    """Evaluates unsuppressed buffer, returning updated buffer and emitted tokens."""
    found, before, after, end_tag = _find_start_tag(buf)
    if found:
        tokens = [before] if before else []
        return after, True, end_tag, tokens

    last_lt = buf.rfind("<")
    if last_lt != -1 and last_lt > len(buf) - 15:
        return buf[last_lt:], False, None, [buf[:last_lt]]
    return "", False, None, [buf]


def _step_buffer(buf: str, suppressing: bool, suppress_end: str | None) -> tuple[str, bool, str | None, list[str], bool]:
    """Processes a single cycle of the streaming buffer."""
    if not suppressing:
        new_buf, new_supp, new_end, tokens = _handle_unsuppressed_buf(buf)
        stop = bool(not new_supp and new_buf)
        return new_buf, new_supp, new_end, tokens, stop

    new_buf, new_supp = _consume_suppressed(buf, suppress_end)
    return new_buf, new_supp, suppress_end, [], new_supp


def strip_internal_tokens(delta_iter: Iterator[str]) -> Iterator[str]:
    """Generator filter consuming raw delta strings and yielding clean user-facing text."""
    buf = ""
    suppressing = False
    suppress_end: str | None = ""

    for delta in delta_iter:
        buf += delta
        while buf:
            buf, suppressing, suppress_end, tokens, stop = _step_buffer(buf, suppressing, suppress_end)
            for token in tokens:
                yield token
            if stop:
                break

    if buf and not suppressing:
        yield buf


def split_tokens(text: str, chunk_size: int = 4) -> Iterator[str]:
    """Split text into small word chunks to simulate streaming for non-stream calls."""
    words = text.split(" ")
    buf: list[str] = []
    for word in words:
        buf.append(word)
        if len(buf) >= chunk_size:
            yield " ".join(buf) + " "
            buf = []
    if buf:
        yield " ".join(buf)


def extract_clarification_options(text: str) -> list[str]:
    """Extract interactive button chip options if the model offered numbered choices."""
    lines = [line.strip() for line in text.split("\n") if line.strip()]
    options: list[str] = []
    for line in lines:
        match = re.match(r"^(?:(?:\d+\.|\*|-)\s+)?\*{0,2}([\w\s]{3,40}?)\*{0,2}(?::|\s-\s|\?|$)", line)
        if match:
            candidate = match.group(1).strip()
            if candidate and 2 < len(candidate) <= 35:
                if not candidate.lower().startswith(("here are", "would you", "option", "i can")):
                    options.append(candidate)
    if 2 <= len(options) <= 4:
        return options
    return []


def compress_for_history(text: str) -> str:
    """Trim assistant response to first 400 chars to save tokens in multi-turn history."""
    clean = re.sub(r"\[.*?\]\(.*?\)", "", text)
    clean = re.sub(r"[*_`#]", "", clean)
    clean = " ".join(clean.split())
    if len(clean) > MAX_HISTORY_CHARS:
        return clean[:MAX_HISTORY_CHARS].rsplit(" ", 1)[0] + "..."
    return clean


def clean_reply_text(text: str) -> str:
    """Removes leaked action chips and trailing whitespace from the final response text."""
    clean = re.sub(
        r"\[(ACTION_CHIPS|QUICK_ACTIONS|SUGGESTED_ACTIONS):.*?\]",
        "",
        text,
        flags=re.DOTALL,
    )
    return clean.strip()


def strip_internal_model_markers(text: str) -> str:
    """Strips internal model reasoning and tool-call markers from final text."""
    clean = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
    clean = re.sub(r"<tool_call>.*?</tool_call>", "", clean, flags=re.DOTALL)
    clean = re.sub(r"<\|.*?\|>", "", clean)
    return clean.strip()