import pandas as pd
from model_handler import ModelHandler
from config import LING_MINI_2_0
from .agent_common_utils import format_df_to_string

def fetch_prompt_suggestions_agent(editor_content: str, style: str, kb_df: pd.DataFrame, short_outline_df: pd.DataFrame, long_outline_df: pd.DataFrame):
    """
    Agent for fetching short prompt suggestions using LING_MINI_2_0.
    """
    print("\n[Agent][fetch_prompt_suggestions_agent] === 推理类型：续写提示推荐 ===")
    try:
        # 1. Format context
        style_context = f"### 整体章程\n{style}\n\n"
        kb_context = format_df_to_string(kb_df, "知识库")
        short_outline_context = format_df_to_string(short_outline_df, "当前章节大纲")

        # 2. Build System Prompt
        system_prompt = (
            "你是一个辅助写作的创意助手。请根据提供的故事背景和知识库，结合“互动”、“冲突”、“发展”、“对话”等动作，生成3个简短的续写提示短语。\n"
            "要求：\n"
            "1. 短语简洁明了，例如：“和Alpha争吵”、“探索废弃的地铁站”、“回忆起旧照片的往事”。\n"
            "2. 尽量使用知识库中的专有名词。\n"
            "3. 请严格遵守以下格式：输出3个短语，用 `|` 分隔。不要包含其他内容。\n"
            "例如：和Alpha争吵|探索废弃的地铁站|回忆起旧照片的往事"
        )
        
        # 3. Build User Prompt
        full_context = style_context + kb_context + short_outline_context
        user_prompt = (
            f"### 背景设定\n{full_context}\n"
            f"### 当前已写内容 (末尾部分)\n{editor_content[-500:]}\n\n" # Only need a little context
            f"### 任务\n生成3个续写提示。"
        )

        # 4. Call LLM
        model_handler = ModelHandler()
        response_generator = model_handler.generate_code(
            system_prompt=system_prompt,
            user_prompt=user_prompt,
            model_choice=LING_MINI_2_0 
        )
        
        full_response = "".join(chunk for chunk in response_generator)
        print("【收到的建议】", full_response)
        
        suggestions = full_response.split("|")
        # Ensure 3 suggestions
        suggestions += ["继续推进剧情"] * (3 - len(suggestions))
        
        return suggestions[0].strip(), suggestions[1].strip(), suggestions[2].strip()

    except Exception as e:
        print(f"[Agent] Error fetching prompt suggestions: {e}")
        return "生成失败", "生成失败", "生成失败"