main.py

import os
import sys
import json
import sqlite3
import time
import asyncio
import shutil
import uuid
import threading
from fastapi import FastAPI, UploadFile, File, Form, Header, HTTPException, Depends, Query, Request
from fastapi.responses import JSONResponse, FileResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from typing import Optional
from contextlib import asynccontextmanager

STORAGE_DIR   = "/data" if os.path.exists("/data") else "."
LOCAL_DB_PATH = os.path.join(STORAGE_DIR, "metadata.db")
WEBHOOK_TOKEN = os.getenv("WEBHOOK_TOKEN", "YOUR_SECURE_PASSWORD_HERE")
DATABASE_URL  = os.getenv("DATABASE_URL")
USE_POSTGRES  = DATABASE_URL is not None
DB_WRITE_LOCK = threading.Lock()
DB_POOL       = None
_thread_local = threading.local()


# ─── DATABASE HELPERS ─────────────────────────────────────────────────────────

def get_sqlite_conn():
    if not hasattr(_thread_local, "conn"):
        _thread_local.conn = sqlite3.connect(LOCAL_DB_PATH, timeout=60.0, check_same_thread=False)
        _thread_local.conn.execute("PRAGMA journal_mode = WAL;")
        _thread_local.conn.execute("PRAGMA synchronous = NORMAL;")
    return _thread_local.conn

def init_db_pool():
    global DB_POOL
    if USE_POSTGRES:
        from psycopg2.pool import ThreadedConnectionPool
        DB_POOL = ThreadedConnectionPool(5, 20, DATABASE_URL)
        conn = DB_POOL.getconn()
        try:
            with conn.cursor() as cursor:
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS images (
                        id SERIAL PRIMARY KEY,
                        room_id TEXT,
                        session_id TEXT,
                        batch_index INTEGER,
                        prompt TEXT,
                        filename TEXT,
                        timestamp REAL
                    );
                """)
                cursor.execute("CREATE INDEX IF NOT EXISTS idx_room ON images(room_id);")
            conn.commit()
        except Exception:
            conn.rollback()
            raise
        finally:
            DB_POOL.putconn(conn)
    else:
        conn = sqlite3.connect(LOCAL_DB_PATH)
        cursor = conn.cursor()
        try:
            conn.execute("PRAGMA journal_mode = WAL;")
            conn.execute("BEGIN IMMEDIATE;")
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS images (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    room_id TEXT,
                    session_id TEXT,
                    batch_index INTEGER,
                    prompt TEXT,
                    filename TEXT,
                    timestamp REAL
                )
            """)
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_room ON images(room_id);")
            conn.commit()
        except Exception:
            conn.rollback()
            raise
        finally:
            cursor.close()
            conn.close()

def execute_thread_query(query_func, *args, is_write=False):
    if USE_POSTGRES:
        conn = DB_POOL.getconn()
        try:
            with conn.cursor() as cursor:
                result = query_func(cursor, *args)
            conn.commit()
            return result
        except Exception:
            try: conn.rollback()
            except Exception: pass
            raise
        finally:
            DB_POOL.putconn(conn)
    else:
        conn = get_sqlite_conn()
        if is_write:
            with DB_WRITE_LOCK:
                conn.execute("BEGIN IMMEDIATE;")
                cursor = conn.cursor()
                try:
                    result = query_func(cursor, *args)
                    conn.commit()
                    return result
                except Exception:
                    try: conn.rollback()
                    except Exception: pass
                    raise
                finally:
                    cursor.close()
        else:
            cursor = conn.cursor()
            try:
                return query_func(cursor, *args)
            finally:
                cursor.close()

def _fetch_expired_worker(cursor, cutoff_time: float, limit_val: int):
    placeholder = "%s" if USE_POSTGRES else "?"
    query = f"SELECT filename FROM images WHERE timestamp < {placeholder} LIMIT {placeholder};"
    cursor.execute(query, (cutoff_time, limit_val))
    return [r[0] for r in cursor.fetchall()]

def _drop_expired_worker(cursor, filenames: list):
    if not filenames: return
    placeholder  = "%s" if USE_POSTGRES else "?"
    placeholders = ",".join([placeholder] * len(filenames))
    query = f"DELETE FROM images WHERE filename IN ({placeholders});"
    cursor.execute(query, tuple(filenames))

def _read_history_worker(cursor, room_id: str):
    placeholder = "%s" if USE_POSTGRES else "?"
    cursor.execute(
        f"SELECT filename FROM images WHERE room_id = {placeholder} ORDER BY timestamp DESC;",
        (room_id,)
    )
    return cursor.fetchall()

