app.py

import os
import time
from fastapi import FastAPI, Request, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response
import httpx
from bs4 import BeautifulSoup
from typing import List, Dict

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["GET"],
    allow_headers=["*"],
)

OLLAMA_LIBRARY_URL = "https://ollama.com/library"

# -----------------------------
# RATE LIMITING (25 req/day/IP)
# -----------------------------
RATE_LIMIT = 25
WINDOW_SECONDS = 60 * 60 * 24  # 24 hours
ip_store = {}  # { ip: { "count": int, "reset": timestamp } }


def check_rate_limit(ip: str):
    now = time.time()

    if ip not in ip_store:
        ip_store[ip] = {"count": 0, "reset": now + WINDOW_SECONDS}

    entry = ip_store[ip]

    if now > entry["reset"]:
        entry["count"] = 0
        entry["reset"] = now + WINDOW_SECONDS

    if entry["count"] >= RATE_LIMIT:
        raise HTTPException(
            status_code=429,
            detail="Daily limit reached: 25 images per IP"
        )

    entry["count"] += 1


# -----------------------------
# IMAGE GENERATION ENDPOINT
# -----------------------------
PKEY = os.getenv("POLLINATIONS_KEY", "")  # ensure this is set in your environment


@app.get("/genimg/{prompt}")
async def generate_image(prompt: str, request: Request):
    client_ip = request.client.host
    check_rate_limit(client_ip)

    url = f"https://gen.pollinations.ai/image/{prompt}?model=zimage&key={PKEY}"

    async with httpx.AsyncClient() as client:
        response = await client.get(url)

    if response.status_code != 200:
        raise HTTPException(
            status_code=500,
            detail=f"Pollinations error: {response.status_code}"
        )

    # Pollinations always returns JPEG
    return Response(
        content=response.content,
        media_type="image/jpeg"
    )



# -----------------------------
# EXISTING MODELS SCRAPER
# -----------------------------
@app.get("/models")
async def get_models() -> List[Dict]:
    async with httpx.AsyncClient() as client:
        response = await client.get(OLLAMA_LIBRARY_URL)
        html = response.text

    soup = BeautifulSoup(html, "html.parser")
    items = soup.select("li[x-test-model]")

    models = []
    for item in items:
        name = item.select_one("[x-test-model-title] span")
        description = item.select_one("p.max-w-lg")
        sizes = [el.get_text(strip=True) for el in item.select("[x-test-size]")]
        pulls = item.select_one("[x-test-pull-count]")
        tags = [t.get_text(strip=True) for t in item.select('span[class*="text-blue-600"]')]
        updated = item.select_one("[x-test-updated]")
        link = item.select_one("a")

        models.append({
            "name": name.get_text(strip=True) if name else "",
            "description": description.get_text(strip=True) if description else "No description",
            "sizes": sizes,
            "pulls": pulls.get_text(strip=True) if pulls else "Unknown",
            "tags": tags,
            "updated": updated.get_text(strip=True) if updated else "Unknown",
            "link": link.get("href") if link else None,
        })

    return models