from collections.abc import AsyncIterator

from app.agents.llm_adapter import LLMResponse, LLMStreamChunk
from app.agents.hint_agent import HintAgent
from app.agents.patient_agent import PatientAgent
from app.agents.report_agent import ReportAgent
from app.agents.scoring_agent import ScoringAgent
from app.models.source_case import CaseBase
from app.models.training import SessionOrder, SessionSubmission, TrainingSession


class MedicalConsultationOrchestrator:
    """主编排器：统一调度 Patient、Scoring、Report 等子 Agent。"""

    def __init__(self) -> None:
        self.patient_agent = PatientAgent()
        self.hint_agent = HintAgent()
        self.scoring_agent = ScoringAgent()
        self.report_agent = ReportAgent()

    async def patient_reply(self, session: TrainingSession, case: CaseBase, memory_messages: list[dict], message: str) -> LLMResponse:
        """问诊编排：调用 Patient Agent 生成 AI 病人回复。"""
        return await self.patient_agent.reply(case, memory_messages, message, session.mode)

    async def patient_stream_reply(
        self,
        session: TrainingSession,
        case: CaseBase,
        memory_messages: list[dict],
        message: str,
    ) -> AsyncIterator[LLMStreamChunk]:
        """流式问诊编排：调用 Patient Agent 并返回流式片段。"""
        async for chunk in self.patient_agent.stream_reply(case, memory_messages, message, session.mode):
            yield chunk

    async def evaluate(
        self,
        session: TrainingSession,
        case: CaseBase,
        memory_messages: list[dict],
        orders: list[SessionOrder],
        submission: SessionSubmission,
        rubric: object | None,
        guideline_refs: list[dict],
        scoring_rules: list | None = None,
    ) -> dict:
        """评价编排：调用 Scoring Agent 后交给 Report Agent 整理报告。"""
        scoring_result = await self.scoring_agent.score(
            session=session,
            case=case,
            memory_messages=memory_messages,
            orders=orders,
            submission=submission,
            rubric=rubric,
            guideline_refs=guideline_refs,
            scoring_rules=scoring_rules or [],
        )
        return self.report_agent.build_report(scoring_result)

    async def generate_hints(
        self,
        session: TrainingSession,
        case: CaseBase,
        memory_messages: list[dict],
        orders: list[SessionOrder],
        last_user_message: str | None = None,
    ) -> dict:
        """新手提示编排：基于当前会话上下文生成轻量训练提醒。"""
        return await self.hint_agent.generate(session, case, memory_messages, orders, last_user_message)