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, TopicSession
from app.models.knowledge import KnowledgeRetrievalLog
from app.models.user import User
from app.core.auth_context import ChatAccessScope
from app.services.ai_request_log_service import AiRequestLogService
from app.services.chat_service import ChatService, _title_from_question
from app.services.chat_context_service import ChatContextService
from app.services.entitlement_service import EntitlementService, entitlement_prompt_context
from app.services.external_errors import ExternalServiceError
from app.services.growth_profile_service import GrowthProfileService
from app.services.human_attention_service import HumanAttentionService
from app.services.model_stream_service import ModelStreamService
from app.services.model_routing_service import ModelRoutingService
from app.services.rag_async_service import AsyncRagService
from app.services.rag_service import RagService
from app.services.topic_session_service import TopicSessionService


class ChatStreamService:
    @staticmethod
    def stream_answer(
        db: Session,
        user: User,
        session_id: int,
        question: str,
        scope: ChatAccessScope | None = None,
    ) -> Iterator[str]:
        user = ChatService.prepare_daily_quota(db, user)
        session = ChatService._get_user_session(db, user, session_id, scope)
        ChatService._ensure_quota(user)
        entitlement = EntitlementService.active_entitlement(
            db,
            user,
            monthly_topic_used=TopicSessionService.monthly_used_count(db, user.id),
        )
        ChatService._ensure_topic_quota(db, user, session, entitlement)

        now = _now()
        normalized_question = question.strip()
        topic = TopicSessionService.get_or_create_active(
            db,
            user=user,
            session=session,
            question=normalized_question,
            deduct_quota=entitlement.deduct_quota,
        )
        user_message = ChatMessage(
            session_id=session.id,
            topic_session_id=topic.id,
            user_id=user.id,
            role="user",
            content=normalized_question,
            message_status="FINISHED",
            created_at=now,
        )
        db.add(user_message)
        TopicSessionService.attach_user_message(user_message, topic)
        db.flush()

        history = list(
            db.scalars(
                select(ChatMessage)
                .where(
                    ChatMessage.session_id == session.id,
                    ChatMessage.topic_session_id == topic.id,
                    ChatMessage.user_id == user.id,
                    ChatMessage.id < user_message.id,
                )
                .order_by(ChatMessage.created_at.asc(), ChatMessage.id.asc())
            )
        )
        summary_result = ChatContextService.update_summary(db, session, history)
        context_trace = [summary_result.trace] if summary_result.trace else None
        growth_context = GrowthProfileService.prompt_context(db, user) if entitlement.enable_growth_profile else None
        if growth_context:
            context_trace = list(context_trace or [])
            context_trace.append({"tool": "load_recent_practice_review", "status": "success", "count": 1})
        topic_context, topic_summary_used = TopicSessionService.prompt_context(db, topic)
        product_context = entitlement_prompt_context(entitlement)
        context_trace = _append_runtime_context_trace(
            context_trace,
            topic=topic,
            topic_summary_used=topic_summary_used,
            entitlement=entitlement,
        )

        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),
                getattr(session, "summary_up_to_message_id", None),
                context_trace,
                growth_context,
                topic_context,
                product_context,
            )
            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)
            failed_route = _model_log_context(db, rag_result, model_response)
            AiRequestLogService.write_failed(
                db,
                session_id=session.id,
                message_id=user_message.id,
                user_id=user.id,
                model_name=failed_route[3],
                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,
                model_id=failed_route[2],
                route_reason=failed_route[0],
                question_type=failed_route[1],
            )
            _mark_retrieval_failed(db, rag_result, str(exc), cost_ms)
            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)
            failed_route = _model_log_context(db, rag_result, model_response)
            AiRequestLogService.write_failed(
                db,
                session_id=session.id,
                message_id=user_message.id,
                user_id=user.id,
                model_name=failed_route[3],
                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,
                model_id=failed_route[2],
                route_reason=failed_route[0],
                question_type=failed_route[1],
            )
            _mark_retrieval_failed(db, rag_result, "模型未返回有效内容", cost_ms)
            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,
            topic_session_id=topic.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)
        TopicSessionService.attach_assistant_message(
            assistant_message,
            topic,
            token_input=model_response.input_token if model_response is not None else None,
            token_output=_rough_token_count(answer),
        )
        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,
            model_id=model_response.model_id if model_response is not None else None,
            route_reason=model_response.route_reason if model_response is not None else None,
            question_type=model_response.question_type if model_response is not None else None,
        )
        db.commit()

    @staticmethod
    async def stream_answer_async(
        db: Session,
        user: User,
        session_id: int,
        question: str,
        *,
        retry_failed_question: bool = False,
        scope: ChatAccessScope | None = None,
    ) -> AsyncIterator[str]:
        user = ChatService.prepare_daily_quota(db, user)
        session = ChatService._get_user_session(db, user, session_id, scope)
        ChatService._ensure_quota(user)
        entitlement = EntitlementService.active_entitlement(
            db,
            user,
            monthly_topic_used=TopicSessionService.monthly_used_count(db, user.id),
        )
        ChatService._ensure_topic_quota(db, user, session, entitlement)

        now = _now()
        normalized_question = question.strip()
        user_message = (
            _retryable_user_message(
                db,
                session_id=session.id,
                user_id=user.id,
                question=normalized_question,
            )
            if retry_failed_question
            else None
        )
        topic = db.get(TopicSession, user_message.topic_session_id) if user_message and user_message.topic_session_id else None
        if user_message is None or topic is None:
            topic = TopicSessionService.get_or_create_active(
                db,
                user=user,
                session=session,
                question=normalized_question,
                deduct_quota=entitlement.deduct_quota,
            )
            user_message = ChatMessage(
                session_id=session.id,
                topic_session_id=topic.id,
                user_id=user.id,
                role="user",
                content=normalized_question,
                message_status="FINISHED",
                created_at=now,
            )
            db.add(user_message)
            TopicSessionService.attach_user_message(user_message, topic)
            db.flush()

        history = list(
            db.scalars(
                select(ChatMessage)
                .where(
                    ChatMessage.session_id == session.id,
                    ChatMessage.topic_session_id == topic.id,
                    ChatMessage.user_id == user.id,
                    ChatMessage.id < user_message.id,
                )
                .order_by(ChatMessage.created_at.asc(), ChatMessage.id.asc())
            )
        )
        summary_result = await ChatContextService.update_summary_async(db, session, history)
        context_trace = [summary_result.trace] if summary_result.trace else None
        growth_context = GrowthProfileService.prompt_context(db, user) if entitlement.enable_growth_profile else None
        if growth_context:
            context_trace = list(context_trace or [])
            context_trace.append({"tool": "load_recent_practice_review", "status": "success", "count": 1})
        topic_context, topic_summary_used = TopicSessionService.prompt_context(db, topic)
        product_context = entitlement_prompt_context(entitlement)
        context_trace = _append_runtime_context_trace(
            context_trace,
            topic=topic,
            topic_summary_used=topic_summary_used,
            entitlement=entitlement,
        )

        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),
                summary_up_to_message_id=getattr(session, "summary_up_to_message_id", None),
                session_id=session.id,
                context_trace=context_trace,
                growth_context=growth_context,
                topic_context=topic_context,
                product_context=product_context,
            )
            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)
            failed_route = _model_log_context(db, rag_result, model_response)
            AiRequestLogService.write_failed(
                db,
                session_id=session.id,
                message_id=user_message.id,
                user_id=user.id,
                model_name=failed_route[3],
                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,
                model_id=failed_route[2],
                route_reason=failed_route[0],
                question_type=failed_route[1],
            )
            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)
            failed_route = _model_log_context(db, rag_result, model_response)
            AiRequestLogService.write_failed(
                db,
                session_id=session.id,
                message_id=user_message.id,
                user_id=user.id,
                model_name=failed_route[3],
                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,
                model_id=failed_route[2],
                route_reason=failed_route[0],
                question_type=failed_route[1],
            )
            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,
            topic=topic,
        )


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


