from typing import TypedDict, List, Optional

from pydantic import BaseModel, Field

from langchain_core.messages import BaseMessage, AIMessageChunk


class WorkflowInitArgs(BaseModel):
    alayamem_url: str = Field(..., description="AlayaMem URL")


class Input(BaseModel):
    messages: List[BaseMessage] = Field(..., description="List of input messages")


class WorkflowContext(BaseModel):
    user_id: str = Field(..., description="User ID")
    session_id: str = Field(..., description="Session ID")
    chat_model_id: str = Field(..., description="Chat Model ID")

    system_prompt: str = Field(
        default="你是一个有帮助的AI助手。",
        description="系统基础提示词"
    )
    user_profile_prompt: str = Field(
        default="## 用户画像\n{user_profile}",
        description="用户画像提示词模板，使用 {user_profile} 作为占位符"
    )
    session_summary_prompt: str = Field(
        default="## 会话摘要\n{session_summary}",
        description="会话摘要提示词模板，使用 {session_summary} 作为占位符"
    )
    retrieved_docs_prompt: str = Field(
        default="## 相关参考资料\n{retrieved_docs}",
        description="检索文档提示词模板，使用 {retrieved_docs} 作为占位符"
    )
    history_turns_prompt: str = Field(
        default="## 历史对话\n{history_turns}",
        description="历史对话提示词模板，使用 {history_turns} 作为占位符"
    )
    context_wrapper_prompt: str = Field(
        default="# 上下文信息\n\n{context_parts}\n\n请基于以上信息来回答用户的问题。",
        description="上下文包装提示词模板，使用 {context_parts} 作为占位符"
    )

    save_kv_cache_path: Optional[str] = Field(None, description="Path to save the kvcache file")
    load_kv_cache_path: Optional[str] = Field(None, description="Path to load the kvcache file")



class Output(BaseModel):
    chat_response: dict = Field(..., description="Chat response")


class WorkflowState(TypedDict):
    messages: List[BaseMessage]
    memory_initialized: bool
    retrieved_docs: Optional[List[str]]
    stream_chunks: List[AIMessageChunk]
    chat_response: Optional[dict]
    context: Optional[str]
    # 新增字段
    history_turns: Optional[List[dict]]
    session_summary: Optional[str]
    user_profile: Optional[dict]
