import gradio as gr
import requests

def check_key_gemini_availability(key, ai_model):
    """
    Checks the availability of a Gemini API key.
    Args:
        key: The API key to check.
    Returns:
        A tuple containing a boolean indicating whether the key is valid and an error message (or None if the key is valid).
    """
    try:
        # First, check if the key allows access to the list of models
        url_getListModel = f"https://generativelanguage.googleapis.com/v1beta/models?key={key}"
        rq = requests.get(url_getListModel)
        rq.raise_for_status()  # Raise HTTPError for bad responses (4xx or 5xx)
        result = rq.json()
        if 'models' not in result:
            return False, f"Invalid key: 'models' field not found in response from list models endpoint."

        # Second, attempt to generate content to further validate the key
        ai_model_ = ai_model.strip() or "gemini-1.5-flash"
        models = [ x["name"].split('/',1)[1] for x in result['models'] ]
        if ai_model_ not in models:
            return False, f"Specified model {ai_model_} is not in available 'models'."
        url_generateContent = f"https://generativelanguage.googleapis.com/v1beta/models/{ai_model_}:generateContent?key={key}"
        # headers = {'Content-Type': 'application/json'}
        data = {"contents": [{"parts": [{"text": "Say \"this is a test\""}]}]}
        rq = requests.post(url_generateContent, json=data) #headers=headers, data=json.dumps(data)

        if 400 <= rq.status_code < 500:
            return False, f"{rq.status_code}"
        else:
            # rq.raise_for_status() # Raise other errors for other status codes
            print(key)
            return True, None

    except requests.exceptions.HTTPError as e:
        return False, f"Invalid: HTTP {e.response.status_code} - {e}"
    except Exception as e:
        return False, f"Error: {e}"

def check_keys(keys_text, ai_model):
    """
    Checks a list of Gemini API keys and categorizes them into valid and invalid.

    Args:
        keys_text: A string containing API keys separated by lines.

    Returns:
        A tuple containing two strings: a list of valid keys and a list of invalid keys with error messages.
    """
    key_list = keys_text.strip().splitlines()
    key_list = list(set(key_list))  # Remove duplicates

    valid_keys = []
    invalid_keys = []

    for key in key_list:
        is_valid, error_msg = check_key_gemini_availability(key, ai_model)
        if is_valid:
            valid_keys.append(key)
        else:
            invalid_keys.append(f"{key} - {error_msg}")

    return "\n".join(valid_keys), "\n".join(invalid_keys)

def clear_inputs():
    """Clears the input text area."""
    return ""

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

    with gr.Row():
        with gr.Column():
            keys = gr.TextArea(label="API Keys (by lines)")
            ai_model = gr.Textbox(label="gemini model")
            with gr.Row():
                clear_button = gr.Button("Clear")
                submit_button = gr.Button("Submit", variant="primary")
        with gr.Column():
            infos = gr.TextArea(label="Alive Keys")
            errs = gr.TextArea(label="Invalid Keys-status code")

    clear_button.click(fn=clear_inputs, outputs=[keys])
    submit_button.click(fn=check_keys, inputs=[keys,ai_model], outputs=[infos, errs], api_name="check_keys")

demo.launch()