# -*- coding: utf-8 -*-

from typing import List, Callable
from enum import Enum

class ContextSource(Enum):
    PROMPT_GENERATED = "prompt_generated"
    EXTERNAL_KNOWLEDGE = "external_knowledge"
    DYNAMIC_ASSEMBLED = "dynamic_assembled"

class ContextRetrieval:
    def __init__(self, context_manager):
        self.context_manager = context_manager
        # 注册外部知识检索器（示例接口）
        self.external_retrievers = {
            "default": self._default_external_retriever
        }

    def generate_prompt_based_context(self, task_type: str, reasoning_framework: str = "basic") -> str:
        """基于提示的生成：创建任务优化的指令框架
        Args:
            task_type: 任务类型，如"text_classification"、"nl2sql"、"summarization"
            reasoning_framework: 推理框架类型，如"basic"、"chain_of_thought"、"tree_of_thought"
        Returns:
            生成的提示上下文字符串
        """
        # 基础提示模板库
        prompt_templates = {
            "text_classification": """你是一个文本分类专家。基于以下上下文信息，将文本分类到预定义类别中：
{context}
文本：{input}
要求：输出最可能的类别名称，无需解释。""",
            "nl2sql": """你是一个SQL生成专家。基于以下数据库模式信息，将自然语言转换为SQL查询：
{context}
用户问题：{input}
要求：输出标准SQL，不包含解释和多余文本。""",
            "summarization": """你是一个文本摘要专家。基于以下上下文信息，生成简洁准确的摘要：
{context}
文本：{input}
要求：摘要长度不超过{max_length}字，保留关键信息。"""
        }

        # 推理框架增强
        reasoning_enhancements = {
            "basic": "",
            "chain_of_thought": "思考过程：让我逐步分析...",
            "tree_of_thought": "可能的思考方向：\n1. ...\n2. ...\n结论：..."
        }

        # 获取基础模板
        base_template = prompt_templates.get(task_type, prompt_templates["text_classification"])
        # 添加推理框架
        enhanced_template = f"{base_template}\n{reasoning_enhancements.get(reasoning_framework, '')}"

        # 存储生成的提示模板到上下文
        self.context_manager.set(f"prompt_template_{task_type}", enhanced_template)
        return enhanced_template

    def retrieve_external_knowledge(self, query: str, retriever_name: str = "default", **kwargs) -> str:
        """外部知识检索：访问动态信息源
        Args:
            query: 检索查询词
            retriever_name: 检索器名称
            **kwargs: 检索参数（如top_k、filter条件等）
        Returns:
            检索到的知识文本
        """
        if retriever_name not in self.external_retrievers:
            raise ValueError(f"Retriever {retriever_name} not registered")

        # 调用注册的检索器
        knowledge = self.external_retrievers[retriever_name](query, **kwargs)
        # 存储检索结果到上下文
        self.context_manager.set(f"external_knowledge_{query[:30]}", knowledge)
        return knowledge

    def assemble_dynamic_context(self, context_ids: List[str], assembly_strategy: str = "sequential") -> str:
        """动态上下文组装：将多个上下文组件整合成优化的输入
        Args:
            context_ids: 上下文组件ID列表
            assembly_strategy: 组装策略，如"sequential"、"priority_based"、"summary_based"
        Returns:
            组装后的完整上下文
        """
        # 获取所有上下文组件
        context_components = []
        for cid in context_ids:
            component = self.context_manager.get(cid)
            if component:
                context_components.append((cid, component))

        # 应用组装策略
        if assembly_strategy == "sequential":
            assembled = "\n\n".join([f"[{cid}]: {comp}" for cid, comp in context_components])
        elif assembly_strategy == "priority_based":
            # 按ID中包含的优先级关键词排序
            context_components.sort(key=lambda x: "priority" in x[0], reverse=True)
            assembled = "\n\n".join([f"[{cid}]: {comp}" for cid, comp in context_components])
        elif assembly_strategy == "summary_based":
            # 生成组件摘要（简化实现）
            summaries = [f"[{cid}摘要]: {comp[:100]}..." for cid, comp in context_components]
            assembled = "\n\n".join(summaries + ["\n完整内容:\n" + "\n\n".join([comp for _, comp in context_components])])
        else:
            assembled = "\n\n".join([comp for _, comp in context_components])

        # 存储组装结果
        self.context_manager.set("assembled_context", assembled)
        return assembled

    def _default_external_retriever(self, query: str, top_k: int = 3) -> str:
        """默认外部知识检索器（示例实现）
        实际应用中可替换为调用搜索引擎、数据库或API
        """
        # 模拟外部检索结果
        mock_knowledge = [
            f"[检索结果1] 关于'{query}'的信息：这是模拟的外部知识，实际应用中应替换为真实API调用。",
            f"[检索结果2] 相关数据：{len(query)}个字符，{query.count(' ')+1}个词。",
            f"[检索结果3] 时间戳：{self.context_manager.get('current_timestamp', '未知')}"
        ]
        return "\n\n".join(mock_knowledge[:top_k])

    def register_retriever(self, name: str, retriever_func: Callable) -> None:
        """注册自定义外部知识检索器"""
        self.external_retrievers[name] = retriever_func


def test():
    print()


if __name__ == '__main__':
    None
