import re
import os
import yaml
import json
import requests
from dotenv import load_dotenv


load_dotenv()

ANTHROPIC_BASE_URL = os.getenv("ANTHROPIC_BASE_URL")
ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")

GEMINI_BASE_URL = os.getenv("GEMINI_BASE_URL")
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

DEEPSEEK_BASE_URL = os.getenv("DEEPSEEK_BASE_URL")
DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")

QWEN_BASE_URL = os.getenv("QWEN_BASE_URL")
QWEN_API_KEY = os.getenv("QWEN_API_KEY")

def read_json_to_string(json_file_path):
    """Read JSON file and convert to string.

    Args:
        json_file_path (str): Path to the JSON file

    Returns:
        str: JSON content as string
    """
    with open(json_file_path, 'r', encoding='utf-8') as file:
        data = json.load(file)
    return json.dumps(data, ensure_ascii=False)

def read_yaml_to_string(yaml_file_path):
    """Read YAML file and convert to string.

    Args:
        yaml_file_path (str): Path to the YAML file

    Returns:
        str: YAML content as string
    """
    with open(yaml_file_path, 'r', encoding='utf-8') as file:
        data = yaml.load(file, Loader=yaml.SafeLoader)
    return yaml.dump(data, allow_unicode=True)

def extract_json_from_string(string: str) -> dict:
    """Extract JSON content from a string with enhanced error handling.

    Args:
        string (str): String containing JSON

    Returns:
        dict: Extracted JSON object or empty list on failure
    """
    # use regex to extract json string
    pattern = re.compile(r'```json(.*?)```', re.DOTALL)
    match = pattern.search(string)
    if match:
        json_str = match.group(1).strip()
        try:
            return json.loads(json_str)
        except json.JSONDecodeError as e:
            print(f"JSON parsing error: {e}")
            print(f"JSON content preview: {json_str[:200]}...")
            
            # try to fix common json errors
            try:
                # try to fix missing comma problem
                fixed_json = json_str.replace('"\n  "', '",\n  "')
                return json.loads(fixed_json)
            except:
                pass
                
            try:
                # try to use more relaxed parsing
                import ast
                # replace double quotes with single quotes
                python_dict_str = json_str.replace('true', 'True').replace('false', 'False').replace('null', 'None')
                return ast.literal_eval(python_dict_str)
            except:
                pass
                
            print(f"Failed to automatically fix JSON, returning original string")
            return json_str
    else:
        print(f"Warning: extract json from string failed: \n{string[:200]}...")
        return []
    
def extract_yaml_from_string(string: str) -> str:
    """Extract YAML content from a string.

    Args:
        string (str): String containing YAML

    Returns:
        str: Extracted YAML content or empty string on failure
    """
    # use regex to extract yaml string
    pattern = re.compile(r'```yaml(.*?)```', re.DOTALL)
    match = pattern.search(string)
    if match:
        yaml_str = match.group(1)
        return yaml_str
    else:
        print(f"Warning: extract yaml from string failed: \n{string}")
        return ""

def get_claude_response(prompt: str, model: str = "claude-3-7-sonnet-20250219", max_tokens: int = 32768, temperature: float = 0.0) -> str:
    """Get response from Claude model.

    Args:
        prompt (str): Input prompt for the model
        model (str, optional): Claude model name. Defaults to "claude-3-7-sonnet-20250219".
        max_tokens (int, optional): Maximum tokens in response. Defaults to 32768.
        temperature (float, optional): Sampling temperature. Defaults to 0.0.

    Returns:
        str: Model response text
    """
    headers = {
        'anthropic-version': '2023-06-01',
        'content-type': 'application/json',
        "Authorization":f"Bearer {ANTHROPIC_API_KEY}"
    }
    data = {
        "model": model,
        "messages": [
            {"role": "user", 
             "content": prompt
            }
        ],
        "max_tokens": max_tokens,
        "temperature": temperature,
        "stream": False
    }
    response = requests.post(
        url=f"{ANTHROPIC_BASE_URL}/v1/messages",
        headers=headers,
        data=json.dumps(data),
        timeout=1800,
    )

    response = response.json()['content'][0]['text']

    return response

