from fastapi import FastAPI, Header, HTTPException, Request
import secrets, sqlite3, requests, os

app = FastAPI()
# Change this line in your main.py
DB_PATH = "/tmp/keys.db"
BRAIN_URL = "https://CooLLaMACEO-Gemma-Nano-Max.hf.space/completion"
HF_TOKEN = os.getenv("HF_TOKEN")

# Initialize Database
conn = sqlite3.connect(DB_PATH, check_same_thread=False)
conn.execute("CREATE TABLE IF NOT EXISTS keys (api_key TEXT PRIMARY KEY, user_id TEXT)")

@app.post("/generate-key")
async def create_key(user_id: str):
    new_key = f"gemma_max_sk_{secrets.token_urlsafe(16)}"
    conn.execute("INSERT INTO keys VALUES (?, ?)", (new_key, user_id))
    conn.commit()
    return {"api_key": new_key}

@app.post("/v1/chat")
async def proxy_to_brain(request: Request, authorization: str = Header(None)):
    # 1. Check API Key
    key = authorization.replace("Bearer ", "")
    res = conn.execute("SELECT user_id FROM keys WHERE api_key=?", (key,)).fetchone()
    if not res:
        raise HTTPException(status_code=401, detail="Invalid API Key")

    # 2. Forward to Private Brain
    user_data = await request.json()
    response = requests.post(
        BRAIN_URL,
        json={"prompt": user_data["prompt"], "n_predict": 512},
        headers={"Authorization": f"Bearer {HF_TOKEN}"}
    )
    return response.json()