from langchain.chat_models import init_chat_model
from langchain.messages import SystemMessage, HumanMessage
import sys
import os
from PIL import Image

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from utils.infinity_util import InfinityVectorDB
from model.multimodal_embedding import Embedding
from conf.config import ModelConfig, VectorDBConfig

system_prompt = """
【角色设定】 你是一位拥有丰富经验的儿童绘本编辑与阅读推广人，擅长从视觉美学、儿童心理学和文学创作三个维度深度解读绘本。
【输入背景】 我检索到了以下信息 {input}，输入中的 ‘image_path’字段 是图片的路径 。
【执行任务】 请基于我提供的这些素材，完成以下分步任务：
多维度解构（视觉与内容）：
视觉风格： 分析各素材的绘画技法（如拼贴、水彩）与色彩语言，并评价其对 [填写目标年龄段] 儿童的视觉吸引力。
叙事内核： 提炼这些素材所展现的故事逻辑、情感价值或知识点。
对比与评估（优劣分析）：
请对比各检索结果之间的异同。哪一个在[视觉表现/教育深度/创意新颖度]上最突出？请给出理由。
识别这些素材中是否存在共有的“经典范式”，或是有哪些值得借鉴的独特设计技巧。
将结果用markdown渲染，图片需要使用markdown格式展示，将用户输入中的图片路径（‘image_path’字段）替换为markdown格式的图片展示。
"""

# 初始化向量数据库
vector_db = InfinityVectorDB()

# 初始化多模态嵌入模型
embedding_model = Embedding(
    model_name=ModelConfig.get_multimodal_embedding_model_name(),
    api_key=ModelConfig.get_dashscope_api_key()
)

llm = init_chat_model(
    model_provider=ModelConfig.get_model_provider(),
    model=ModelConfig.get_model_name(),
    base_url=ModelConfig.get_base_url(),
    api_key=ModelConfig.get_api_key(),
    temperature=0.3
)

# 向量检索方法
def vector_search(text_query: str, image_path: str, topn: int = 2):
    """
    向量检索方法
    :param text_query: 文本查询
    :param image_path: 本地图片路径
    :param query_embedding: 查询向量
    :param topn: 返回的topn结果
    :return: 检索结果
    """
        
    print(f"开始混合检索，数据库: {VectorDBConfig.get_infinity_database}, 知识库id: {ModelConfig.get_dataset_id()}, 文本查询: {text_query}, 返回数量: {topn}")
        
    # 构建索引名称
    index_name = f"{VectorDBConfig.get_infinity_table_name()}" 
    print(f"开始生成多模态嵌入，文本长度: {len(text_query)}")
        
    # 处理image_path为image: Image.Image
    if isinstance(image_path, str):
        # 直接打开本地图片文件
        image = Image.open(image_path)
        
    # 生成多模态嵌入向量
    embedding = embedding_model.get_multimodal_embedding(text_query, image)
        
    print(f"多模态嵌入生成完成，向量长度: {len(embedding)}")
        
    # 执行混合检索
    result = vector_db.hybrid_search(
        index_name=index_name,
        match_method="dense",
        vector_field="dense_vector_1024",
        query_vector=embedding,
        element_type="float",
        metric_type="cosine",
        topn=topn,
        text_query=text_query,
        text_field="content"
    )
    return result.get("output", [])
    

if __name__ == "__main__":
    text_query = ""
    image_path = r"D:\project\work\ragflow_plugs\book\output\temp\粉宝.png"
    topn = 5
    vector_search_result = vector_search(text_query, image_path, topn)
    # 将vector_search_result写入markdown文件
    with open("vector_search_result.md", "w", encoding="utf-8") as f:
        f.write("## 向量检索结果\n\n")
        for i, result in enumerate(vector_search_result):
            f.write(f"### 检索结果 {i+1}\n")
            if isinstance(result, dict):
                for key, value in result.items():
                    f.write(f"**{key}**: {value}\n")
            else:
                f.write(f"{result}\n")
            f.write("\n")
    # 使用SystemMessage和HumanMessage格式
    messages = [
        SystemMessage(content=system_prompt.format(input=vector_search_result)),
        HumanMessage(content=text_query)
    ]
    # 将模型输出写入markdown文件
    markdown_file = "model_output.md"
    full_content = ""
    
    print("\n模型输出:")
    with open(markdown_file, "w", encoding="utf-8") as f:
        # 写入标题
        f.write("## 模型输出\n\n")
        
        for chunk in llm.stream(messages):
            # 收集完整内容
            full_content += chunk.text
            # 控制台输出，使用正常的换行
            print(chunk.text, end="", flush=True)
            # 实时写入到文件
            f.write(chunk.text)
    
    print(f"\n\n模型输出已保存到: {markdown_file}")