def get_claude_reasoning_response(prompt: str, model: str = "claude-3-7-sonnet-20250219", max_tokens: int = 128000) -> str:
    """Get reasoning response from Claude model with thinking enabled.

    Args:
        prompt (str): Input prompt for the model
        model (str, optional): Claude model name. Defaults to "claude-3-7-sonnet-20250219".
        max_tokens (int, optional): Maximum tokens in response. Defaults to 128000.

    Returns:
        str: Model response text with reasoning
    """
    headers = {
        'anthropic-version': '2023-06-01',
        'anthropic-beta': 'output-128k-2025-02-19',
        'content-type': 'application/json',
        "Authorization":f"Bearer {ANTHROPIC_API_KEY}"
    }
    data = {
        "model": model,
        "max_tokens": max_tokens,
        "thinking":{
            "type": "enabled",
            "budget_tokens": 32768
        },
        "messages": [
            {
                "role": "user", 
                "content": prompt
            }
        ],
        "stream": False
    }
    response = requests.post(
        url=f"{ANTHROPIC_BASE_URL}/v1/messages",
        headers=headers,
        data=json.dumps(data),
        timeout=1800,
    )
    
    response = response.json()['content'][1]['text']

    # # use regex to extract content between <think> and </think>
    # pattern = re.compile(r'<think>(.*?)</think>(.*?)<think>(.*?)</think>', re.DOTALL)
    # match = pattern.search(response)
    # if match:
    #     response = match.group(1)
    # else:
    #     response = response

    return response

def get_deepseek_response(prompt: str, model: str = "deepseek-v3", max_tokens: int = 32768, temperature: float = 0.0) -> str:
    """Get response from DeepSeek model.

    Args:
        prompt (str): Input prompt for the model
        model (str, optional): DeepSeek model name. Defaults to "deepseek-v3".
        max_tokens (int, optional): Maximum tokens in response. Defaults to 32768.
        temperature (float, optional): Sampling temperature. Defaults to 0.0.

    Returns:
        str: Model response text
    """
    url = f"{DEEPSEEK_BASE_URL}/v1/chat/completions"
    headers = {
        "Content-Type": "application/json; charset=utf-8",
        "Authorization": f"Bearer {DEEPSEEK_API_KEY}"
    }
    data = {
        "model": model,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "messages": [
            {
                "content": prompt,
                "role": "user"
            }
        ]
    }
    
    response = requests.post(url=url, headers=headers, data=json.dumps(data), timeout=180).json()
    return response['choices'][0]['message']['content']



def get_gemini_response(prompt: str, model: str = "gemini-2.5-pro", max_tokens: int = 32768, temperature: float = 0.0) -> str:
    """Get response from Gemini model.

    Args:
        prompt (str): Input prompt for the model
        model (str, optional): Gemini model name. Defaults to "gemini-2.5-pro".
        max_tokens (int, optional): Maximum tokens in response. Defaults to 32768.
        temperature (float, optional): Sampling temperature. Defaults to 0.0.

    Returns:
        str: Model response text
    """
    url = f"{GEMINI_BASE_URL}/v1beta/models/{model}:generateContent?key={GEMINI_API_KEY}"
    headers = {
        "Content-Type": "application/json",
    }
    data = {
        "contents": [
            {
                "parts": [
                    {
                        "text": prompt
                    }
                ]
            }
        ],
        "generationConfig": {
            "maxOutputTokens": max_tokens,
            "temperature": temperature
        }
    }
    response = requests.post(url=url, headers=headers, data=json.dumps(data), timeout=180).json()
    return response['candidates'][0]['content']['parts'][0]['text']



def get_openai_response(prompt: str, model: str = "gpt-4o", max_tokens: int = 16384, temperature: float = 0.0) -> str:
    """Get response from OpenAI model.

    Args:
        prompt (str): Input prompt for the model
        model (str, optional): OpenAI model name. Defaults to "gpt-4o".
        max_tokens (int, optional): Maximum tokens in response. Defaults to 16384.
        temperature (float, optional): Sampling temperature. Defaults to 0.0.

    Returns:
        str: Model response text
    """
    headers = {"Content-Type": "application/json", "Authorization": f"Bearer {OPENAI_API_KEY}"}
    
    data = {
        "model": model,
        "messages": [
            {
                "role": "user",
                "content": prompt
            }
        ],
        "max_tokens": max_tokens,
        "temperature": temperature,
        "stream": False
    }
    
    response = requests.post(
        url=f"{OPENAI_BASE_URL}/v1/chat/completions",
        headers=headers,
        data=json.dumps(data),
        timeout=180
    )
    
    return response.json()["choices"][0]["message"]["content"]

