import requests

from flask import Flask, request, jsonify, Response
import json
import pickle
import base64
import argparse


class VLLMClient:
    def __init__(self, base_url: str, authorization_token: str = None):
        """
        初始化 vLLM 客户端。
        
        :param base_url: vLLM 服务的基础 URL，例如 "http://localhost:8000"
        """
        self.base_url = base_url.rstrip("/")
        self.authorization_token = authorization_token
        self.headers = None
        if self.authorization_token:
            self.headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.authorization_token}"
            }
        else:
            self.headers = {"Content-Type": "application/json"}

    def generate_completion(self, model: str, prompt: str, max_tokens: int = 100, temperature: float = 1.0):
        """
        调用 /v1/completions 接口生成文本补全。
        
        :param model: 模型名称。
        :param prompt: 输入的提示文本。
        :param max_tokens: 生成的最大 token 数。
        :param temperature: 控制生成文本的随机性（越低越确定）。
        :return: 完成的生成结果。
        """
        
        url = f"{self.base_url}/v1/completions"
        payload = {
            "model": model,  # 模型名称
            "prompt": prompt,
            "max_tokens": max_tokens,
            "temperature": temperature,
        }



        try:
            response = requests.post(url, json=payload, headers=self.headers)
            response.raise_for_status()
            data = response.json()  # 返回 JSON 格式的响应
            return data
        except requests.RequestException as e:
            print(f"Error during request: {e}")
            return None

    def generate_chat_completion(self, model: str, messages: list, max_tokens: int = 100, temperature: float = 1.0):
        """
        调用 /v1/chat/completions 接口生成聊天补全。
        
        :param model: 模型名称。
        :param messages: 聊天消息列表。
        :param max_tokens: 生成的最大 token 数。
        :param temperature: 控制生成文本的随机性（越低越确定）。
        :return: 完成的生成结果。
        """
        url = f"{self.base_url}/v1/chat/completions"
        payload = {
            "model": model,  # 模型名称
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
        }
        
        try:
            response = requests.post(url, json=payload, headers=self.headers)
            response.raise_for_status()
            data = response.json()  # 返回 JSON 格式的响应
            return data
        except requests.RequestException as e:
            print(f"Error during request: {e}")
            return None

    def generate_completion_stream(self, model: str, prompt: str, max_tokens: int = 100, temperature: float = 1.0):
        """
        调用 /v1/completions 接口以流式方式生成文本补全。
        
        :param model: 模型名称。
        :param prompt: 输入的提示文本。
        :param max_tokens: 生成的最大 token 数。
        :param temperature: 控制生成文本的随机性（越低越确定）。
        :return: 完成的生成结果。
        """
        url = f"{self.base_url}/v1/completions"
        payload = {
            "model": model,  # 模型名称
            "prompt": prompt,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "stream": True  # 启用流式响应
        }

        # 加密数据并发送
        try:
            # 启用流式响应
            with requests.post(url, json=payload, headers=self.headers, stream=True) as response:
                response.raise_for_status()
                for chunk in response.iter_lines():
                    if chunk:
                        print(f"Received chunk: {chunk}")
                        # 你可以在这里处理每个 chunk，例如解析 JSON 或直接打印
        except requests.RequestException as e:
            print(f"Error during request: {e}")

    def generate_chat_completion_stream(self, model: str, messages: list, max_tokens: int = 100, temperature: float = 1.0):
        """
        调用 /v1/chat/completions 接口以流式方式生成聊天补全。
        
        :param model: 模型名称。
        :param messages: 聊天消息列表。
        :param max_tokens: 生成的最大 token 数。
        :param temperature: 控制生成文本的随机性（越低越确定）。
        :return: 完成的生成结果。
        """
        url = f"{self.base_url}/v1/chat/completions"
        payload = {
            "model": model,  # 模型名称
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "stream": True  # 启用流式响应
        }

        try:
            # 启用流式响应
            with requests.post(url, json=payload, headers=self.headers, stream=True) as response:
                response.raise_for_status()
                for chunk in response.iter_lines():
                    if chunk:
                        print(f"Received chunk: {chunk}")
                        # 你可以在这里处理每个 chunk，例如解析 JSON 或直接打印
        except requests.RequestException as e:
            print(f"Error during request: {e}")

# 示例用法
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="这是一个简单的命令行参数解析器示例")

    # 添加参数
    parser.add_argument('--ip', '-i', type=str, default="localhost", help='服务的IP地址')
    parser.add_argument('--port', '-p', type=str, default='8080', help='服务的端口')
    parser.add_argument('--endpoint', '-e', type=str, help='服务的接入点')
    parser.add_argument('--model_name', '-m', type=str, default="decrypted_model/Qwen32B", help='模型名称')
    parser.add_argument('--authorization_token', '-auth', type=str, default=None, help='授权令牌')
    args = parser.parse_args()

    base_url = f"http://{args.ip}:{args.port}"
    if args.endpoint:
        base_url = args.endpoint

    client = VLLMClient(base_url, args.authorization_token)  # 替换为你的 vLLM 服务地址
    prompt = "Once upon a time in a faraway land,"
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Tell me a story about San Francisco."}
    ]

    response_completion = client.generate_completion(args.model_name, prompt, max_tokens=50, temperature=0.7)
    print("Generated Text:", response_completion)
    
    response_chat_completion = client.generate_chat_completion(args.model_name, messages, max_tokens=100, temperature=0.7)
    print("Generated Chat Response:", response_chat_completion)

    client.generate_completion_stream(args.model_name, prompt, max_tokens=50, temperature=0.7)

    client.generate_chat_completion_stream(args.model_name, messages, max_tokens=100, temperature=0.7)