import requests
import json
# import openai
import os
from typing import Dict, Any, Optional
import time

class Client:
    def __init__(self, 
                 model_id: str, 
                 server_url: str, 
                 top_p: float = 0.9, 
                 temperature: float = 0.7, 
                 max_tokens: int = 1024,
                 api_key: str = None):
        """
        初始化客户端
        
        参数:
            model_id: 模型ID
            server_url: API服务器地址
            top_p: 采样概率阈值，控制多样性
            temperature: 温度参数，控制随机性
            max_tokens: 最大生成token数
        """
        self.model_id = model_id
        self.server_url = f"{server_url}/chat/completions"
        # self.server_url = server_url
        self.top_p = top_p
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.api_key = api_key

    def interact(self, prompt: str, use_system_prompt: bool = False, use_thinking: bool = False) -> str:
        """
        与大模型进行交互
        
        参数:
            prompt: 用户输入的提示文本
            use_system_prompt: 是否添加系统提示，默认为False
            
        返回:
            str: 模型生成的响应
        """
        # 构建消息列表
        messages = []
        
        # 添加系统提示（如果启用）
        if use_system_prompt:
            messages.append({
                "role": "system", 
                "content": "You are an expert automated testing assistant."
            })
        
        # 添加用户提示
        messages.append({"role": "user", "content": prompt})
        
        # 最大重试次数
        max_retries = 10000
        retry_count = 0
        while retry_count <= max_retries:
            # 构建请求数据
            request_data = {
                "model": self.model_id,
                "messages": messages,  # 使用构建的消息列表
                "temperature": self.temperature,
                "top_p": self.top_p,
                "max_tokens": self.max_tokens,
            }
            # 构建请求头
            headers = {
                "Content-Type": "application/json",
                "Authorization": "Bearer dumpy"
            }
            if self.api_key:
                headers["Authorization"] = f"Bearer {self.api_key}"

            try:
                # 发送请求
                response = requests.post(
                    self.server_url,
                    headers=headers,
                    data=json.dumps(request_data),
                    timeout=3600
                )
                
                # 检查响应状态
                response.raise_for_status()
                
                # 解析响应
                result = response.json()
                # print(result)
                reasoning_content = result.get("choices", [{}])[0].get("message", {}).get("reasoning_content", "抱歉，模型没有返回有效思考内容")
                content = result.get("choices", [{}])[0].get("message", {}).get("content", "抱歉，模型没有返回有效内容")
                # 检查是否为无效内容
                if content != "抱歉，模型没有返回有效内容" and len(content) > 0:
                    if use_thinking:
                        return [content, reasoning_content]
                    else:
                        return content
                else:
                    if retry_count == 1:
                        # print(f"prompt is {prompt}")
                        print(result.get("choices", [{}])[0].get("message", {}).get("reasoning_content", "抱歉，模型没有返回有效内容"))
                        print(f"content is {content}")
                    retry_count += 1
                    print(f"第 {retry_count} 次重试，模型返回无效内容")
                    if retry_count <= max_retries:
                        time.sleep(10)  # 等待10秒后重试
            except requests.exceptions.RequestException as e:
                print(f"请求错误: {e}")
                time.sleep(20)
                retry_count += 1
            except json.JSONDecodeError:
                print("无法解析模型响应")
                time.sleep(20)
                retry_count += 1
            # except openai.RateLimitError:
            #     print("Rate limit exceeded. Waiting...")
            #     time.sleep(20)
            #     retry_count += 1
            # except openai.APIConnectionError:
            #     print("API connection error. Waiting...")
            #     time.sleep(20)
            #     retry_count += 1
            # except openai.APIError as e:
            #     print(e)
            #     time.sleep(20)
            #     retry_count += 1
            except Exception as e:
                print("Unknown error. Waiting...")
                print(e)
                time.sleep(20)
                retry_count += 1
        return "抱歉，模型没有返回有效内容"

    def test_interact(self, use_system_prompt: bool = False) -> None:
        """
        在终端中进行交互式测试
        """
        print(f"开始与模型 {self.model_id} 进行交互测试")
        print("输入 'exit' 或 'quit' 结束对话")
        
        while True:
            user_input = input("\n你: ")
            
            if user_input.lower() in {'exit', 'quit'}:
                print("对话已结束")
                break
                
            if not user_input.strip():
                continue
                
            # 获取模型响应
            response = self.interact(user_input, use_system_prompt=use_system_prompt)
            
            # 显示响应
            print(f"model: {response}")