import json, re, copy
import importlib.util, os
from openai import AzureOpenAI  
from azure.identity import AzureCliCredential, get_bearer_token_provider  
        

def call_gpt(prompt, system_prompt="", model="o3-mini", **kwargs):
    if not isinstance(prompt, list):
        prompt = [{"role": "user", "content": prompt}] if not system_prompt else \
            [{"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}]
    else:
        role_0 = prompt[0].get("role", "")
        if role_0 != "system":
            prompt = [{"role": "system", "content": system_prompt}] + prompt if system_prompt else prompt 
        
    if model == "gpt-4o":
        spec = importlib.util.spec_from_file_location("call_gpt4", "/home/superbench/yaoxiang/code/DataSynthesis/call_api_local/call_api_3.py")
        test_api = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(test_api)
        call_gpt = test_api.call_gpt4                                                                                                                                                                                                                                                                                               
        try:
            response = call_gpt(prompt)[0]
            response = response.strip()
            return [response]
        except Exception as e:
            print(f"提取过程中出错: {e}")
            return None
    elif model == "o3-mini":
        endpoint = os.getenv("ENDPOINT_URL", "https://aims-oai-research-inference-uks.openai.azure.com/")  
        deployment = os.getenv("DEPLOYMENT_NAME", "o3-mini")  
            
        # Initialize Azure OpenAI client with Entra ID authentication
        token_provider = get_bearer_token_provider(  
            AzureCliCredential(),  
            "https://cognitiveservices.azure.com/.default"  
        )  
        
        client = AzureOpenAI(  
            azure_endpoint=endpoint,  
            azure_ad_token_provider=token_provider,  
            api_version="2025-01-01-preview",  
        )  
        
        try:
            completion = client.chat.completions.create(  
                model=deployment,  
                messages=prompt,
                max_completion_tokens=100000,
                stop=None,  
                stream=False  
            )  
            return [completion.choices[0].message.content]
        except Exception as e:
            print(f"提取过程中出错: {e}")
            return None
