from __future__ import annotations

import re
from collections.abc import AsyncIterator, Iterator
from dataclasses import dataclass
from typing import Literal

from sqlalchemy import select
from sqlalchemy.orm import Session

from app.models.ai_config import SystemConfig


REASONING_VISIBILITY_KEY = "show_model_reasoning"
_OPEN_TAG = re.compile(r"^<think(?:\s[^>]*)?>", re.IGNORECASE)
_CLOSE_TAG = re.compile(r"^</think\s*>", re.IGNORECASE)


@dataclass(frozen=True)
class ReasoningSegment:
    kind: Literal["reasoning", "content"]
    content: str


class ReasoningStreamParser:
    """Incrementally separates model reasoning tags from answer content."""

    def __init__(self) -> None:
        self._buffer = ""
        self._depth = 0

    def feed(self, chunk: str) -> list[ReasoningSegment]:
        if not chunk:
            return []
        self._buffer += chunk
        return list(self._drain(final=False))

    def finish(self) -> list[ReasoningSegment]:
        return list(self._drain(final=True))

    def _drain(self, *, final: bool) -> Iterator[ReasoningSegment]:
        while self._buffer:
            tag_start = self._buffer.find("<")
            if tag_start < 0:
                yield self._segment(self._buffer)
                self._buffer = ""
                return
            if tag_start > 0:
                yield self._segment(self._buffer[:tag_start])
                self._buffer = self._buffer[tag_start:]
                continue

            open_tag = _OPEN_TAG.match(self._buffer)
            if open_tag:
                self._depth += 1
                self._buffer = self._buffer[open_tag.end():]
                continue
            close_tag = _CLOSE_TAG.match(self._buffer)
            if close_tag:
                self._depth = max(0, self._depth - 1)
                self._buffer = self._buffer[close_tag.end():]
                continue

            lower = self._buffer.lower()
            incomplete_tag = (
                "<think".startswith(lower)
                or "</think>".startswith(lower)
                or (lower.startswith("<think") and ">" not in lower)
                or (lower.startswith("</think") and ">" not in lower)
            )
            if incomplete_tag and not final:
                return
            if incomplete_tag and final:
                self._buffer = ""
                return

            yield self._segment("<")
            self._buffer = self._buffer[1:]

    def _segment(self, content: str) -> ReasoningSegment:
        return ReasoningSegment("reasoning" if self._depth else "content", content)


class ReasoningPolicyService:
    @staticmethod
    def is_visible(db: Session) -> bool:
        config = db.scalar(select(SystemConfig).where(SystemConfig.config_key == REASONING_VISIBILITY_KEY))
        if config is None:
            return False
        return config.config_value.strip().lower() in {"1", "true", "yes", "on"}

    @staticmethod
    def set_visible(db: Session, visible: bool, admin_id: int | None) -> SystemConfig:
        config = db.scalar(select(SystemConfig).where(SystemConfig.config_key == REASONING_VISIBILITY_KEY))
        if config is None:
            config = SystemConfig(config_key=REASONING_VISIBILITY_KEY, config_value="false")
        config.config_value = "true" if visible else "false"
        config.description = "是否向用户端及后台预览展示模型思考过程"
        config.updated_by = admin_id
        db.add(config)
        return config

    @staticmethod
    def split_text(text: str) -> tuple[str, str]:
        parser = ReasoningStreamParser()
        segments = [*parser.feed(text), *parser.finish()]
        answer = "".join(item.content for item in segments if item.kind == "content")
        reasoning = "".join(item.content for item in segments if item.kind == "reasoning")
        return answer, reasoning

    @classmethod
    def strip_reasoning(cls, text: str) -> str:
        return cls.split_text(text)[0]

    @staticmethod
    async def iter_segments(chunks: AsyncIterator[str]) -> AsyncIterator[ReasoningSegment]:
        parser = ReasoningStreamParser()
        async for chunk in chunks:
            for segment in parser.feed(chunk):
                if segment.content:
                    yield segment
        for segment in parser.finish():
            if segment.content:
                yield segment
