"""
title: Reasoning Stripper
id: reasoning_stripper
author: p1s4
description: Strips LLM reasoning/thinking blocks from older messages to save context tokens. Runs before Async Context Compression (prio 15) to clean stale reasoning before summarization.
version: 1.0.0
license: MIT

═══════════════════════════════════════════════════════════════════════
📌 Overview
═══════════════════════════════════════════════════════════════════════

This filter removes LLM "thinking" / "reasoning" blocks from conversation
messages BEFORE they are sent to the LLM, dramatically reducing token usage
in long conversations.

WHY: Modern reasoning models (DeepSeek R1, Qwen3, Gemini, o1, etc.) emit
large thinking/reasoning blocks before their actual answers. These blocks:
  • Consume 500-5000+ tokens EACH
  • Are useful for the immediate response but useless after a few turns
  • Accumulate rapidly, filling the context window with stale reasoning
  • Get summarized by Async Context Compression, wasting even MORE tokens

PLACEMENT: Priority 5 — before Async Context Compression (prio 15).
This means:
  1. (optional) earlier inlet filters inject context
  2. ⭐ Reasoning Stripper cleans old thoughts (prio 5)  ← THIS FILTER
  3. Async Context Compression summarizes already-clean context (prio 15)

═══════════════════════════════════════════════════════════════════════
🛡️ Safety Rules
═══════════════════════════════════════════════════════════════════════

1. NEVER strips from the last N messages (configurable, default 4)
   → Some models need their own recent thoughts to maintain coherence
2. NEVER strips from user messages — reasoning only appears in assistant msgs
3. NEVER strips from system messages — instructions must be preserved
4. NEVER strips from messages with tool_calls — reasoning before tool use
   may be part of the decision chain
5. NEVER modifies the original chat history in the database — only the
   in-flight messages sent to the LLM
6. Only strips KNOWN reasoning patterns — no false positives

═══════════════════════════════════════════════════════════════════════
🔍 Supported Reasoning Patterns
═══════════════════════════════════════════════════════════════════════

1. <details type="reasoning">...</details>  — OpenWebUI + Gemini format
2. <details type="thinking">...</details>   — Alternative OWUI format
3. <details type="thought">...</details>    — Another variant
4. <details ...>...</details> (generic)     — Any OWUI details block
   (only when strip_all_details is enabled)
5. <think>...</think>                       — DeepSeek R1, Qwen3, etc.
6. <thinking>...</thinking>                 — Some Anthropic-style models
7. <|channel|>analysis<|message|>...<|channel|> — Interleaved thinking

═══════════════════════════════════════════════════════════════════════
⚙️ Configuration
═══════════════════════════════════════════════════════════════════════

priority:               Filter priority (default: 5, before ACC prio 15)
enabled:                Master switch (default: true)
protect_recent:         Number of recent messages to never strip (default: 4)
strip_reasoning_details: Strip <details type="reasoning/thinking/thought"> (default: true)
strip_think_tags:       Strip <think>...</think> and <thinking>...</thinking> (default: true)
strip_all_details:      Strip ALL <details> blocks, not just reasoning ones (default: false)
                        ⚠️ WARNING: This removes tool_calls details too!
                        Only enable if you don't use native function calling
                        or if ACC already handles tool output trimming.
strip_channel_analysis: Strip <|channel|>analysis blocks (default: true)
min_messages_before_strip: Don't strip anything until chat has this many
                           messages (default: 6, i.e. ~3 turns)
show_status:            Show status notification about stripping results (default: true)
debug_mode:             Enable detailed logging (default: false)

═══════════════════════════════════════════════════════════════════════
📝 Design Notes
═══════════════════════════════════════════════════════════════════════

This filter is intentionally MINIMAL and STATELESS:
  • No database — nothing to persist or break
  • No LLM calls — zero latency, zero cost
  • No async background tasks — runs synchronously in inlet
  • Pure regex-based — deterministic, no surprises
  • Only modifies in-flight messages — DB history is untouched

The "protect_recent" mechanism is the key safety valve:
  - Recent thoughts help the model maintain chain-of-thought coherence
  - After 3-4 messages, the model no longer needs its own old thoughts
  - The token savings are enormous: each stripped reasoning block saves
    500-5000+ tokens from being sent, and later from being summarized

"""

