import APIKey
import requests


def check_azure(key: APIKey):
    line = key.api_key.split(':')
    key.endpoint = line[0]
    api_key = line[1]

    deployments = get_deployments(key, api_key)
    if deployments is None:
        return

    # deal with dall-e separately
    key.dalle_deployments = [deployment['id'] for deployment in deployments if deployment['model'] == 'dall-e-3']
    key.deployments = [(deployment['id'], deployment['model'], test_deployment(key, api_key, deployment['id'], deployment['model'])) for deployment in deployments if deployment['model'].startswith('gpt')]

    if key.deployments is None or not key.deployments:
        if not deployments:
            return
        key.best_deployment = deployments[0]['id']
        key.model = deployments[0]['model']
        key.deployments = key.deployments[1:]
        return True

    key.deployments = [(deployment_id, 'gpt-4-turbo' if model == 'gpt-4' and deployment_id in key.has_gpt4_turbo else model, is_unfiltered) for deployment_id, model, is_unfiltered in key.deployments]
    key.deployments = sorted(key.deployments, key=sort_deployments)
    key.best_deployment, key.model, is_unfiltered = key.deployments[0]
    key.deployments = key.deployments[1:]

    if is_unfiltered:
        key.unfiltered = True
    return True


model_priority = ['gpt-4o', 'gpt-4-turbo', 'gpt-4-32k', 'gpt-4', 'gpt-35-turbo-16k', 'gpt-35-turbo']


def sort_deployments(deployment):
    deployment_id, model, is_unfiltered = deployment
    for i, base_model in enumerate(model_priority):
        if model.startswith(base_model):
            return (i - 0.5) if is_unfiltered else i
    return len(model_priority)


def get_deployments(key: APIKey, api_key):
    deployment_endpoint = f'https://{key.endpoint}.openai.azure.com/openai/deployments?api-version=2022-12-01'
    response = requests.get(deployment_endpoint, headers={'api-key': api_key})
    if response.status_code != 200:
        return
    return response.json()['data']


def test_deployment(key: APIKey, api_key, deployment_id, deployment_model):
    mp = 'hi'
    check_turbo = deployment_model == 'gpt-4'
    completions_endpoint = f'https://{key.endpoint}.openai.azure.com/openai/deployments/{deployment_id}/chat/completions?api-version=2024-02-01'
    data = {
        'messages': [{'role': 'user', 'content': f'{mp}'}],
        'max_tokens': 8200 if check_turbo else 1,
    }
    response = requests.post(completions_endpoint, headers={'api-key': api_key, 'accept': 'application/json'}, json=data)

    if response.status_code == 200:
        return 'choices' not in response.json() or 'content_filter_results' not in response.json().get('choices', [{}])[0]
    elif response.status_code == 400:
        if check_turbo:
            if response.json()["error"]["code"] != "context_length_exceeded":
                key.has_gpt4_turbo.append(deployment_id)
            data['max_tokens'] = 1
            response = requests.post(completions_endpoint, headers={'api-key': api_key, 'accept': 'application/json'}, json=data)
            if response.status_code == 200:
                return 'choices' not in response.json() or 'content_filter_results' not in response.json().get('choices', [{}])[0]
            if response.status_code == 400:
                return False
        return False
    return


def pretty_print_azure_keys(keys):
    print('-' * 90)
    print(f'Validated {len(keys)} Azure keys:')
    unfiltered = 0
    keys = sorted(keys, key=lambda x: (x.unfiltered, bool(x.dalle_deployments)), reverse=True)
    for key in keys:
        if key.unfiltered:
            unfiltered += 1
        key_string = (f'{key.api_key}'
                      + f' | best deployment - {key.best_deployment}'
                      + f' | top model - {key.model}')
        if key.deployments:
            key_string += ' | other chat deployments - ['
            for deployment_id, model, filter_status in key.deployments:
                key_string += (f"'{deployment_id}'" + (' - unfiltered' if filter_status else '') + ', ')
            key_string = key_string.rstrip(', ') + ']'
        key_string += (' | !!!UNFILTERED!!!' if key.unfiltered else '')
        key_string += (f' | dall-e 3 deployments found on - {key.dalle_deployments}' if key.dalle_deployments else '')
        print(key_string)
    print(f'\n--- Total Valid Azure Keys: {len(keys)} ({unfiltered} unfiltered) ---\n')