from pydantic import BaseModel, Field
from typing import Literal, List
from ..utils.prompt_template import PromptTemplate
from ..utils.extract import (
    default_extract_fields_to_string,
    default_extract_strings,
    modality_type_to_name,
    event_entity_to_name,
    understood_data_get_main_content,
)


class MemoryQueryPlan(BaseModel):
    query_type: Literal["none", "long_term_cached", "long_term_fresh"] = Field(
        default="none",
        description="""
查询类型，根据当前消息判断：
- none: 不查询任何记忆
- long_term_cached: 需要缓存的长期记忆的相关信息，复用之前查询过的长期记忆
- long_term_fresh: 需要从长期记忆中重新查询相关信息，要求最新信息或信息可能已变化
""",
    )

    query_strategy: Literal["semantic", "keyword"] = Field(
        default="semantic",
        description="""
查询策略，根据当前消息判断：
- semantic: 语义搜索，基于情境自动联想（**忽略query_triggers**）
- keyword: 关键词搜索，严格使用query_triggers精确匹配
""",
    )

    # TODO: 关键词不准确
    query_triggers: List[str] = Field(
        default_factory=list,
        description="数组，包含字符串。用于搜索记忆的关键词列表，可以添加更多同义词、相关的联想词。**仅当query_strategy为'keyword'时有效**",
    )

    time_range: List[int] = Field(
        default_factory=list, description="查询时间范围[起始天数, 结束天数]"
    )
    importance_score_filter: int = Field(
        default=0, description="重要性分数阈值(0-100)，只查询分数大于等于此值的记忆"
    )


"""event_type: Literal["user_command", "sensor_alert", "object_detected", "other"] = (
    Field(
        default="other",
        description="事件类型: user_command(用户指令)、sensor_alert(传感器预警)、object_detected(识别到特定物体或人)、other(其他)",
    )
)
- `event_type`: 字符串，枚举类型。必须是以下之一：
  - `"user_command"`: 用户指令
  - `"sensor_alert"`: 传感器预警  
  - `"object_detected"`: 识别到特定物体或人
  - `"other"`: 其他类型事件"""


class UnderstoodData(BaseModel):
    response_priority: Literal["low", "medium", "high", "critical"] = Field(
        default="low",
        description="根据安全性、紧急性判断响应紧急程度: low(低)、medium(中)、high(高)、critical(极高)",
    )
    main_content: str = Field(..., description="用一句话清晰概括当前信息的核心内容")
    current_situation: str | None = Field(
        ...,
        description="综合当前信息与历史上下文，生成对整体情境的连贯理解。形成完整的情境认知",
    )
    event_entity: str = Field(..., description="触发事件的主体")
    key_entities: List[str] = Field(
        default_factory=list, description="从信息中提取的重要名词或实体"
    )
    importance_score: int = Field(
        default=0, description="当前事件的重要程度分数(0-100)"
    )
    action_categorys: List[str] = Field(
        default_factory=list,
        description="根据当前情境可能需要执行的动作分类KEY。从可选动作分类中选择",
    )
    memory_query_plan: MemoryQueryPlan | None = Field(
        default=None, description="制定从长期记忆中查询相关信息的计划"
    )


understand_system_template = """下面是当前的信息，请根据你的角色将杂乱的多模态信息整理成一条结构化的“工作记忆”：

你可能会收到来自以下来源的原始信息：
- [asr]： 自动语音识别文本，可能包含错误或歧义。
- [text]： 文本信息，但可能有错别字。

【需要你理解的信息】
[{{understand_event_type}}]{{understand_event}}

【当前情境】
{{current_situation}}

【刚才的对话和事件】
{{recent_events}}

【可选动作分类】
{{action_categories}}"""


understand_template = f"""
<|im_start|>system
{understand_system_template}
<|im_end|>
<|im_start|>assistant
"""

