# Copyright 2023 The OPRO Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""The utility functions for prompting GPT and Google Cloud models."""
import openai
import time
try:
    import google.generativeai as palm
except Exception:
    palm = None
try:
    from vllm import LLM, SamplingParams
except Exception:
    LLM = None
    SamplingParams = None
from pathlib import Path


# 缓存 vLLM 实例，避免重复加载
_llm_instance = None

def get_llm(local_model_path):
    if LLM is None:
        raise RuntimeError("vLLM not available")
    global _llm_instance
    if _llm_instance is None:
        assert local_model_path is not None, "model_path cannot be None"
        local_model_path = str(Path(local_model_path).resolve())
        _llm_instance = LLM(
           model=local_model_path,
            dtype="bfloat16",
            tensor_parallel_size=8,
            max_num_batched_tokens=8192,
            max_num_seqs=64,
            gpu_memory_utilization=0.7,
            enforce_eager=True,
            block_size=16,
            enable_chunked_prefill=True,
            trust_remote_code=True,
        )
    return _llm_instance

def call_local_server_single_prompt(prompt, local_model_path=None, temperature=0.8, max_decode_steps=512, **kwargs):
    """
    使用本地 vLLM 模型生成单个 prompt 的响应，替代原本 OpenAI API。
    """
    llm = get_llm(local_model_path)
    sampling_params = SamplingParams(
        temperature=temperature,
        top_p=0.9,
        max_tokens=max_decode_steps,
        skip_special_tokens=True  # 避免特殊字符触发协议错误
    )
    outputs = llm.generate([prompt], sampling_params)
    return outputs[0].outputs[0].text

def call_local_server_func(inputs, local_model_path=None, temperature=0.8, max_decode_steps=512, **kwargs):
    """
    批量处理多个输入 prompt。
    """
    
    assert local_model_path is not None, "local_model_path must be provided"
    # 强制类型检查
    if isinstance(inputs, bytes):
        inputs = inputs.decode('utf-8')
    outputs = []
    for input_str in inputs:
        output = call_local_server_single_prompt(
            input_str,
            local_model_path=local_model_path,
            temperature=temperature,
            max_decode_steps=max_decode_steps
        )
        outputs.append(output)
    return outputs


def call_openai_server_single_prompt(
    prompt, model="gpt-3.5-turbo", max_decode_steps=20, temperature=0.8
):
  """The function to call OpenAI server with an input string."""
  try:
    completion = openai.ChatCompletion.create(
        model=model,
        temperature=temperature,
        max_tokens=max_decode_steps,
        messages=[
            {"role": "user", "content": prompt},
        ],
    )
    return completion.choices[0].message.content
# 函数捕获了 ​​6 类常见异常​​，并在遇到错误时自动重试：
  except openai.error.Timeout as e: # API 请求超时
    retry_time = e.retry_after if hasattr(e, "retry_after") else 30
    print(f"Timeout error occurred. Retrying in {retry_time} seconds...")
    time.sleep(retry_time)
    return call_openai_server_single_prompt(
        prompt, max_decode_steps=max_decode_steps, temperature=temperature
    )

  except openai.error.RateLimitError as e: #请求频率超限（Rate Limit）
    retry_time = e.retry_after if hasattr(e, "retry_after") else 30
    print(f"Rate limit exceeded. Retrying in {retry_time} seconds...")
    time.sleep(retry_time)
    return call_openai_server_single_prompt(
        prompt, max_decode_steps=max_decode_steps, temperature=temperature
    )

  except openai.error.APIError as e: # API 错误（如服务器错误）
    retry_time = e.retry_after if hasattr(e, "retry_after") else 30
    print(f"API error occurred. Retrying in {retry_time} seconds...")
    time.sleep(retry_time)
    return call_openai_server_single_prompt(
        prompt, max_decode_steps=max_decode_steps, temperature=temperature
    )

  except openai.error.APIConnectionError as e: # API 连接错误（如网络问题）
    retry_time = e.retry_after if hasattr(e, "retry_after") else 30
    print(f"API connection error occurred. Retrying in {retry_time} seconds...")
    time.sleep(retry_time)
    return call_openai_server_single_prompt(
        prompt, max_decode_steps=max_decode_steps, temperature=temperature
    )

  except openai.error.ServiceUnavailableError as e: # 服务不可用（如服务器维护）
    retry_time = e.retry_after if hasattr(e, "retry_after") else 30
    print(f"Service unavailable. Retrying in {retry_time} seconds...")
    time.sleep(retry_time)
    return call_openai_server_single_prompt(
        prompt, max_decode_steps=max_decode_steps, temperature=temperature
    )

  except OSError as e: # 操作系统级连接错误（如网络中断）
    retry_time = 5  # Adjust the retry time as needed
    print(
        f"Connection error occurred: {e}. Retrying in {retry_time} seconds..."
    )
    time.sleep(retry_time)
    return call_openai_server_single_prompt(
        prompt, max_decode_steps=max_decode_steps, temperature=temperature
    )


def call_openai_server_func( #​​批量处理多个输入提示（prompts），通过 OpenAI API 并行或顺序获取多个生成结果​​。
    inputs, model="gpt-3.5-turbo", max_decode_steps=20, temperature=0.8
):
  """The function to call OpenAI server with a list of input strings."""
  if isinstance(inputs, str): # 将单个字符串转为列表，统一处理
    inputs = [inputs]
  outputs = []
  for input_str in inputs:
    output = call_openai_server_single_prompt(
        input_str,
        model=model,
        max_decode_steps=max_decode_steps,
        temperature=temperature,
    )
    outputs.append(output)
  return outputs

#通过 Google PaLM API（Cloud 版）调用 text-bison模型生成文本​​，并包含基本错误处理和自动重试机制。
def call_palm_server_from_cloud(
    input_text, model="text-bison-001", max_decode_steps=20, temperature=0.8
):
  if palm is None:
    raise RuntimeError("google.generativeai not available")
  assert isinstance(input_text, str)
  assert model == "text-bison-001"
  all_model_names = [
      m
      for m in palm.list_models()
      if "generateText" in m.supported_generation_methods
  ]
  model_name = all_model_names[0].name
  try:
    completion = palm.generate_text(
        model=model_name,
        prompt=input_text,
        temperature=temperature,
        max_output_tokens=max_decode_steps,
    )
    output_text = completion.result
    return [output_text]
  except Exception:
    retry_time = 10
    time.sleep(retry_time)
    return call_palm_server_from_cloud(
        input_text, max_decode_steps=max_decode_steps, temperature=temperature
    )

def refine_instruction(query: str) -> str:
    return f"""
你是一个“问题澄清与重写助手”。
请根据用户的原始问题：
【{query}】
生成不少于20条多角度、可直接执行的问题改写，每行一条。
"""

def refine_instruction_with_history(query: str, rejected_list: list) -> str:
    rejected_text = "\n".join(f"- {r}" for r in rejected_list) if rejected_list else ""
    return f"""
你是一个“问题澄清与重写助手”。
原始问题：
{query}

以下改写已被否定：
{rejected_text}

请从新的角度重新生成至少20条不同的改写问题，每条单独一行。
"""