import re
import logging
from typing import Optional, Callable, Awaitable, Any
from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)

# ═══════════════════════════════════════════════════════════════════════════
# Regex Patterns for Reasoning Block Detection
# ═══════════════════════════════════════════════════════════════════════════

# Pattern 1: <details type="reasoning">...</details> (OpenWebUI + Gemini)
# Also matches type="thinking" and type="thought"
REASONING_DETAILS_PATTERN = re.compile(
    r'<details[^>]*\btype\s*=\s*["\'](?:reasoning|thinking|thought)["\'][^>]*>.*?</details>',
    re.DOTALL | re.IGNORECASE,
)

# Pattern 2: ANY <details>...</details> block (aggressive — only when strip_all_details is on)
ALL_DETAILS_PATTERN = re.compile(
    r'<details[^>]*>.*?</details>',
    re.DOTALL | re.IGNORECASE,
)

# Pattern 3: <think>...</think> (DeepSeek R1, Qwen3, etc.)
THINK_TAG_PATTERN = re.compile(
    r'<think>.*?</think>',
    re.DOTALL | re.IGNORECASE,
)

# Pattern 4: <thinking>...</thinking> (some Anthropic-style models)
THINKING_TAG_PATTERN = re.compile(
    r'<thinking>.*?</thinking>',
    re.DOTALL | re.IGNORECASE,
)

# Pattern 5: <|channel|>analysis<|message|>...<|channel|> (interleaved)
CHANNEL_ANALYSIS_PATTERN = re.compile(
    r'<\|channel\|>analysis<\|message\|>.*?(?=<\|channel\|>|$)',
    re.DOTALL,
)


