import requests
import gradio as gr
import os

# Environment-based configuration
# For Hugging Face Spaces, set OLLAMA_BASE_URL in Space secrets
# For local development, it defaults to localhost
OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
OLLAMA_URL = f"{OLLAMA_BASE_URL}/api/generate"

def generate_response(model, prompt):
    try:
        response = requests.post(
            OLLAMA_URL,
            json={
                "model": model,
                "prompt": prompt,
                "stream": False
            },
            timeout=30
        )
        response.raise_for_status()
        return response.json()["response"]
    except requests.exceptions.RequestException as e:
        return f"Error: Could not connect to Ollama server at {OLLAMA_BASE_URL}. Details: {str(e)}"
    except Exception as e:
        return f"Error: {str(e)}"

# 🔥 Available models (you can update this)
models = [
    "mistral",
    "llama3",
    "tinyllama",
    "hf.co/bartowski/Llama-3.2-1B-Instruct-GGUF"
]

# UI
with gr.Blocks() as app:
    gr.Markdown(f"# 🚀 Your Ollama UI")
    gr.Markdown(f"**Connected to:** `{OLLAMA_BASE_URL}`")

    model_dropdown = gr.Dropdown(models, value="mistral", label="Select Model")
    prompt_input = gr.Textbox(lines=5, label="Enter Prompt")

    generate_btn = gr.Button("Generate")

    output = gr.Textbox(label="Response")

    generate_btn.click(
        fn=generate_response,
        inputs=[model_dropdown, prompt_input],
        outputs=output
    )

app.launch(server_name="0.0.0.0", server_port=7860)