from __future__ import annotations

import asyncio
from collections.abc import AsyncIterator, Iterator
from datetime import UTC, datetime
from inspect import signature
from time import perf_counter

from fastapi import HTTPException, status
from sqlalchemy import select
from sqlalchemy.orm import Session

from app.models.chat import ChatMessage
from app.models.user import User
from app.services.ai_request_log_service import AiRequestLogService
from app.services.chat_service import ChatService, _title_from_question
from app.services.external_errors import ExternalServiceError
from app.services.model_stream_service import ModelStreamService
from app.services.rag_async_service import AsyncRagService
from app.services.rag_service import RagService


class ChatStreamService:
    @staticmethod
    def stream_answer(db: Session, user: User, session_id: int, question: str) -> Iterator[str]:
        session = ChatService._get_user_session(db, user, session_id)
        ChatService._ensure_quota(user)

        now = _now()
        normalized_question = question.strip()
        user_message = ChatMessage(
            session_id=session.id,
            user_id=user.id,
            role="user",
            content=normalized_question,
            message_status="FINISHED",
            created_at=now,
        )
        db.add(user_message)
        db.flush()

        history = list(
            db.scalars(
                select(ChatMessage)
                .where(
                    ChatMessage.session_id == session.id,
                    ChatMessage.user_id == user.id,
                    ChatMessage.id < user_message.id,
                )
                .order_by(ChatMessage.created_at.asc(), ChatMessage.id.asc())
            )
        )
        _maybe_update_summary(db, session, history)

        started_at = perf_counter()
        rag_result = None
        model_response = None
        answer_parts: list[str] = []

        try:
            rag_result = _build_rag_result(db, user, normalized_question, history, getattr(session, "summary", None))
            model_response = ModelStreamService.stream(db, rag_result)
            for chunk in model_response.chunks:
                if chunk:
                    answer_parts.append(chunk)
                    yield chunk
        except GeneratorExit:
            db.rollback()
            raise
        except ExternalServiceError as exc:
            cost_ms = int((perf_counter() - started_at) * 1000)
            AiRequestLogService.write_failed(
                db,
                session_id=session.id,
                message_id=user_message.id,
                user_id=user.id,
                model_name=model_response.model_name if model_response is not None else None,
                prompt=rag_result.prompt if rag_result is not None else normalized_question,
                knowledge_ids=rag_result.knowledge_ids if rag_result is not None else None,
                retrieve_count=len(rag_result.chunks) if rag_result is not None else 0,
                cost_ms=cost_ms,
                error_message=str(exc),
                retrieved_chunks=rag_result.chunks if rag_result is not None else None,
            )
            db.commit()
            raise HTTPException(status_code=status.HTTP_502_BAD_GATEWAY, detail=str(exc)) from exc

        answer = "".join(answer_parts)
        if not answer.strip():
            cost_ms = int((perf_counter() - started_at) * 1000)
            AiRequestLogService.write_failed(
                db,
                session_id=session.id,
                message_id=user_message.id,
                user_id=user.id,
                model_name=model_response.model_name if model_response is not None else None,
                prompt=rag_result.prompt if rag_result is not None else normalized_question,
                knowledge_ids=rag_result.knowledge_ids if rag_result is not None else None,
                retrieve_count=len(rag_result.chunks) if rag_result is not None else 0,
                cost_ms=cost_ms,
                error_message="模型未返回有效内容",
                retrieved_chunks=rag_result.chunks if rag_result is not None else None,
            )
            db.commit()
            raise HTTPException(status_code=status.HTTP_502_BAD_GATEWAY, detail="模型未返回有效内容")

        cost_ms = int((perf_counter() - started_at) * 1000)
        assistant_message = ChatMessage(
            session_id=session.id,
            user_id=user.id,
            role="assistant",
            content=answer,
            message_status="FINISHED",
            token_input=model_response.input_token if model_response is not None else None,
            token_output=_rough_token_count(answer),
            response_time_ms=cost_ms,
            model_id=model_response.model_id if model_response is not None else None,
            created_at=now,
        )
        db.add(assistant_message)
        db.flush()

        session.message_count += 2
        session.last_message_at = now
        if session.title == "新聊天":
            session.title = _title_from_question(normalized_question)

        user.daily_chat_used += 1
        db.add_all([session, user])
        AiRequestLogService.write_success(
            db,
            session_id=session.id,
            message_id=assistant_message.id,
            user_id=user.id,
            model_name=model_response.model_name if model_response is not None else "unknown",
            prompt=rag_result.prompt if rag_result is not None else normalized_question,
            knowledge_ids=rag_result.knowledge_ids if rag_result is not None else "",
            retrieve_count=len(rag_result.chunks) if rag_result is not None else 0,
            input_token=model_response.input_token if model_response is not None else 0,
            output_token=_rough_token_count(answer),
            cost_ms=cost_ms,
            retrieved_chunks=rag_result.chunks if rag_result is not None else None,
        )
        db.commit()

    @staticmethod
    async def stream_answer_async(db: Session, user: User, session_id: int, question: str) -> AsyncIterator[str]:
        session = ChatService._get_user_session(db, user, session_id)
        ChatService._ensure_quota(user)

        now = _now()
        normalized_question = question.strip()
        user_message = ChatMessage(
            session_id=session.id,
            user_id=user.id,
            role="user",
            content=normalized_question,
            message_status="FINISHED",
            created_at=now,
        )
        db.add(user_message)
        db.flush()

        history = list(
            db.scalars(
                select(ChatMessage)
                .where(
                    ChatMessage.session_id == session.id,
                    ChatMessage.user_id == user.id,
                    ChatMessage.id < user_message.id,
                )
                .order_by(ChatMessage.created_at.asc(), ChatMessage.id.asc())
            )
        )
        _maybe_update_summary(db, session, history)

        started_at = perf_counter()
        rag_result = None
        model_response = None
        answer_parts: list[str] = []

        try:
            rag_result = await AsyncRagService.build_result(
                db,
                user,
                normalized_question,
                history=history,
                session_summary=getattr(session, "summary", None),
            )
            model_response = ModelStreamService.stream_async(db, rag_result)
            async for chunk in model_response.chunks:
                if chunk:
                    answer_parts.append(chunk)
                    yield chunk
        except asyncio.CancelledError:
            db.rollback()
            raise
        except ExternalServiceError as exc:
            cost_ms = int((perf_counter() - started_at) * 1000)
            AiRequestLogService.write_failed(
                db,
                session_id=session.id,
                message_id=user_message.id,
                user_id=user.id,
                model_name=model_response.model_name if model_response is not None else None,
                prompt=rag_result.prompt if rag_result is not None else normalized_question,
                knowledge_ids=rag_result.knowledge_ids if rag_result is not None else None,
                retrieve_count=len(rag_result.chunks) if rag_result is not None else 0,
                cost_ms=cost_ms,
                error_message=str(exc),
                retrieved_chunks=rag_result.chunks if rag_result is not None else None,
            )
            db.commit()
            raise HTTPException(status_code=status.HTTP_502_BAD_GATEWAY, detail=str(exc)) from exc

        answer = "".join(answer_parts)
        if not answer.strip():
            cost_ms = int((perf_counter() - started_at) * 1000)
            AiRequestLogService.write_failed(
                db,
                session_id=session.id,
                message_id=user_message.id,
                user_id=user.id,
                model_name=model_response.model_name if model_response is not None else None,
                prompt=rag_result.prompt if rag_result is not None else normalized_question,
                knowledge_ids=rag_result.knowledge_ids if rag_result is not None else None,
                retrieve_count=len(rag_result.chunks) if rag_result is not None else 0,
                cost_ms=cost_ms,
                error_message="模型未返回有效内容",
                retrieved_chunks=rag_result.chunks if rag_result is not None else None,
            )
            db.commit()
            raise HTTPException(status_code=status.HTTP_502_BAD_GATEWAY, detail="模型未返回有效内容")

        _write_success(
            db,
            user=user,
            session=session,
            question=normalized_question,
            answer=answer,
            rag_result=rag_result,
            model_response=model_response,
            started_at=started_at,
            now=now,
        )