def _retryable_user_message(
    db: Session,
    *,
    session_id: int,
    user_id: int,
    question: str,
) -> ChatMessage | None:
    """Reuse only the latest unanswered user turn to avoid duplicate retry messages."""
    latest = db.scalar(
        select(ChatMessage)
        .where(ChatMessage.session_id == session_id, ChatMessage.user_id == user_id)
        .order_by(ChatMessage.id.desc())
        .limit(1)
    )
    if latest is None or latest.role != "user":
        return None
    if latest.content.strip() != question.strip():
        return None
    return latest


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,
    topic,
) -> None:
    cost_ms = int((perf_counter() - started_at) * 1000)
    assistant_message = ChatMessage(
        session_id=session.id,
        topic_session_id=topic.id if topic is not None else None,
        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)
    if topic is not None:
        TopicSessionService.attach_assistant_message(
            assistant_message,
            topic,
            token_input=model_response.input_token if model_response is not None else None,
            token_output=_rough_token_count(answer),
        )
    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,
        model_id=model_response.model_id if model_response is not None else None,
        route_reason=model_response.route_reason if model_response is not None else None,
        question_type=model_response.question_type if model_response is not None else None,
    )
    if rag_result is not None and rag_result.retrieval_log_id:
        retrieval_log = db.get(KnowledgeRetrievalLog, rag_result.retrieval_log_id)
        if retrieval_log:
            attention = HumanAttentionService.create_if_needed(
                db,
                session_id=session.id,
                message_id=user_message_id_for_attention(db, session.id, assistant_message.id),
                user_id=user.id,
                question=question,
                answer=answer,
                knowledge_missing=not rag_result.allow_general_knowledge and not rag_result.is_hit,
            )
            retrieval_log.message_id = assistant_message.id
            retrieval_log.final_answer = answer
            retrieval_log.status = "success"
            retrieval_log.total_cost_ms = cost_ms
            retrieval_log.attention_created = 1 if attention else 0
            db.add(retrieval_log)
    db.commit()


