"""
Streaming guardrail for AI Service.
[AC-AISVC-82] Streaming mode forbidden word detection with sliding window buffer.
"""

import logging
from collections.abc import AsyncIterator
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

from app.models.entities import (
    ForbiddenWord,
    ForbiddenWordStrategy,
)
from app.services.guardrail.word_service import ForbiddenWordService

logger = logging.getLogger(__name__)


class StreamingGuardrailState(str, Enum):
    """State of streaming guardrail."""
    ACTIVE = "active"
    BLOCKED = "blocked"
    COMPLETED = "completed"


@dataclass
class StreamingGuardrailResult:
    """Result from streaming guardrail processing."""
    delta: str
    should_stop: bool = False
    fallback_reply: str | None = None
    triggered_words: list[str] = field(default_factory=list)
    triggered_categories: list[str] = field(default_factory=list)


class StreamingGuardrail:
    """
    [AC-AISVC-82] Streaming guardrail with sliding window buffer.

    Features:
    - Maintains a sliding window buffer for forbidden word detection
    - Detects forbidden words across chunk boundaries
    - Applies mask/replace strategies incrementally
    - Stops streaming and returns fallback for block strategy
    - Final check at stream end
    """

    DEFAULT_FALLBACK_REPLY = "抱歉，让我换个方式回答您"
    DEFAULT_WINDOW_SIZE = 50

    def __init__(
        self,
        word_service: ForbiddenWordService,
        window_size: int = DEFAULT_WINDOW_SIZE,
    ):
        self._word_service = word_service
        self._window_size = window_size
        self._buffer = ""
        self._state = StreamingGuardrailState.ACTIVE
        self._triggered_words: list[str] = []
        self._triggered_categories: list[str] = []
        self._fallback_reply = self.DEFAULT_FALLBACK_REPLY
        self._words_cache: list[ForbiddenWord] | None = None
        self._max_word_length = 0

    async def initialize(self, tenant_id: str) -> None:
        """
        Initialize the guardrail with tenant's forbidden words.
        Must be called before processing chunks.
        """
        self._words_cache = await self._word_service.get_enabled_words_for_filtering(tenant_id)

        if self._words_cache:
            self._max_word_length = max(len(w.word) for w in self._words_cache)
            effective_window = max(self._window_size, self._max_word_length)
            if effective_window > self._window_size:
                logger.info(
                    f"[AC-AISVC-82] Adjusted window size from {self._window_size} to {effective_window} "
                    f"to accommodate max word length {self._max_word_length}"
                )
                self._window_size = effective_window

        logger.debug(
            f"[AC-AISVC-82] StreamingGuardrail initialized: "
            f"words_count={len(self._words_cache) if self._words_cache else 0}, "
            f"window_size={self._window_size}"
        )

    async def process_chunk(
        self,
        delta: str,
        tenant_id: str,
    ) -> StreamingGuardrailResult:
        """
        Process a streaming chunk with sliding window detection.

        Args:
            delta: New text chunk from LLM
            tenant_id: Tenant ID for isolation

        Returns:
            StreamingGuardrailResult with processed delta and state
        """
        if self._state == StreamingGuardrailState.BLOCKED:
            return StreamingGuardrailResult(
                delta="",
                should_stop=True,
                fallback_reply=self._fallback_reply,
            )

        if self._state == StreamingGuardrailState.COMPLETED:
            return StreamingGuardrailResult(delta="")

        if not self._words_cache:
            await self.initialize(tenant_id)

        if not self._words_cache:
            return StreamingGuardrailResult(delta=delta)

        self._buffer += delta

        result = self._check_buffer()

        if result.should_stop:
            self._state = StreamingGuardrailState.BLOCKED
            logger.warning(
                f"[AC-AISVC-82] Streaming blocked: tenant={tenant_id}, "
                f"triggered_words={result.triggered_words}"
            )

        return result

    def _check_buffer(self) -> StreamingGuardrailResult:
        """
        Check buffer for forbidden words and apply strategies.
        """
        triggered_words: list[str] = []
        triggered_categories: list[str] = []
        filtered_buffer = self._buffer
        blocked = False
        fallback_reply = self.DEFAULT_FALLBACK_REPLY

        if not self._words_cache:
            return StreamingGuardrailResult(delta=self._buffer)

        for word in self._words_cache:
            if word.word in filtered_buffer:
                triggered_words.append(word.word)
                if word.category not in triggered_categories:
                    triggered_categories.append(word.category)

                if word.strategy == ForbiddenWordStrategy.BLOCK.value:
                    blocked = True
                    fallback_reply = word.fallback_reply or self.DEFAULT_FALLBACK_REPLY
                    break

                elif word.strategy == ForbiddenWordStrategy.MASK.value:
                    filtered_buffer = filtered_buffer.replace(word.word, "*" * len(word.word))

                elif word.strategy == ForbiddenWordStrategy.REPLACE.value:
                    replacement = word.replacement or ""
                    filtered_buffer = filtered_buffer.replace(word.word, replacement)

        self._triggered_words.extend([w for w in triggered_words if w not in self._triggered_words])
        self._triggered_categories.extend([c for c in triggered_categories if c not in self._triggered_categories])

        if blocked:
            return StreamingGuardrailResult(
                delta="",
                should_stop=True,
                fallback_reply=fallback_reply,
                triggered_words=triggered_words,
                triggered_categories=triggered_categories,
            )

        safe_output, remaining = self._split_buffer(filtered_buffer)

        self._buffer = remaining

        return StreamingGuardrailResult(
            delta=safe_output,
            should_stop=False,
            triggered_words=triggered_words,
            triggered_categories=triggered_categories,
        )

    def _split_buffer(self, buffer: str) -> tuple[str, str]:
        """
        Split buffer into safe output and remaining buffer.

        Safe output: characters that are definitely safe (before window boundary)
        Remaining: characters that might be part of a forbidden word
        """
        if len(buffer) <= self._window_size:
            return "", buffer

        safe_length = len(buffer) - self._window_size
        safe_output = buffer[:safe_length]
        remaining = buffer[safe_length:]

        return safe_output, remaining

    async def finalize(self, tenant_id: str) -> StreamingGuardrailResult:
        """
        Finalize the stream and process remaining buffer.
        Must be called at the end of streaming.

        Args:
            tenant_id: Tenant ID for isolation

        Returns:
            StreamingGuardrailResult with final delta
        """
        if self._state == StreamingGuardrailState.BLOCKED:
            self._state = StreamingGuardrailState.COMPLETED
            return StreamingGuardrailResult(
                delta="",
                should_stop=True,
                fallback_reply=self._fallback_reply,
                triggered_words=self._triggered_words,
                triggered_categories=self._triggered_categories,
            )

        if not self._words_cache:
            await self.initialize(tenant_id)

        result = self._final_check()

        self._state = StreamingGuardrailState.COMPLETED

        if self._triggered_words:
            logger.info(
                f"[AC-AISVC-82] Streaming finalized with triggers: tenant={tenant_id}, "
                f"triggered_words={self._triggered_words}"
            )

            for word_entity in self._get_triggered_word_entities():
                try:
                    await self._word_service.increment_hit_count(tenant_id, word_entity.id)
                except Exception as e:
                    logger.warning(
                        f"Failed to increment hit count for word {word_entity.id}: {e}"
                    )

        return result

    def _final_check(self) -> StreamingGuardrailResult:
        """
        Final check of remaining buffer.
        """
        if not self._buffer:
            return StreamingGuardrailResult(
                delta="",
                triggered_words=self._triggered_words,
                triggered_categories=self._triggered_categories,
            )

        if not self._words_cache:
            return StreamingGuardrailResult(
                delta=self._buffer,
                triggered_words=self._triggered_words,
                triggered_categories=self._triggered_categories,
            )

        filtered_buffer = self._buffer
        blocked = False
        fallback_reply = self.DEFAULT_FALLBACK_REPLY

        for word in self._words_cache:
            if word.word in filtered_buffer:
                if word.word not in self._triggered_words:
                    self._triggered_words.append(word.word)
                if word.category not in self._triggered_categories:
                    self._triggered_categories.append(word.category)

                if word.strategy == ForbiddenWordStrategy.BLOCK.value:
                    blocked = True
                    fallback_reply = word.fallback_reply or self.DEFAULT_FALLBACK_REPLY
                    break

                elif word.strategy == ForbiddenWordStrategy.MASK.value:
                    filtered_buffer = filtered_buffer.replace(word.word, "*" * len(word.word))

                elif word.strategy == ForbiddenWordStrategy.REPLACE.value:
                    replacement = word.replacement or ""
                    filtered_buffer = filtered_buffer.replace(word.word, replacement)

        if blocked:
            return StreamingGuardrailResult(
                delta="",
                should_stop=True,
                fallback_reply=fallback_reply,
                triggered_words=self._triggered_words,
                triggered_categories=self._triggered_categories,
            )

        self._buffer = ""

        return StreamingGuardrailResult(
            delta=filtered_buffer,
            triggered_words=self._triggered_words,
            triggered_categories=self._triggered_categories,
        )

    def _get_triggered_word_entities(self) -> list[ForbiddenWord]:
        """Get word entities for triggered words."""
        if not self._words_cache:
            return []
        return [w for w in self._words_cache if w.word in self._triggered_words]

    def get_triggered_info(self) -> dict[str, Any]:
        """Get triggered words and categories info."""
        return {
            "triggered_words": self._triggered_words,
            "triggered_categories": self._triggered_categories,
            "guardrail_triggered": len(self._triggered_words) > 0,
        }

    def reset(self) -> None:
        """Reset the guardrail state for reuse."""
        self._buffer = ""
        self._state = StreamingGuardrailState.ACTIVE
        self._triggered_words = []
        self._triggered_categories = []
        self._fallback_reply = self.DEFAULT_FALLBACK_REPLY
        self._words_cache = None
        self._max_word_length = 0


async def wrap_stream_with_guardrail(
    stream: AsyncIterator[str],
    guardrail: StreamingGuardrail,
    tenant_id: str,
) -> AsyncIterator[tuple[str, bool, str | None]]:
    """
    Wrap an async stream with guardrail processing.

    Args:
        stream: Original LLM output stream
        guardrail: StreamingGuardrail instance
        tenant_id: Tenant ID for isolation

    Yields:
        Tuple of (delta, should_stop, fallback_reply)
    """
    await guardrail.initialize(tenant_id)

    async for delta in stream:
        result = await guardrail.process_chunk(delta, tenant_id)

        if result.delta:
            yield (result.delta, False, None)

        if result.should_stop:
            yield ("", True, result.fallback_reply)
            return

    final_result = await guardrail.finalize(tenant_id)

    if final_result.delta:
        yield (final_result.delta, False, None)

    if final_result.should_stop:
        yield ("", True, final_result.fallback_reply)
