import os
from openai import OpenAI
import json
from pprint import pprint
import random
from json.decoder import JSONDecodeError
import hashlib
from h4rm3l.model_adapters import olmo_adapter, openai_adapter, anthropic_adapter, anyscale_adapter
import logging
import sys
import traceback

MODEL_ADAPTERS = {
    'gpt-4': openai_adapter,
    'gpt-3.5-turbo': openai_adapter,
    'gpt-4o-2024-05-13': openai_adapter,
    'gpt-4-0125-preview': openai_adapter,
    'gpt-4-turbo-2024-04-09': openai_adapter,
    'claude-3-haiku-20240307': anthropic_adapter,
    'claude-3-sonnet-20240229': anthropic_adapter,
    'claude-3-opus-20240229': anthropic_adapter,
    'meta-llama/Llama-2-7b-chat-hf': anyscale_adapter,
    'meta-llama/Llama-2-13b-chat-hf': anyscale_adapter,
    'meta-llama/Llama-2-70b-chat-hf': anyscale_adapter,
    'meta-llama/Meta-Llama-3-8B-Instruct': anyscale_adapter,
    'meta-llama/Meta-Llama-3-70B-Instruct': anyscale_adapter,
    'mistralai/Mistral-7B-Instruct-v0.1': anyscale_adapter,
    'allenai/OLMo-7B': olmo_adapter,
    'allenai/OLMo-7B-SFT': olmo_adapter,
    'allenai/OLMo-7B-Instruct': olmo_adapter
}

def prompt_model(credentials, model_name, prompt=None, messages=None, max_tokens=None, temperature=None):
    if model_name not in MODEL_ADAPTERS:
        raise ValueError(f"Unknown model: {model_name}")
    
    try:
        return MODEL_ADAPTERS[model_name].prompt_model(
            credentials, 
            model_name, 
            prompt, 
            messages, 
            max_tokens, 
            temperature
        )
    except Exception as e:
        logging.critical("Uncaught prompting error " + str(e))
        logging.critical(traceback.format_exc())
        return ""


def extract_possible_json(txt):
    starts = []
    ends = []
    for start_char in ['{', '[']:
        if start_char in txt:
            starts.append(txt.index(start_char))
    for end_char in ['}', ']']:
        if end_char in txt:
            ends.append(txt.rfind(end_char))
    if not starts and not ends:
        return ""
    
    start = min(starts)
    end = max(ends)
    return txt[start:end+1]


def prompt_model_json(credentials, model_name, prompt=None, messages=None, max_tokens=None, temperature=None):
    text_response = ""
    try:
        text_response = prompt_model(credentials, model_name, prompt, messages, max_tokens, temperature)
        
        doc = json.loads(extract_possible_json(text_response))
        return doc
    except (KeyError, JSONDecodeError, ValueError, TypeError) as ex:
        logging.warning(f"json response parsing error {ex}\n" + text_response)
        return None

