"""pydantic-ai Agent 应用模块

使用 2026 pydantic-ai 最佳实践：
- deps_type 依赖注入
- @agent.instructions 动态指令
- 结构化输出 (Pydantic models)
"""

import asyncio
import os
from dataclasses import dataclass
from typing import Protocol

from dotenv import load_dotenv
from pydantic_ai import Agent, RunContext

from src.features import StudentFeatures, StudyGuidance
from src.infer import explain_prediction, predict_pass_prob

load_dotenv()


# --- 1. 定义依赖协议和数据类 ---


class MLModelProtocol(Protocol):
    """ML 模型接口协议"""

    def predict(self, features: StudentFeatures) -> float:
        """预测通过概率"""
        ...

    def explain(self) -> str:
        """获取模型解释"""
        ...


@dataclass
class AgentDeps:
    """Agent 依赖项

    封装 ML 模型和学生特征，通过依赖注入传递给 Agent。
    """

    student: StudentFeatures
    model_path: str = "models/model.pkl"


# --- 2. 定义 Agent ---


study_advisor = Agent(
    "deepseek:deepseek-chat",
    deps_type=AgentDeps,
    output_type=StudyGuidance,
    instructions=(
        "你是一个严谨的学业数据分析师。你的任务是根据学生的具体情况预测其考试通过率，并给出建议。\n"
        "【重要规则】\n"
        "1. 必须先调用 `predict_pass_probability` 获取概率。\n"
        "2. 必须调用 `get_model_explanation` 获取模型认为最重要的特征，并在 `key_factors` 中引用这些特征。\n"
        "3. 你的建议必须针对那些最重要的特征（例如，如果模型说睡眠很重要，就给睡眠建议）。\n"
        "4. 严禁凭空编造数值。所有数据必须来自工具返回。\n"
        "5. `rationale` 必须引用 `key_factors` 中的具体因素。"
    ),
)


@study_advisor.instructions
async def add_student_context(ctx: RunContext[AgentDeps]) -> str:
    """动态添加学生信息到系统提示"""
    s = ctx.deps.student
    return (
        f"当前学生信息：\n"
        f"- 每周学习时长: {s.study_hours} 小时\n"
        f"- 每晚睡眠时长: {s.sleep_hours} 小时\n"
        f"- 出勤率: {s.attendance_rate:.0%}\n"
        f"- 压力等级: {s.stress_level}/5\n"
        f"- 学习方式: {s.study_type}"
    )


# --- 3. 注册工具 ---


@study_advisor.tool
async def predict_pass_probability(ctx: RunContext[AgentDeps]) -> float:
    """调用 ML 模型预测学生通过概率

    Returns:
        float: 预测通过率 (0-1)
    """
    s = ctx.deps.student
    return predict_pass_prob(
        study_hours=s.study_hours,
        sleep_hours=s.sleep_hours,
        attendance_rate=s.attendance_rate,
        stress_level=s.stress_level,
        study_type=s.study_type,
    )


@study_advisor.tool
async def get_model_explanation(ctx: RunContext[AgentDeps]) -> str:
    """获取 ML 模型的特征重要性解释

    Returns:
        str: 特征重要性排名说明
    """
    return explain_prediction()


# --- 4. 咨询师 Agent (多轮对话) ---


counselor_agent = Agent(
    "deepseek:deepseek-chat",
    deps_type=AgentDeps,
    instructions=(
        "你是一位富有同理心且专业的大学心理咨询师。\n"
        "你的目标是倾听学生的学业压力和生活烦恼，提供情感支持。\n"
        "【交互风格】\n"
        "1. 同理心：首先通过复述或确认学生的感受来表达理解。\n"
        "2. 引导性：不要急于给出解决方案，先通过提问了解更多背景。\n"
        "3. 数据驱动（可选）：如果学生询问具体通过率，请调用工具。\n"
        "4. 语气：温暖、支持、专业，像朋友一样交谈。"
    ),
)


@counselor_agent.tool
async def predict_student_pass(ctx: RunContext[AgentDeps]) -> float:
    """获取学生通过率预测（用于咨询过程提供客观数据）"""
    s = ctx.deps.student
    return predict_pass_prob(
        study_hours=s.study_hours,
        sleep_hours=s.sleep_hours,
        attendance_rate=s.attendance_rate,
        stress_level=s.stress_level,
        study_type=s.study_type,
    )


@counselor_agent.tool
async def explain_factors(ctx: RunContext[AgentDeps]) -> str:
    """获取模型特征重要性解释"""
    return explain_prediction()


# --- 5. 运行示例 ---


async def main():
    """运行 Agent 示例"""
    if not os.getenv("DEEPSEEK_API_KEY"):
        print("❌ 错误: 未设置 DEEPSEEK_API_KEY")
        print("请在 .env 文件中设置密钥，或 export DEEPSEEK_API_KEY='...'")
        return

    # 构建学生特征
    student = StudentFeatures(
        study_hours=12,
        sleep_hours=4,
        attendance_rate=0.9,
        stress_level=4,
        study_type="Self",
    )

    # 创建依赖
    deps = AgentDeps(student=student)

    # 用户查询
    query = (
        "我最近压力很大 (等级4)，每天只睡 4 小时，不过我每周自学(Self) 12 小时，"
        "出勤率大概 90%。请帮我分析一下我会挂科吗？基于模型告诉我怎么做最有效。"
    )

    print(f"用户: {query}\n")
    print("Agent 正在思考并调用模型工具...\n")

    try:
        result = await study_advisor.run(query, deps=deps)

        print("--- 结构化分析报告 ---")
        print(result.output.model_dump_json(indent=2))

    except Exception as e:
        print(f"❌ 运行失败: {e}")


if __name__ == "__main__":
    asyncio.run(main())