import gradio as gr
import anthropic
from api_usage import get_subscription, check_key_availability, check_key_ant_availability, check_key_gemini_availability, check_key_azure_availability, get_azure_status, get_azure_deploy

def sort_key(key):
    _key = key.strip()
    if _key.startswith("sk-ant-"):
        return get_key_ant_info(_key)
    elif _key.startswith("AIzaSy"):
        return get_key_gemini_info(_key)
    elif "openai.azure.com" in _key.split(';')[0]:
        return get_key_azure_info(_key)
    elif _key.startswith("AKIA"):
        return get_key_aws_info(_key)
    else:
        return get_key_oai_info(_key)

def get_key_oai_info(key):
    # Return a dictionary containing key information
    key_avai = check_key_availability(key)
    info_dict = {#"account_name": "",
                 "key_type": "OpenAI",
                 "key_availability": True if key_avai else False,
                 "gpt4_availability": "",
                 "gpt4_32k_availability": "",
                 "default_org": "",                 
                 "org_description": "",
                 "organization": "",
                 "models": "",
                 "requests_per_minute": "",
                 "tokens_per_minute": "",
                 #"tokens_per_minute_left": "",
                 "quota": ""}
    if key_avai:
        info = get_subscription(key, key_avai)
        info_dict["gpt4_availability"] = info["has_gpt4"]
        info_dict["gpt4_32k_availability"] = info["has_gpt4_32k"]
        info_dict["default_org"] = info["default_org"]        
        info_dict["org_description"] = info["org_description"]
        info_dict["organization"] = info["organization"]      
        info_dict["models"] = info["models"]
        info_dict["requests_per_minute"] = info["rpm"]
        info_dict["tokens_per_minute"] = info["tpm"]
        info_dict["quota"] = info["quota"]
    return info_dict

def get_key_ant_info(key):
    # Return a dictionary containing key information
    ant = anthropic.Anthropic(api_key=key)
    key_avai = check_key_ant_availability(ant)
    info_dict = {#"account_name": "",
                 "key_type": "Anthropic Claude",
                 "key_availability": key_avai[0],
                 "status": key_avai[1],
                 "filter_response": key_avai[2]}
    return info_dict

def get_key_gemini_info(key):
    key_avai = check_key_gemini_availability(key)
    info_dict = {#"account_name": "",
                 "key_type": "Google Gemini",
                 "key_availability": key_avai[0],
                 "models": key_avai[1]}
    return info_dict

def get_key_azure_info(key):
    key_avai = check_key_azure_availability(key)
    if key_avai[0]:        
        azure_deploy = get_azure_deploy(key)
        status = get_azure_status(key, azure_deploy)
        info_dict = {#"account_name": "",
                     "key_type": "Microsoft Azure OpenAI",
                     "key_availability": key_avai[0],
                     "gpt4_32k_availability": status[1],
                     "gpt4_turbo_availability": status[2],
                     "gpt4_availability": status[3],
                     "gpt35_availability": status[4],
                     "models": key_avai[1],
                     "deployments": azure_deploy,
                     "moderation_response": status[0]}
    else:
        info_dict = {#"account_name": "",
                     "key_type": "Microsoft Azure OpenAI",
                     "key_availability": False,
                     "models": None,
                     "deployments": None,
                     "moderation_response": ""}
    return info_dict

def get_key_aws_info(key):
    info_dict = {#"account_name": "",
                     "key_type": "Amazon AWS Claude",
                     "status": " Soon™ "}
    return info_dict
    
def clear_inputs(text):
    return ""

with gr.Blocks() as demo:
    gr.Markdown('''
    # OpenAI/Anthropic/Google Gemini/Microsoft Azure API Key Status Checker
    
    *(Based on shaocongma, CncAnon1 and Drago key checkers)*
    
    (Azure endpoint's format: (https://)YOUR_RESOURCE_NAME.openai.azure.com;YOUR_API_KEY)
    ''')

    with gr.Row():
        with gr.Column():
            key = gr.Textbox(lines=1, max_lines=1, label="OpenAI/Anthropic/Gemini/Azure API Key")
            with gr.Row():
                clear_button = gr.Button("Clear")
                submit_button = gr.Button("Submit", variant="primary")
        with gr.Column():
            info = gr.JSON(label="OpenAI/Anthropic/Gemini/Azure API Key Information")

    clear_button.click(fn=clear_inputs, inputs=[key], outputs=[key])
    submit_button.click(fn=sort_key, inputs=[key], outputs=[info], api_name="sort_key")
demo.launch()