#!/usr/bin/env python3
import os
import sys
import io

from typing import List, Dict, Any, Optional, Tuple
from pathlib import Path
import time

# from langchain.agents import initialize_agent, AgentType
# from langchain.tools import BaseTool, Tool
# from langchain.llms import LlamaCpp
# from langchain.callbacks.manager import CallbackManager
# from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
# from langchain.prompts import PromptTemplate
# from langchain.chains import LLMChain
# import magic  # 用于识别文件类型（辅助文件搜索工具）

# from langchain.vectorstores import Chroma
# from langchain.embeddings import LlamaCppEmbeddings
# from langchain.text_splitter import RecursiveCharacterTextSplitter
# from langchain.document_loaders import (
#     TextLoader, UnstructuredFileLoader, PDFMinerLoader,
#     UnstructuredMarkdownLoader, Docx2txtLoader
# )
# from langchain.chains import RetrievalQA
# from langchain.memory import ConversationBufferWindowMemory
# from langchain.prompts import PromptTemplate
# from langchain.schema import HumanMessage, AIMessage

import logging

# import argparse
# import base64

# from fastapi import FastAPI, HTTPException, APIRouter, Path
# from fastapi.responses import JSONResponse

# from fastapi_utils.cbv import cbv
# from fastapi_utils.inferring_router import InferringRouter

from pydantic import BaseModel
from PIL import Image
import numpy as np
import requests
import json

# 获取当前脚本的路径（Path 对象）
current_script = Path(__file__).resolve()
# 向上两级找到项目根目录（根据实际结构调整层级）
project_root = current_script.parents[1]  # parents[0] 是当前目录，parents[1] 是父目录，以此类推
print("project_root: ",project_root)

#添加到搜索路径
sys.path.append(str(project_root))

#查看系统环境
print("sys path:",sys.path)

# from tools.core.common_functions import read_json_file
from tools.data_process.image_process.image_process import PIL_image_to_base64
from models.model_definition import (TextInferenceRequest,MultimodalInferenceRequest,InferenceRequest,InferenceResponse,HealthResponse)

# 配置日志记录器
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