def _now() -> datetime:
    return datetime.now(UTC).replace(tzinfo=None)


def _rough_token_count(text: str) -> int:
    return max(1, len(text.strip()) // 2)


def _write_success(
    db: Session,
    *,
    user: User,
    session,
    question: str,
    answer: str,
    rag_result,
    model_response,
    started_at: float,
    now: datetime,
) -> None:
    cost_ms = int((perf_counter() - started_at) * 1000)
    assistant_message = ChatMessage(
        session_id=session.id,
        user_id=user.id,
        role="assistant",
        content=answer,
        message_status="FINISHED",
        token_input=model_response.input_token if model_response is not None else None,
        token_output=_rough_token_count(answer),
        response_time_ms=cost_ms,
        model_id=model_response.model_id if model_response is not None else None,
        created_at=now,
    )
    db.add(assistant_message)
    db.flush()

    session.message_count += 2
    session.last_message_at = now
    if session.title == "新聊天":
        session.title = _title_from_question(question)

    user.daily_chat_used += 1
    db.add_all([session, user])
    AiRequestLogService.write_success(
        db,
        session_id=session.id,
        message_id=assistant_message.id,
        user_id=user.id,
        model_name=model_response.model_name if model_response is not None else "unknown",
        prompt=rag_result.prompt if rag_result is not None else question,
        knowledge_ids=rag_result.knowledge_ids if rag_result is not None else "",
        retrieve_count=len(rag_result.chunks) if rag_result is not None else 0,
        input_token=model_response.input_token if model_response is not None else 0,
        output_token=_rough_token_count(answer),
        cost_ms=cost_ms,
        retrieved_chunks=rag_result.chunks if rag_result is not None else None,
    )
    db.commit()


def _build_rag_result(db: Session, user: User, question: str, history: list[ChatMessage], summary: str | None):
    parameters = signature(RagService.build_result).parameters
    if "history" in parameters:
        return RagService.build_result(db, user, question, history=history, session_summary=summary)
    return RagService.build_result(db, user, question)


def _maybe_update_summary(db: Session, session, history: list[ChatMessage]) -> None:
    summary_handler = getattr(ChatService, "_maybe_update_summary", None)
    if not callable(summary_handler) or not hasattr(session, "summary_up_to_message_id"):
        return
    try:
        from app.services import rag_service

        trigger_rounds = getattr(rag_service, "SUMMARY_TRIGGER_ROUNDS", None)
        if trigger_rounds is None:
            return
        if len(history) // 2 >= trigger_rounds:
            summary_handler(db, session, history)
    except Exception:
        return
