from __future__ import annotations

import asyncio
from collections.abc import AsyncIterator
from types import SimpleNamespace

from sqlalchemy.orm import Session

from app.models.user import User
from app.models.admin import Admin
from app.models.ai_config import ModelConfig
from app.schemas.admin import AgentDebugRequest
from app.services.admin_service import OperationLogService
from app.services.entitlement_service import EntitlementService, entitlement_dict
from app.services.growth_profile_service import GrowthProfileService
from app.services.knowledge_agent_service import KnowledgeAgentService
from app.services.model_stream_service import ModelStreamService
from app.services.reasoning_policy_service import ReasoningPolicyService
from app.services.rag_service import RagResult
from app.services.topic_session_service import TopicSessionService


class AgentDebugService:
    @staticmethod
    async def build_result(db: Session, payload: AgentDebugRequest) -> RagResult:
        history = [
            SimpleNamespace(id=index + 1, role=item.role, content=item.content)
            for index, item in enumerate(payload.history)
        ]
        preview_knowledge_ids = payload.knowledgeIds or None
        version_overrides = payload.knowledgeVersions or None
        debug_context = AgentDebugService._debug_user_context(db, payload.userId)
        rag_result = await KnowledgeAgentService.build_result(
            db,
            question=payload.question,
            history=history,
            version_overrides=version_overrides,
            preview_knowledge_ids=preview_knowledge_ids,
            user_id=payload.userId,
            context_trace=debug_context["trace"],
            prompt_override=payload.promptContent,
            response_depth=payload.responseDepth,
            growth_context=debug_context["growth_context"],
        )
        return RagResult(
            question=rag_result.question,
            knowledge_scopes=rag_result.knowledge_scopes,
            chunks=rag_result.chunks,
            prompt=rag_result.prompt,
            allow_general_knowledge=rag_result.allow_general_knowledge,
            retrieval_log_id=rag_result.retrieval_log_id,
            tool_trace=rag_result.tool_trace,
            messages=rag_result.messages,
        )

    @staticmethod
    def _debug_user_context(db: Session, user_id: int | None) -> dict:
        if user_id is None:
            return {"growth_context": None, "trace": []}
        user = db.get(User, user_id)
        if user is None or user.is_deleted:
            return {
                "growth_context": None,
                "trace": [
                    {
                        "tool": "load_debug_user_context",
                        "order": 1,
                        "request": {"userId": user_id},
                        "status": "failed",
                        "durationMs": 0,
                        "response": None,
                        "error": "模拟学员不存在或已删除",
                    }
                ],
            }
        entitlement = EntitlementService.active_entitlement(
            db,
            user,
            monthly_topic_used=TopicSessionService.monthly_used_count(db, user.id),
        )
        growth_context = GrowthProfileService.prompt_context(db, user) if entitlement.enable_growth_profile else None
        return {
            "growth_context": growth_context,
            "trace": [
                {
                    "tool": "load_debug_user_context",
                    "order": 1,
                    "request": {"userId": user.id},
                    "status": "success",
                    "durationMs": 0,
                    "response": {
                        "userId": user.id,
                        "userName": user.name,
                        "entitlement": entitlement_dict(entitlement),
                        "growthProfileUsed": bool(growth_context),
                    },
                }
            ],
        }

    @staticmethod
    def overrides(payload: AgentDebugRequest) -> dict:
        return {
            "temperature": payload.temperature,
            "top_p": payload.topP,
            "top_k": payload.topK,
            "presence_penalty": payload.presencePenalty,
            "frequency_penalty": payload.frequencyPenalty,
            "max_token": payload.maxToken,
            "stream_enabled": payload.streamEnabled,
        }

    @classmethod
    async def stream(
        cls,
        payload: AgentDebugRequest,
        db: Session,
        current_admin: Admin,
    ) -> AsyncIterator[dict]:
        model = db.get(ModelConfig, payload.modelId)
        if model is None:
            yield {"type": "error", "message": "模型不存在"}
            return
        try:
            reasoning_visible = payload.reasoningVisible == 1
            yield {
                "type": "status",
                "stage": "retrieving",
                "message": "思考中",
                "reasoningVisible": reasoning_visible,
            }
            rag_result = await cls.build_result(db, payload)
            model_response = ModelStreamService.debug_stream_async(
                model,
                rag_result,
                cls.overrides(payload),
            )
            async for segment in ReasoningPolicyService.iter_segments(model_response.chunks):
                if segment.kind == "content":
                    yield {"type": "content", "content": segment.content}
                elif reasoning_visible:
                    yield {"type": "reasoning", "content": segment.content}
            OperationLogService.write(
                db,
                admin_id=current_admin.id,
                module="agent",
                action="debug_stream",
                target_id=model.id,
            )
            db.commit()
            yield {
                "type": "complete",
                "message": "Agent 调试完成",
                "modelName": model_response.model_name,
                "retrieveCount": len(rag_result.chunks),
                "knowledgeIds": rag_result.knowledge_ids,
                "retrievalTrace": rag_result.tool_trace or [],
                "retrievalLogId": rag_result.retrieval_log_id,
            }
        except asyncio.CancelledError:
            db.rollback()
            raise
        except Exception as exc:
            db.rollback()
            yield {"type": "error", "message": str(exc) or "Agent 调试失败"}