class Filter:
    """
    Open WebUI Filter that strips reasoning/thinking blocks from older
    messages in the conversation to reduce context token usage.
    """

    class Valves(BaseModel):
        priority: int = Field(
            default=5,
            description=(
                "Filter priority. Default 5 runs before Async Context Compression (prio 15) "
                "so stale reasoning is removed before summarization."
            ),
        )
        enabled: bool = Field(
            default=True,
            description="Master switch. Disable to skip all reasoning stripping.",
        )
        protect_recent: int = Field(
            default=4,
            ge=0,
            description=(
                "Number of recent messages to NEVER strip reasoning from. "
                "Some models need their own recent thoughts for coherence. "
                "Default 4 protects the last ~2 turns of reasoning. "
                "Set to 0 to strip from all messages (aggressive)."
            ),
        )
        strip_reasoning_details: bool = Field(
            default=True,
            description=(
                'Strip <details type="reasoning">, <details type="thinking">, '
                'and <details type="thought"> blocks. '
                'These are the OpenWebUI-wrapped reasoning tokens from Gemini and similar models.'
            ),
        )
        strip_think_tags: bool = Field(
            default=True,
            description=(
                "Strip <think>...</think> and <thinking>...</thinking> blocks. "
                "These are the raw reasoning tokens from DeepSeek R1, Qwen3, and similar models."
            ),
        )
        strip_all_details: bool = Field(
            default=False,
            description=(
                "Strip ALL <details>...</details> blocks, not just reasoning ones. "
                "⚠️ WARNING: This also removes <details type='tool_calls'> blocks! "
                "Only enable if you don't use native function calling or if "
                "Async Context Compression already handles tool output trimming."
            ),
        )
        strip_channel_analysis: bool = Field(
            default=True,
            description=(
                "Strip <|channel|>analysis<|message|>... blocks "
                "(interleaved thinking format used by some models)."
            ),
        )
        min_messages_before_strip: int = Field(
            default=6,
            ge=0,
            description=(
                "Minimum number of messages in the conversation before stripping begins. "
                "Prevents stripping in very short chats where it provides no benefit. "
                "Default 6 means stripping starts after ~3 turns. Set to 0 to always strip."
            ),
        )
        show_status: bool = Field(
            default=True,
            description="Show a status notification about how many reasoning blocks were stripped.",
        )
        debug_mode: bool = Field(
            default=False,
            description="Enable detailed logging for debugging.",
        )

    def __init__(self):
        self.valves = self.Valves()

    def _strip_reasoning_from_content(self, content: str) -> tuple[str, int]:
        """
        Remove reasoning/thinking blocks from a single content string.

        Returns:
            (cleaned_content, blocks_removed_count)
        """
        if not isinstance(content, str) or not content:
            return content, 0

        original = content
        blocks_removed = 0

        # 1. Strip <details type="reasoning/thinking/thought">
        if self.valves.strip_reasoning_details and not self.valves.strip_all_details:
            new_content, count = REASONING_DETAILS_PATTERN.subn('', content)
            if count > 0:
                blocks_removed += count
                content = new_content

        # 2. Strip ALL <details> blocks (aggressive mode)
        if self.valves.strip_all_details:
            new_content, count = ALL_DETAILS_PATTERN.subn('', content)
            if count > 0:
                blocks_removed += count
                content = new_content

        # 3. Strip <think>...</think>
        if self.valves.strip_think_tags:
            new_content, count = THINK_TAG_PATTERN.subn('', content)
            if count > 0:
                blocks_removed += count
                content = new_content

        # 4. Strip <thinking>...</thinking>
        if self.valves.strip_think_tags:
            new_content, count = THINKING_TAG_PATTERN.subn('', content)
            if count > 0:
                blocks_removed += count
                content = new_content

        # 5. Strip <|channel|>analysis blocks
        if self.valves.strip_channel_analysis:
            new_content, count = CHANNEL_ANALYSIS_PATTERN.subn('', content)
            if count > 0:
                blocks_removed += count
                content = new_content

        # Clean up excessive whitespace left behind by removal
        if blocks_removed > 0:
            # Remove multiple consecutive blank lines (max 2 newlines)
            content = re.sub(r'\n{3,}', '\n\n', content)
            # Remove leading/trailing whitespace from the whole content
            content = content.strip()
            # If we stripped everything and content is now empty, leave a minimal marker
            # so the message doesn't become an empty assistant message
            if not content:
                content = "[reasoning stripped]"

        return content, blocks_removed

    def _should_protect_message(self, message: dict) -> bool:
        """
        Return True if this message should NOT have its reasoning stripped.

        Protection rules:
        1. Non-assistant messages (user, system, tool) are never modified
        2. Messages with tool_calls are protected (reasoning before tool use matters)
        3. Messages that are already summary markers are protected
        """
        role = message.get("role", "")

        # Never strip from non-assistant messages
        if role != "assistant":
            return True

        # Never strip from messages with tool_calls
        # (the reasoning before a tool call may be part of the decision chain)
        tool_calls = message.get("tool_calls")
        if isinstance(tool_calls, list) and tool_calls:
            return True

        # Never strip from summary markers (from Async Context Compression)
        metadata = message.get("metadata", {})
        if isinstance(metadata, dict) and metadata.get("is_summary"):
            return True

        return False

    def _extract_text_content(self, content: Any) -> str:
        """Extract text from string or multimodal list content."""
        if isinstance(content, str):
            return content
        if isinstance(content, list):
            text_parts = []
            for part in content:
                if isinstance(part, dict):
                    text_value = part.get("text")
                    if isinstance(text_value, str) and text_value:
                        text_parts.append(text_value)
            return " ".join(text_parts)
        if isinstance(content, dict):
            text_value = content.get("text") or content.get("content", "")
            return str(text_value) if text_value else ""
        return str(content) if content is not None else ""

    def _strip_reasoning_from_multimodal(self, content: list) -> tuple[list, int]:
        """
        Strip reasoning from multimodal content (list of parts).
        Only modifies text parts; leaves image/audio/etc. parts untouched.

        Returns:
            (cleaned_content_list, blocks_removed_count)
        """
        total_blocks = 0
        cleaned = []
        for part in content:
            if isinstance(part, dict) and part.get("type") == "text":
                text = part.get("text", "")
                if isinstance(text, str):
                    new_text, count = self._strip_reasoning_from_content(text)
                    total_blocks += count
                    cleaned.append({**part, "text": new_text})
                else:
                    cleaned.append(part)
            else:
                cleaned.append(part)
        return cleaned, total_blocks

    async def inlet(
        self,
        body: dict,
        __user__: Optional[dict] = None,
        __metadata__: Optional[dict] = None,
        __event_emitter__: Optional[Callable[[Any], Awaitable[None]]] = None,
        __event_call__: Optional[Callable[[Any], Awaitable[Any]]] = None,
    ) -> dict:
        """
        Process messages before they are sent to the LLM.
        Strip reasoning blocks from older assistant messages.
        """
        # Master switch check
        if not self.valves.enabled:
            return body

        messages = body.get("messages", [])
        if not messages:
            return body

        # Check minimum message count
        if len(messages) < self.valves.min_messages_before_strip:
            if self.valves.debug_mode:
                logger.info(
                    f"[Reasoning Stripper] Skipping: only {len(messages)} messages "
                    f"(min={self.valves.min_messages_before_strip})"
                )
            return body

        # Calculate the cutoff index: messages before this index are eligible for stripping
        # The last `protect_recent` messages are always protected
        protect_recent = self.valves.protect_recent
        cutoff_index = max(0, len(messages) - protect_recent)

        if self.valves.debug_mode:
            logger.info(
                f"[Reasoning Stripper] Processing {len(messages)} messages, "
                f"cutoff_index={cutoff_index}, protect_recent={protect_recent}"
            )

        total_blocks_stripped = 0
        messages_modified = 0
        chars_removed = 0

        for i, message in enumerate(messages):
            # Skip messages in the protected zone (last N messages)
            if i >= cutoff_index:
                if self.valves.debug_mode:
                    logger.info(
                        f"[Reasoning Stripper] Protecting message {i} "
                        f"(in recent zone, cutoff={cutoff_index})"
                    )
                continue

            # Check if this message type should be protected
            if self._should_protect_message(message):
                continue

            content = message.get("content")

            # Handle string content
            if isinstance(content, str):
                original_len = len(content)
                new_content, blocks = self._strip_reasoning_from_content(content)
                if blocks > 0:
                    message["content"] = new_content
                    total_blocks_stripped += blocks
                    messages_modified += 1
                    chars_removed += original_len - len(new_content)
                    if self.valves.debug_mode:
                        logger.info(
                            f"[Reasoning Stripper] Message {i}: stripped {blocks} "
                            f"reasoning blocks ({original_len - len(new_content)} chars)"
                        )

            # Handle multimodal (list) content
            elif isinstance(content, list):
                original_text = self._extract_text_content(content)
                original_len = len(original_text)
                new_content, blocks = self._strip_reasoning_from_multimodal(content)
                if blocks > 0:
                    message["content"] = new_content
                    total_blocks_stripped += blocks
                    messages_modified += 1
                    new_text = self._extract_text_content(new_content)
                    chars_removed += original_len - len(new_text)
                    if self.valves.debug_mode:
                        logger.info(
                            f"[Reasoning Stripper] Message {i} (multimodal): stripped "
                            f"{blocks} reasoning blocks ({original_len - len(new_text)} chars)"
                        )

        # Emit status notification
        if total_blocks_stripped > 0 and self.valves.show_status and __event_emitter__:
            status_msg = (
                f"🧹 Stripped {total_blocks_stripped} reasoning block(s) "
                f"from {messages_modified} message(s) "
                f"({chars_removed:,} chars saved)"
            )
            await __event_emitter__(
                {
                    "type": "status",
                    "data": {
                        "description": status_msg,
                        "done": True,
                    },
                }
            )

        if self.valves.debug_mode:
            logger.info(
                f"[Reasoning Stripper] Complete: {total_blocks_stripped} blocks stripped, "
                f"{messages_modified} messages modified, {chars_removed} chars removed"
            )

        return body