import requests
import gradio as gr
import os

# ----------- Helper Functions -----------
def test_openai_key(key: str):
    try:
        headers = {"Authorization": f"Bearer {key}"}
        resp = requests.get("https://api.openai.com/v1/models", headers=headers)
        if resp.status_code == 200:
            return True, [m.get("id") for m in resp.json().get("data", [])]
        return False, resp.text
    except Exception as e:
        return False, str(e)

def test_gemini_key(key: str):
    try:
        url = f"https://generativelanguage.googleapis.com/v1beta/models?key={key}"
        resp = requests.get(url)
        if resp.status_code == 200:
            models = [m.get("name") for m in resp.json().get("models", [])]
            return True, models
        return False, resp.text
    except Exception as e:
        return False, str(e)

# ----------- Gradio Interface -----------
def validate_keys(openai_key: str, gemini_key: str):
    parts = ["## Results"]
    openai_models = None
    gemini_models = None

    if openai_key:
        valid, data = test_openai_key(openai_key)
        if valid:
            parts.append("### OpenAI Key: ✅ Valid")
            openai_models = data
        else:
            parts.append("### OpenAI Key: ❌ Invalid")
            parts.append(f"```\n{str(data)}\n```")

    if gemini_key:
        valid, data = test_gemini_key(gemini_key)
        if valid:
            parts.append("### Gemini Key: ✅ Valid")
            gemini_models = data
        else:
            parts.append("### Gemini Key: ❌ Invalid")
            parts.append(f"```\n{str(data)}\n```")

    if not openai_key and not gemini_key:
        parts.append("Enter at least one key to test.")

    summary_md = "\n\n".join(parts)
    return summary_md, (openai_models or []), (gemini_models or [])

custom_css = """
:root { --brand1: #2b6cb0; --brand2: #38b2ac; }
.header {text-align:center; padding: 1rem 0;}
.subtitle {text-align:center; color: #555; margin-bottom: 1rem;}
.card {border: 1px solid #e6e6e6; border-radius: 12px; padding: 16px; background: #fafafa;}
.btn-primary {background: var(--brand1) !important; color: white !important}
.models {min-height: 280px}
"""

with gr.Blocks(title="API Key Validator for OpenAI & Gemini", css=custom_css, theme=gr.themes.Soft(primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.teal)) as demo:
    gr.Markdown("""
    <div class="header">
      <h1>🔑 API Key Validator</h1>
    </div>
    <div class="subtitle">Validate keys and view available models for OpenAI & Gemini</div>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 🔐 Enter API Keys")
            with gr.Group():
                openai_key_box = gr.Textbox(label="OpenAI API Key", type="password", placeholder="sk-...", show_label=True)
                gemini_key_box = gr.Textbox(label="Gemini API Key", type="password", placeholder="AI Studio key...", show_label=True)
            test_btn = gr.Button("Test Keys", elem_classes=["btn-primary"])

        with gr.Column(scale=1):
            gr.Markdown("### ℹ️ Tips")
            gr.Markdown("- Keys are not stored.\n- Invalid keys show provider error JSON.\n- Enter either or both keys.")

    with gr.Row():
        with gr.Column():
            gr.Markdown("### ✅ Results")
            summary = gr.Markdown(elem_classes=["card"])

    with gr.Row():
        with gr.Column():
            gr.Markdown("### OpenAI Models")
            openai_models_json = gr.JSON(elem_classes=["card","models"])
        with gr.Column():
            gr.Markdown("### Gemini Models")
            gemini_models_json = gr.JSON(elem_classes=["card","models"])

    test_btn.click(
        fn=validate_keys,
        inputs=[openai_key_box, gemini_key_box],
        outputs=[summary, openai_models_json, gemini_models_json],
    )

if __name__ == "__main__":
    share = os.getenv("GRADIO_SHARE", "false").lower() in {"1", "true", "yes"}
    port = int(os.getenv("PORT", "7860"))
    demo.launch(server_name="0.0.0.0", server_port=port, share=share)