from openai import OpenAI
import anthropic
from together import Together
import json
import re

# Initialize clients
anthropic_client = anthropic.Anthropic()
openai_client = OpenAI()
together_client = Together()

SYSTEM_PROMPT = """Please act as an impartial judge and evaluate based on the user's instruction. Your output format should strictly adhere to JSON as follows: {"feedback": "<write feedback>", "result": <numerical score>}. Ensure the output is valid JSON, without additional formatting or explanations."""

def get_openai_response(model_name, prompt):
    """Get response from OpenAI API"""
    try:
        response = openai_client.chat.completions.create(
            model=model_name,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt}
            ]
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"Error with OpenAI model {model_name}: {str(e)}"

def get_anthropic_response(model_name, prompt):
    """Get response from Anthropic API"""
    try:
        response = anthropic_client.messages.create(
            model=model_name,
            max_tokens=1000,
            temperature=0,
            system=SYSTEM_PROMPT,
            messages=[
                {"role": "user", "content": [{"type": "text", "text": prompt}]}
            ]
        )
        return response.content[0].text
    except Exception as e:
        return f"Error with Anthropic model {model_name}: {str(e)}"

def get_together_response(model_name, prompt):
    """Get response from Together API"""
    try:
        response = together_client.chat.completions.create(
            model=model_name,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt}
            ],
            stream=False
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"Error with Together model {model_name}: {str(e)}"

def get_model_response(model_name, model_info, prompt):
    """Get response from appropriate API based on model organization"""
    if not model_info:
        return "Model not found or unsupported."
    
    api_model = model_info['api_model']
    organization = model_info['organization']
    
    try:
        if organization == 'OpenAI':
            return get_openai_response(api_model, prompt)
        elif organization == 'Anthropic':
            return get_anthropic_response(api_model, prompt)
        else:
            # All other organizations use Together API
            return get_together_response(api_model, prompt)
    except Exception as e:
        return f"Error with {organization} model {model_name}: {str(e)}" 

def parse_model_response(response):
    try:
        # Debug print
        print(f"Raw model response: {response}")
        
        # First try to parse the entire response as JSON
        try:
            data = json.loads(response)
            return str(data.get('result', 'N/A')), data.get('feedback', 'N/A')
        except json.JSONDecodeError:
            # If that fails (typically for smaller models), try to find JSON within the response
            json_match = re.search(r'{.*}', response)
            if json_match:
                data = json.loads(json_match.group(0))
                return str(data.get('result', 'N/A')), data.get('feedback', 'N/A')
            else:
                return 'Error', f"Failed to parse response: {response}"
                
    except Exception as e:
        # Debug print for error case
        print(f"Failed to parse response: {str(e)}")
        return 'Error', f"Failed to parse response: {response}"