import gradio as gr
import requests, re
from api_usage import (
    get_subscription, check_key_availability, get_orgs_me,
    check_key_ant_availability, check_ant_rate_limit, check_key_gemini_availability,
    check_key_azure_availability, get_azure_status, get_azure_deploy,
    check_key_mistral_availability, check_mistral_quota, check_key_replicate_availability,
    check_key_aws_availability, check_key_or_availability, check_key_or_limits,
    check_gcp_anthropic, check_groq_status, check_nai_status,
    check_elevenlabs_status, check_xai_status, check_stability_status,
    check_deepseek_status, not_supported
)

async def sort_key(key, rate_limit, claude_model):
    """
    Check a single API key and return its info dict.
    """
    _key = key.strip()

    # OpenRouter
    if re.match(r"^sk-or-v1-[a-z0-9]{64}$", _key):    
        return get_key_openrouter_info(_key)
    # Anthropic Claude
    if (re.match(r"^sk-ant-api03-[A-Za-z0-9\-_]{93}AA$", _key)
        or (_key.startswith("sk-ant-") and len(_key) == 93)
        or (len(_key) == 89 and re.match(r"^sk-[A-Za-z0-9]{86}$", _key))):
        return await get_key_ant_info(_key, rate_limit, claude_model)
    # Stability AI
    if re.match(r"^sk-[A-Za-z0-9]{48}$", _key) and len(_key) == 51 and 'T3BlbkFJ' not in _key:
        return get_key_stability_info(_key)
    # Deepseek
    if re.match(r"^sk-[a-f0-9]{32}$", _key):
        return get_key_deepseek_info(_key)
    # OpenAI
    if _key.startswith("sk-"):
        return get_key_oai_info(_key)
    # Google Gemini
    if _key.startswith("AIzaSy"):
        return get_key_gemini_info(_key)
    # NovelAI
    if _key.startswith("pst-"):
        return get_key_nai_info(_key)
    # Replicate
    if ((_key.startswith("r8_") and len(_key) == 40)
        or (_key.islower() and len(_key) == 40)):
        return get_key_replicate_info(_key)
    # xAI Grok
    if _key.startswith("xai-"):
        return get_key_xai_info(_key)
    # Azure OpenAI
    if len(_key.split(':')) == 2 and _key.split(':')[1].islower() and len(_key.split(':')[1]) == 32:
        endpoint, api_key = _key.split(':')
        return get_key_azure_info(f"{endpoint}.openai.azure.com", api_key)
    if "openai.azure.com" in _key:
        endpoint, api_key = _key.split(';')
        return get_key_azure_info(endpoint, api_key)
    # AWS
    if _key.startswith("AKIA") and len(_key.split(':')[0]) == 20 and _key.split(':')[0].isupper():
        return await get_key_aws_info(_key)
    # ElevenLabs
    if re.match(r"^[a-f0-9]{32}$", _key) or re.match(r"^sk_[a-f0-9]{48}$", _key):
        return get_key_elevenlabs_info(_key)
    # Mistral AI
    if re.match(r"^[A-Za-z0-9]{32}$", _key):
        return get_key_mistral_info(_key)
    # Groq
    if re.match(r"^gsk_[A-Za-z0-9]{20}WGdyb3FY[A-Za-z0-9]{24}$", _key):
        return get_key_groq_info(_key)
    # GCP Anthropic
    if re.match(r"[\w\-]+:[\w\-@\.]+:[\w\-]+:.+", _key):
        return await get_key_gcp_info(_key, 0)
    if re.match(r"[\w\-]+:[\w\-@\.]+:.+\\n", _key):
        return await get_key_gcp_info(_key, 1)

    # Fallback for unsupported formats
    return not_supported(_key)

async def sort_keys(keys_text, rate_limit, claude_model):
    """
    Split input by newline and lookup each key in turn.
    Returns a mapping from each key to its info dict.
    """
    keys = [k.strip() for k in keys_text.splitlines() if k.strip()]
    results = {}
    for k in keys:
        results[k] = await sort_key(k, rate_limit, claude_model)
    return results


def clear_inputs(text):
    return ""

with gr.Blocks() as demo:
    gr.Markdown('''
    # Multi-Key API Key Status Checker

    Enter one API key per line; this tool will check each in turn.
    ''')

    claude_options = [
        'claude-3-haiku-20240307',
        'claude-3-sonnet-20240229',
        'claude-3-opus-20240229',
        'claude-3-5-sonnet-20240620',
        'claude-3-5-sonnet-20241022',
        'claude-3-5-haiku-20241022'
    ]

    with gr.Row():
        with gr.Column():
            key_input = gr.Textbox(
                lines=5,
                placeholder="sk-... or AKIA...",
                label="API Keys (one per line)"
            )
            claude_model = gr.Dropdown(
                claude_options,
                value="claude-3-haiku-20240307",
                label="Claude API model",
                info="Model for filter_response and concurrent rate check"
            )
            rate_limit = gr.Checkbox(
                label="Check concurrent rate limit (Anthropic Claude)"
            )
            with gr.Row():
                clear_button = gr.Button("Clear")
                submit_button = gr.Button("Submit", variant="primary")
        with gr.Column():
            info = gr.JSON(label="API Keys Information", open=True)

    clear_button.click(fn=clear_inputs, inputs=[key_input], outputs=[key_input])
    submit_button.click(
        fn=sort_keys,
        inputs=[key_input, rate_limit, claude_model],
        outputs=[info],
        api_name="sort_keys"
    )

demo.launch()