import gradio as gr
import requests
import os
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

api_keys = []
last_api_key_used = 0
num_api_keys = 0


def find_api_keys():
    global num_api_keys
    index = 0

    while True:
        index += 1
        api_key = os.getenv(f"OLLAMA_{index}")

        if api_key is None:
            if index == 1:
                print("0 API keys set. Format: OLLAMA_<api key number>")
            break

        api_keys.append(api_key)

    num_api_keys = len(api_keys)
    print(f"Found {num_api_keys} API keys.")


find_api_keys()


def get_api_key():
    global last_api_key_used

    if num_api_keys == 0:
        raise Exception("No API keys found. Format: OLLAMA_<number>")

    last_api_key_used += 1
    if last_api_key_used >= num_api_keys:
        last_api_key_used = 0

    return api_keys[last_api_key_used]


def perform_search(query):

    api_key = get_api_key()

    url = "https://ollama.com/api/web_search"

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }

    payload = {
        "query": query,
        "max_results": 1
    }

    response = requests.post(url, headers=headers, json=payload)
    response.raise_for_status()

    return response.json()


# -----------------------
# FastAPI
# -----------------------

app = FastAPI(title="Ollama Web Search API")


class SearchRequest(BaseModel):
    query: str
    max_results: int = 1


@app.post("/api/search")
def api_search(request: SearchRequest):

    try:
        api_key = get_api_key()

        url = "https://ollama.com/api/web_search"

        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }

        payload = {
            "query": request.query,
            "max_results": request.max_results
        }

        response = requests.post(url, headers=headers, json=payload)
        response.raise_for_status()

        return response.json()

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/health")
def health():
    return {"status": "ok"}


# -----------------------
# Gradio UI
# -----------------------

demo = gr.Interface(
    fn=perform_search,
    inputs=gr.Textbox(label="Search Query"),
    outputs=gr.JSON(label="Results"),
    title="Ollama Web Search API"
)

# Mount Gradio at root
app = gr.mount_gradio_app(app, demo, path="/")