def get_openai_reasoning_response(prompt: str, model: str = "o3", max_tokens: int = 100000) -> str:
    """Get reasoning response from OpenAI model.

    Args:
        prompt (str): Input prompt for the model
        model (str, optional): OpenAI model name. Defaults to "o3".
        max_tokens (int, optional): Maximum tokens in response. Defaults to 100000.

    Returns:
        str: Model response text with reasoning
    """
    headers = {"Content-Type": "application/json", "Authorization": f"Bearer {OPENAI_API_KEY}"}
    
    data = {
        "model": model,
        "max_completion_tokens": max_tokens,
        "messages": [
            {
                "role": "user",
                "content": prompt
            }
        ],
        "stream": False
    }
    
    response = requests.post(
        url=f"{OPENAI_BASE_URL}/v1/chat/completions",
        headers=headers,
        data=json.dumps(data),
        timeout=180
    )

    return response.json()["choices"][0]["message"]["content"]


def get_qwen_response(prompt: str, model: str = "qwen2.5-72B", max_tokens: int = 32768, temperature: float = 0.0) -> str:
    """Get response from Qwen model.

    Args:
        prompt (str): Input prompt for the model
        model (str, optional): Qwen model name. Defaults to "qwen2.5-72B".
        max_tokens (int, optional): Maximum tokens in response. Defaults to 32768.
        temperature (float, optional): Sampling temperature. Defaults to 0.0.

    Returns:
        str: Model response text
    """
    url = f"{QWEN_BASE_URL}/v1/chat/completions"
    headers = {
        "Content-Type": "application/json; charset=utf-8",
        "Authorization": f"Bearer {QWEN_API_KEY}"
    }
    data = {
        "model": model,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "messages": [
            {
                "content": prompt,
                "role": "user"
            }
        ]
    }
    
    response = requests.post(url=url, headers=headers, data=json.dumps(data), timeout=180).json()
    return response['choices'][0]['message']['content']



def get_llm_response(prompt: str, model_name: str, temperature: float = 0.0) -> str:
    """Get response from a specified language model.

    Args:
        prompt (str): Input prompt for the model
        model_name (str): Name of the language model to use
        temperature (float, optional): Sampling temperature. Defaults to 0.0.

    Returns:
        str: Model response text

    Raises:
        ValueError: If the model name is not supported
    """
    if model_name == "claude-3-7-sonnet-20250219":
        return get_claude_response(prompt, model_name, temperature=temperature)
    elif model_name == "gpt-4o" or model_name == "gpt-4o-mini" or model_name == "gpt-4.1" or model_name == "gpt-4.5-preview" or model_name == "gpt-4-turbo" or model_name == "gpt-4" or model_name == "gpt-3.5-turbo":
        return get_openai_response(prompt, model_name, temperature=temperature)
    elif model_name == "deepseek-r1" or model_name == "deepseek-r1-32B" or model_name == "deepseek-v3":
        return get_deepseek_response(prompt, model_name, temperature=temperature)
    elif model_name == "gemini-2.5-pro-preview-06-05" or model_name == "gemini-2.5-flash-preview-04-17" or model_name == "gemini-2.0-flash" or model_name == "gemini-2.0-flash-lite" or model_name == "gemini-1.5-flash" or model_name == "gemini-1.5-pro":
        return get_gemini_response(prompt, model_name, temperature=temperature)
    elif model_name == "qwen2.5-72B":
        return get_qwen_response(prompt, model_name, temperature=temperature)
    else:
        raise ValueError(f"Model not supported: {model_name}")

def get_llm_reasoning_response(prompt: str, model_name: str) -> str:
    """Get reasoning response from a specified language model.

    Args:
        prompt (str): Input prompt for the model
        model_name (str): Name of the language model to use

    Returns:
        str: Model response text with reasoning

    Raises:
        ValueError: If the model name is not supported for reasoning
    """
    if model_name == "o1" or model_name == "o1-mini" or model_name == "o3" or model_name == "o3-mini" or model_name == "o4-mini":
        return get_openai_reasoning_response(prompt)
    elif model_name == "claude-3-7-sonnet-20250219":
        return get_claude_reasoning_response(prompt)
    else:
        raise ValueError(f"Model not supported for reasoning: {model_name}")


if __name__ == "__main__":
    print(get_llm_reasoning_response("hello", "o3"))
    print(get_llm_response("hello", "deepseek-v3"))