def _insert_image_worker(cursor, room_id: str, session_id: str, batch_index: int, prompt: str, filename: str):
    placeholder = "%s" if USE_POSTGRES else "?"
    query = (
        f"INSERT INTO images (room_id, session_id, batch_index, prompt, filename, timestamp) "
        f"VALUES ({placeholder}, {placeholder}, {placeholder}, {placeholder}, {placeholder}, {placeholder});"
    )
    cursor.execute(query, (room_id, session_id, batch_index, prompt, filename, time.time()))

def disk_clear_files(expired_files: list):
    for old_file in expired_files:
        file_path = os.path.join(STORAGE_DIR, old_file)
        try:
            if os.path.exists(file_path): os.remove(file_path)
        except Exception as e:
            print(f"[Cleanup File Error] Storage unlink failure {old_file}: {e}", file=sys.stderr)


# ─── CLEANUP LOOP ─────────────────────────────────────────────────────────────

async def decoupled_cleanup_loop():
    error_backoff = 1.0
    while True:
        try:
            await asyncio.sleep(900)
            cutoff_time = time.time() - (2 * 3600)
            while True:
                expired_files = await asyncio.to_thread(
                    execute_thread_query, _fetch_expired_worker, cutoff_time, 500, is_write=False
                )
                if not expired_files: break
                await asyncio.to_thread(disk_clear_files, expired_files)
                await asyncio.to_thread(execute_thread_query, _drop_expired_worker, expired_files, is_write=True)
                await asyncio.sleep(0.1)
            error_backoff = 1.0
        except Exception as e:
            print(f"[Cleanup Master Loop Error]: {e}. Backing off for {error_backoff}s...", file=sys.stderr)
            await asyncio.sleep(error_backoff)
            error_backoff = min(error_backoff * 2, 300.0)


# ─── EVENT BROKER (SSE) ───────────────────────────────────────────────────────
# Manages in-memory SSE queues, one per room. Each connected browser tab holds
# one asyncio.Queue. Single-tab enforcement: when a new client subscribes to a
# room that already has a listener, the old queue receives a "superseded" event
# before being replaced, causing the old EventSource to close itself.

class EventBroker:
    def __init__(self):
        # room_id → single active Queue (one subscriber per room enforced)
        self._queues: dict[str, asyncio.Queue] = {}

    async def subscribe(self, room_id: str) -> asyncio.Queue:
        """Register a new subscriber. Evicts any existing subscriber for this room."""
        if room_id in self._queues:
            old_q = self._queues[room_id]
            try:
                # Tell the old tab to close its EventSource and stop retrying.
                await old_q.put({"event": "superseded", "data": {}})
            except Exception:
                pass  # old queue may already be garbage-collected

        q = asyncio.Queue()
        self._queues[room_id] = q
        return q

    def unsubscribe(self, room_id: str, q: asyncio.Queue):
        """Remove a subscriber only if it is still the current one for this room."""
        if self._queues.get(room_id) is q:
            del self._queues[room_id]

    async def broadcast(self, room_id: str, event: str, data: dict):
        """Push an event to the active subscriber for this room, if any."""
        q = self._queues.get(room_id)
        if q:
            try:
                await q.put({"event": event, "data": data})
            except Exception:
                pass

broker = EventBroker()


# ─── APP LIFESPAN ─────────────────────────────────────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    init_db_pool()
    cleanup_task = asyncio.create_task(decoupled_cleanup_loop())
    yield
    cleanup_task.cancel()
    if USE_POSTGRES and DB_POOL: DB_POOL.closeall()