understand_output_json_template = PromptTemplate(
    template="""请你严格按照指定的JSON格式输出。

# 输出要求
你必须输出一个JSON对象，包含以下字段：

## 一级字段说明：
- `response_priority`: 字符串，枚举类型。根据安全性、紧急性判断响应紧急程度，必须是：
  - `"low"`: 低优先级
  - `"medium"`: 中优先级
  - `"high"`: 高优先级
  - `"critical"`: 极高优先级

- `main_content`: 字符串。用**一句话**清晰概括当前信息的核心内容。

- `current_situation`: 字符串。综合当前信息与历史上下文，生成对整体情境的连贯理解，形成完整的情境认知。

- `event_entity`: 字符串。触发事件的主体（谁或什么触发了这个事件）。

- `key_entities`: 数组，包含字符串。从信息中提取的重要名词或实体。

- `importance_score`: 整数，范围0-100。当前事件的重要程度分数。

- `action_categorys`: 数组，根据当前情境可能需要执行的动作分类KEY。从可选动作分类中选择。

- `memory_query_plan`: 对象，包含记忆查询计划的详细信息。

## memory_query_plan 子对象字段说明：
- `query_type`: 字符串，枚举类型。必须是：
  - `"none"`: 不查询任何记忆
  - `"long_term_cached"`: 需要缓存的长期记忆的相关信息，复用之前查询过的长期记忆
  - `"long_term_fresh"`: 需要从长期记忆中重新查询相关信息，要求最新信息或信息可能已变化

- `query_strategy`: 字符串，枚举类型。必须是：
  - `"semantic"`: 语义搜索，基于情境自动联想（**忽略query_triggers**）
  - `"keyword"`: 关键词搜索，严格使用query_triggers精确匹配

- `query_triggers`: 数组，包含字符串。用于精确搜索记忆的关键词组合，可以添加更多同义词、相关的联想词。**仅当query_strategy为'keyword'时有效**

- `time_range`: 数组，包含两个整数。查询时间范围[起始天数, 结束天数]，如[0, 2]表示最近2天。

- `importance_score_filter`: 整数，范围0-100。重要性分数阈值，只查询分数大于等于此值的记忆。

# 输出示例
```json
{examples}""",
    variables={
        "examples": """{
  "response_priority": "medium", 
  "main_content": "用户要求打开客厅的灯光",
  "current_situation": "用户在晚上进入客厅后发出了开灯指令，表明需要照明",
  "event_entity": "用户",
  "key_entities": ["客厅", "灯光", "用户"],
  "importance_score": 30,
  "action_categorys: ["locomotion", "gestures"],
  "memory_query_plan": {
    "query_type": "long_term_fresh",
    "query_triggers": ["客厅", "灯光"],
    "time_range": [0, 2],
    "importance_score_filter": 0
  }
}"""
    },
)


def understand_task_format_inputs(inputs):
    return {
        "current_situation": inputs.get("current_situation", "未知"),
        "understand_event_type": inputs.get("understand_event", {}).get("type", "未知"),
        "understand_event": inputs.get("understand_event", {}).get("data", "无"),
        "recent_events": default_extract_fields_to_string(
            data_list=inputs.get("recent_events", []),
            field_configs=[
                {
                    "key": "modality_type",
                    "display": "类型",
                    "default": "未知",
                    "processor": modality_type_to_name,
                },
                {
                    "key": "understood_data",
                    "display": "来源",
                    "default": "未知",
                    "processor": event_entity_to_name,
                },
                {
                    "key": "understood_data",
                    "display": "内容",
                    "default": "未知",
                    "processor": understood_data_get_main_content,
                },
            ],
            list_name="无",
        ),
        "action_categories": default_extract_fields_to_string(
            data_list=inputs.get("action_categories", []),
            field_configs=[
                {
                    "key": "category_key",
                    "display": "KEY",
                    "default": "未知",
                },
                {
                    "key": "description",
                    "display": "描述",
                    "default": "未知",
                },
            ],
            list_name="无",
        ),
        "active_goals": default_extract_strings(
            inputs.get("active_goals", []), "description"
        ),
    }