def _model_log_context(
    db: Session,
    rag_result,
    model_response,
) -> tuple[str | None, str | None, int | None, str | None]:
    if model_response is not None:
        return (
            model_response.route_reason,
            model_response.question_type,
            model_response.model_id,
            model_response.model_name,
        )
    route = ModelRoutingService.resolve_chat(
        db,
        [chunk.knowledge_type for chunk in rag_result.chunks] if rag_result is not None else [],
        rag_result.question if rag_result is not None else "",
    )
    return (
        route.reason,
        route.question_type,
        route.model.id if route.model is not None else None,
        route.model.model_name if route.model is not None else None,
    )


def user_message_id_for_attention(db: Session, session_id: int, assistant_message_id: int) -> int:
    message_id = db.scalar(
        select(ChatMessage.id)
        .where(
            ChatMessage.session_id == session_id,
            ChatMessage.role == "user",
            ChatMessage.id < assistant_message_id,
        )
        .order_by(ChatMessage.id.desc())
        .limit(1)
    )
    return int(message_id or assistant_message_id)


def _mark_retrieval_failed(db: Session, rag_result, error_message: str, cost_ms: int) -> None:
    if rag_result is None or not rag_result.retrieval_log_id:
        return
    retrieval_log = db.get(KnowledgeRetrievalLog, rag_result.retrieval_log_id)
    if retrieval_log:
        retrieval_log.status = "failed"
        retrieval_log.error_message = error_message
        retrieval_log.total_cost_ms = cost_ms
        db.add(retrieval_log)


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


def _append_runtime_context_trace(context_trace, *, topic, topic_summary_used: bool, entitlement) -> list[dict]:
    result = list(context_trace or [])
    result.extend(
        [
            {
                "tool": "load_topic_context",
                "status": "success",
                "response": {"topicSessionId": topic.id, "summaryUsed": topic_summary_used},
            },
            {
                "tool": "load_product_entitlement",
                "status": "success",
                "response": {
                    "planId": entitlement.plan_id,
                    "allowHelpCard": entitlement.allow_help_card,
                    "allowShareDraft": entitlement.allow_share_draft,
                },
            },
        ]
    )
    return result