app = FastAPI(lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "https://jiyuudayo-image.prankmoneyku.workers.dev",
        "https://api-gateway.prankmoneyku.workers.dev",
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ─── AUTH ─────────────────────────────────────────────────────────────────────

def verify_webhook_token(x_webhook_token: Optional[str] = Header(None)):
    if not x_webhook_token or x_webhook_token != WEBHOOK_TOKEN:
        raise HTTPException(status_code=403, detail="Unauthorized upload source")


# ─── ROUTES ───────────────────────────────────────────────────────────────────

@app.get("/")
async def root_status():
    return {"status": "online", "service": "Webhook Receiver Backend"}


@app.get("/stream/{room_id}")
async def stream_events(room_id: str):
    """
    Server-Sent Events endpoint. Each room_id supports one active subscriber;
    connecting a second client (new tab) evicts the first via a 'superseded' event.
    A ': ping' heartbeat is emitted every 15 seconds to keep both the Cloudflare
    Worker proxy connection and the HF infrastructure socket alive.
    """
    async def generator():
        q = await broker.subscribe(room_id)
        try:
            while True:
                try:
                    msg = await asyncio.wait_for(q.get(), timeout=15.0)
                    event_name = msg["event"]
                    event_data = json.dumps(msg["data"])
                    yield f"event: {event_name}\ndata: {event_data}\n\n"

                    # After delivering superseded, stop this generator — the client
                    # will close its EventSource on receipt and won't reconnect.
                    if event_name == "superseded":
                        return

                except asyncio.TimeoutError:
                    # Heartbeat — keeps the Cloudflare Worker's idle-subrequest
                    # timeout from closing the proxied connection. SSE comment lines
                    # (starting with ':') are valid per spec and ignored by clients.
                    yield ": ping\n\n"
        finally:
            broker.unsubscribe(room_id, q)

    return StreamingResponse(
        generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control":     "no-cache",
            "X-Accel-Buffering": "no",   # disable nginx response buffering on HF side
        }
    )


@app.post("/webhook", dependencies=[Depends(verify_webhook_token)])
async def receive_webhook(
    file:           UploadFile = File(...),
    x_room_id:      str = Header(...),
    x_session_id:   str = Form(...),
    x_batch_index:  int = Form(...),
    x_image_prompt: str = Form(...),
):
    await file.seek(0)
    secure_uuid     = uuid.uuid4().hex
    unique_filename = f"room_{x_room_id}_{secure_uuid}_{x_batch_index}.png"
    full_path       = os.path.join(STORAGE_DIR, unique_filename)

    try:
        def save_to_nas():
            with open(full_path, "wb") as buffer:
                shutil.copyfileobj(file.file, buffer)

        await asyncio.to_thread(save_to_nas)
        await asyncio.to_thread(
            execute_thread_query,
            _insert_image_worker,
            x_room_id, x_session_id, x_batch_index, x_image_prompt, unique_filename,
            is_write=True
        )
    except Exception as e:
        if os.path.exists(full_path): os.remove(full_path)
        raise HTTPException(status_code=500, detail=f"Pipeline ingestion failure: {str(e)}")

    # Broadcast to any connected SSE client for this room.
    # Fire-and-forget via create_task — the GPU script gets its 200 OK immediately
    # regardless of whether a browser tab is currently subscribed.
    asyncio.create_task(broker.broadcast(x_room_id, "new_image", {
        "filename":    unique_filename,
        "room_id":     x_room_id,
        "session_id":  x_session_id,
        "batch_index": x_batch_index,
        "prompt":      x_image_prompt,
        "task_id":     f"task-{secure_uuid}",
    }))

    return {"status": "success", "saved_as": unique_filename}


@app.post("/registry-webhook", dependencies=[Depends(verify_webhook_token)])
async def registry_webhook(request: Request):
    """
    Receives registry updates from the Cloudflare Worker (forwarded from
    GPU scripts) and broadcasts them via SSE to any connected browser tab
    subscribed to this room.  Fire-and-forget — the GPU script's 200 OK
    is not delayed by SSE delivery.
    """
    body = await request.json()
    room_id = body.get("room_id", "")
    if not room_id:
        return JSONResponse(content={"error": "Missing room_id"}, status_code=400)

    asyncio.create_task(broker.broadcast(room_id, "registry_update", {
        "provider":      body.get("provider", ""),
        "cached_models": body.get("cached_models", []),
        "total_space":   body.get("total_space", 0),
        "used_space":    body.get("used_space", 0),
    }))
    return {"ok": True}


@app.get("/history")
async def get_history(x_room_id: str = Header(...)):
    rows    = await asyncio.to_thread(execute_thread_query, _read_history_worker, x_room_id, is_write=False)
    history = {}
    for idx, (filename,) in enumerate(rows):
        fake_id         = f"task-{idx}"
        history[fake_id] = {
            "outputs": {
                "9": {
                    "images": [{"filename": filename, "type": "output", "subfolder": ""}]
                }
            }
        }
    return JSONResponse(content=history)


@app.get("/view")
async def view_image(filename: str = Query(...), room_id: str = Query(...)):
    safe_path = os.path.basename(filename)
    if not safe_path.startswith(f"room_{room_id}_") or not safe_path.endswith(".png"):
        raise HTTPException(status_code=403, detail="Access denied")
    full_path = os.path.join(STORAGE_DIR, safe_path)
    if os.path.exists(full_path):
        return FileResponse(full_path, media_type="image/png")
    return JSONResponse(content={"error": "Not Found"}, status_code=404)