import replicate
import os
import openai
import json
import time
import re
import requests
from openai import OpenAI
os.environ["OPENAI_API_KEY"] = "<Your API Key>"


def find_answer(input_string):
    # 使用正则表达式匹配 "Answer": 后的字符串
    match = re.search(r'"Answer":\s*"([^"]*)"', input_string)
    if match:
        return match.group(1)
    else:
        return None
    
    
def clean_and_convert_to_dict(json_str):
    clean_str = json_str.replace('```json', '').replace('```', '').replace('\n', '').strip()
    try:
        json_dict = json.loads(clean_str)
        return json_dict
    except json.JSONDecodeError as e:
        print("JSON decoding error:", e)
        error_position = e.pos
        if "Expecting ',' delimiter" in str(e):
            clean_str = clean_str[:error_position] + ',' + clean_str[error_position:]
        elif "Expecting ':' delimiter" in str(e):
            clean_str = clean_str[:error_position] + ':' + clean_str[error_position:]

        try:
            json_dict = json.loads(clean_str)
            return json_dict
        except json.JSONDecodeError:
            return "Error parsing JSON after attempting to fix."


def retry(attempts=3, delay=10):
    def decorator(func):
        def wrapper(*args, **kwargs):
            for i in range(attempts):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    print(f"Attempt {i+1}/{attempts} failed: {e}")
                    if i < attempts - 1:
                        time.sleep(delay) 
            return None
        return wrapper
    return decorator


@retry(attempts=3, delay=10)
def llama3_70b(system_prompt, prompt, clean: bool=False):
    url = "https://api.together.xyz/v1/chat/completions"
    content = [
                {
                    "role": "system",
                    "content": system_prompt
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ]
    payload = {
        "model": "meta-llama/Llama-3-70b-chat-hf",
        "temperature": 0.7,
        "frequency_penalty": 0,
        "presence_penalty": 1.15,
        "top_k": 0,
        "top_p": 0.9,
        "max_tokens": 1024,
        "messages": content
    }
    headers = {
        "accept": "application/json",
        "content-type": "application/json",
        "Authorization": "Bearer <Your API Key>"
    }

    response = requests.post(url, json=payload, headers=headers)
    result = json.loads(response.text)
    print(result['choices'][0]['message']['content'])
    if clean:
        return clean_and_convert_to_dict(result['choices'][0]['message']['content'])
    else:
        return result['choices'][0]['message']['content']
    
@retry(attempts=3, delay=10)
def GPT4(system_prompt, prompt, clean: bool=False):
    client = OpenAI()

    chat_completion = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content":system_prompt},
            {"role": "user", "content": prompt}
        ]
    )
    if clean:
        print(chat_completion.choices[0].message.content)
        result = clean_and_convert_to_dict(chat_completion.choices[0].message.content)
    else:
        result = chat_completion.choices[0].message.content
    print(result)
    return result