#定义模型Client端
class Models_Client:
    def __init__(self, server_url: str = "http://localhost:8000"):
        """
        初始化多模态客户端
        
        参数:
            server_url: 服务器地址，格式为 "http://host:port"
        """
        self.server_url = server_url
        self.text_endpoint = f"{server_url}/text/inference"
        self.multimodal_endpoint = f"{server_url}/multimodal/inference"
        self.model_inference_endpoint = f"{server_url}/model/inference"
        self.health_endpoint = f"{server_url}/model/health"
    
    def wait_for_server(self, timeout: int = 300, check_interval: int = 5) -> bool:
        """
        等待服务器加载完成
        
        参数:
            timeout: 最大等待时间（秒）
            check_interval: 检查间隔（秒）
            
        返回:
            服务器是否准备就绪
        """
        start_time = time.time()
        
        while time.time() - start_time < timeout:
            try:
                health = self.check_health()
                if health.get("status") == "healthy":
                    print("服务器已准备就绪")
                    return True
                print(f"服务器正在加载中，已等待 {int(time.time() - start_time)} 秒...")
                time.sleep(check_interval)
            except Exception as e:
                print(f"检查服务器状态失败: {str(e)}")
                time.sleep(check_interval)
        
        print(f"超时: 服务器在 {timeout} 秒内未准备就绪")
        return False
    
    def check_health(self) -> Dict[str, Any]:
        """检查服务器健康状态"""
        try:
            response = requests.get(self.health_endpoint)
            response.raise_for_status()
            return response.json()
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def text_inference(self, request:TextInferenceRequest) -> Dict[str, Any]:
        """
        发送纯文本推理请求
        
        参数:
            prompt: 提示文本
            max_tokens: 最大生成token数
            temperature: 生成温度，控制随机性
            top_p: 核采样参数
            stop: 停止序列
            
        返回:
            包含推理结果的字典
        """
        try:
            # 将 Pydantic 模型转换为 JSON 字符串
            if hasattr(request, 'model_dump_json'):
                # Pydantic v2
                data = request.model_dump_json()
            elif hasattr(request, 'json'):
                # Pydantic v1
                data = request.json()
            else:
                # 降级方案：转换为字典再序列化
                data = json.dumps(request.dict() if hasattr(request, 'dict') else request.model_dump())
            
            response = requests.post(
                self.text_endpoint,
                headers={"Content-Type": "application/json"},
                data=data
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            return {"error": str(e)}
    
    def text_inference2(self, 
                      prompt: str, 
                      max_tokens: int = 200, 
                      temperature: float = 0.7, 
                      top_p: float = 0.95,
                      system_prompt: str="",
                      stop: Optional[List[str]] = None) -> Dict[str, Any]:
        """
        发送纯文本推理请求
        
        参数:
            prompt: 提示文本
            max_tokens: 最大生成token数
            temperature: 生成温度，控制随机性
            top_p: 核采样参数
            stop: 停止序列
            
        返回:
            包含推理结果的字典
        """
        payload = {
            "user_prompt": prompt,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "system_prompt":system_prompt,
            "stop": stop
        }
        
        try:
            response = requests.post(
                self.text_endpoint,
                headers={"Content-Type": "application/json"},
                data=json.dumps(payload)
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            return {"error": str(e)}
        
    def multimodal_inference(self, request:MultimodalInferenceRequest) -> Dict[str, Any]:
        """
        发送多模态推理请求（文本+图像）
        
        参数:
            prompt: 提示文本
            image_paths: 图像文件路径列表
            images: jpeg base64对象列表
            max_tokens: 最大生成token数
            temperature: 生成温度，控制随机性
            top_p: 核采样参数
            stop: 停止序列
            
        返回:
            包含推理结果的字典
        """
        # # 处理图像
        # image_data = []

        # if images:
        #     # for img in images:
        #     try:
        #         # image_data.append(PIL_image_to_base64(image=img))
        #         image_data= images
        #     except Exception as e:
        #         return {"error": f"处理图像base64对象失败: {str(e)}"}
        
        # elif image_paths:
        #     for path in image_paths:
        #         try:
        #             image_data.append(PIL_image_to_base64(image_path=path))
        #         except Exception as e:
        #             return {"error": f"处理图像 {path} 失败: {str(e)}"}
        
        
        
        # if not image_data:
        #     return {"error": "未提供有效的图像数据"}
        
        # payload = {
        #     "user_prompt": prompt,
        #     "image_data": image_data,
        #     "max_tokens": max_tokens,
        #     "temperature": temperature,
        #     "top_p": top_p,
        #     "system_prompt": system_prompt,
        #     "stop": stop
        # }
        
        try:
            # 将 Pydantic 模型转换为 JSON 字符串
            if hasattr(request, 'model_dump_json'):
                # Pydantic v2
                data = request.model_dump_json()
            elif hasattr(request, 'json'):
                # Pydantic v1
                data = request.json()
            else:
                # 降级方案：转换为字典再序列化
                data = json.dumps(request.dict() if hasattr(request, 'dict') else request.model_dump())
            
            response = requests.post(
                url=self.multimodal_endpoint,
                headers={"Content-Type": "application/json"},
                data=data
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            return {"error": str(e)}
        
    def multimodal_inference2(self, 
                            prompt: str, 
                            image_paths: Optional[List[str]] = None, 
                            images: Optional[List[str]] = None,
                            max_tokens: int = 300, 
                            temperature: float = 0.7, 
                            top_p: float = 0.95,
                            system_prompt: str="",
                            stop: Optional[List[str]] = None) -> Dict[str, Any]:
        """
        发送多模态推理请求（文本+图像）
        
        参数:
            prompt: 提示文本
            image_paths: 图像文件路径列表
            images: jpeg base64对象列表
            max_tokens: 最大生成token数
            temperature: 生成温度，控制随机性
            top_p: 核采样参数
            stop: 停止序列
            
        返回:
            包含推理结果的字典
        """
        # 处理图像
        image_data = []

        if images:
            # for img in images:
            try:
                # image_data.append(PIL_image_to_base64(image=img))
                image_data= images
            except Exception as e:
                return {"error": f"处理图像base64对象失败: {str(e)}"}
        
        elif image_paths:
            for path in image_paths:
                try:
                    image_data.append(PIL_image_to_base64(image_path=path))
                except Exception as e:
                    return {"error": f"处理图像 {path} 失败: {str(e)}"}
        
        
        
        if not image_data:
            return {"error": "未提供有效的图像数据"}
        
        payload = {
            "user_prompt": prompt,
            "image_data": image_data,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "system_prompt": system_prompt,
            "stop": stop
        }
        
        try:
            response = requests.post(
                url=self.multimodal_endpoint,
                headers={"Content-Type": "application/json"},
                data=json.dumps(payload)
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            return {"error": str(e)}
        
    
        
    def model_inference(self, 
                            prompt: str, 
                            image_paths: Optional[List[str]] = None, 
                            images: Optional[List[str]] = None,
                            max_tokens: int = 300, 
                            temperature: float = 0.7, 
                            top_p: float = 0.95,
                            system_prompt: str="",
                            stop: Optional[List[str]] = None) -> Dict[str, Any]:
        """
        发送多模态推理请求（文本+图像）
        
        参数:
            prompt: 提示文本
            image_paths: 图像文件路径列表
            images: jpeg base64对象列表
            max_tokens: 最大生成token数
            temperature: 生成温度，控制随机性
            top_p: 核采样参数
            stop: 停止序列
            
        返回:
            包含推理结果的字典
        """
        # 处理图像
        image_data = []

        if images:
            # for img in images:
            try:
                # image_data.append(PIL_image_to_base64(image=img))
                image_data= images
            except Exception as e:
                return {"error": f"处理图像base64对象失败: {str(e)}"}
        
        elif image_paths:
            for path in image_paths:
                try:
                    image_data.append(PIL_image_to_base64(image_path=path))
                except Exception as e:
                    return {"error": f"处理图像 {path} 失败: {str(e)}"}
        
        payload = {
            "prompt": prompt,
            "image_data": None,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "system_prompt": system_prompt,
            "stop": stop
        }
        
        #如果有image_data，就加入
        if image_data:
            payload["image_data"] = image_data

        logger.info(f"payload: {payload}")
        
        try:
            response = requests.post(
                self.model_inference_endpoint,
                headers={"Content-Type": "application/json"},
                data=json.dumps(payload)
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            return {"error": str(e)}
        
def model_client_test(inference_case=1):
    
    config_json_file="/home/ubuntu/Workspace/Projects/VLM_VLA/UAV_AI/src/Model/AI_Agent/scripts/ai_agent/models/model_config.json"
    from tools.core.common_functions import read_json_file
    model_client_config = read_json_file(config_json_file)["models_client"][0]
    model_client = Models_Client(model_client_config["server_url"])
    
    if 2==inference_case:
        test_prompt = "检测图中的蓝色气球，并按照固定的结构(<|object_ref_start|>蓝色气球<|object_ref_end|><|box_start|>(xmin,ymin),(xamx,ymax)<|box_end|>)给出2D框坐标"
        test_image_path = "/home/ubuntu/Workspace/Projects/VLM_VLA/UAV_AI/datas/test1/balloon2.jpg"
        reponse =  model_client.multimodal_inference(test_prompt,image_paths=[test_image_path])
    elif 1==inference_case:
        test_prompt = "请用中文对图片进行场景描述！"
        test_image_path = "/home/ubuntu/Workspace/Projects/VLM_VLA/UAV_AI/datas/test1/fj1.jpg"
        reponse =  model_client.multimodal_inference(test_prompt,image_paths=[test_image_path])
    elif 0==inference_case:
        test_prompt = "请你说一下《西游记》中孙悟空的性格特点！"
        reponse =  model_client.text_inference(test_prompt)

    print("reponse: ",reponse)



def main():
    model_client_test(inference_case=2)



if __name__ == "__main__":
    main()


    
