import re

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from loguru import logger

from config.config import load_config

config = load_config("config.toml")


def generate_comment(name, age_group, traits):
    """
    生成评语
    :param name: 学生姓名
    :param age_group: 所在班级
    :param traits: 表现特征
    :return: 评语
    """

    ai_config = config["ai"]
    llm = ChatOpenAI(
        base_url=ai_config["api_url"],
        api_key=ai_config["api_key"],
        model=ai_config["model"],
        temperature=0.7,
    )
    # 2. 构建 Prompt Template
    prompt = ChatPromptTemplate.from_messages([
        ("system", ai_config["prompt"]),
        ("human", "学生姓名：{name}\n所在班级：{age_group}\n表现特征：{traits}\n\n请开始撰写评语：")
    ])

    # 3. 组装链 (Prompt -> Model -> OutputParser)
    chain = prompt | llm | StrOutputParser()

    # 4. 执行
    try:
        comment = chain.invoke({
            "name": name,
            "age_group": age_group,
            "traits": traits
        })
        cleaned_text = re.sub(r'\s+', '', comment)
        logger.success(f"学生：{name} =>生成评语成功: {cleaned_text}")
        return cleaned_text
    except Exception as e:
        print(f"生成评语失败: {e}")
        return "生成失败，请检查网络或Key。"