server.py

#!/usr/bin/env python3
"""Echoverse — FastAPI server."""

import asyncio
import json
import os
import traceback
import threading
import uuid
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from dotenv import load_dotenv
load_dotenv()

from google.oauth2 import id_token as google_id_token
from google.auth.transport import requests as google_requests
from fastapi import FastAPI, UploadFile, File, Request
from fastapi.staticfiles import StaticFiles
from fastapi.responses import JSONResponse
from pydantic import BaseModel

from engine.pipeline import SoundToWorldPipeline
from engine import storage, gcs

app = FastAPI(title="Echoverse")
pipeline = SoundToWorldPipeline()

# Local temp dir for generation (uploaded to GCS afterwards in cloud)
OUTPUT_DIR = Path("/tmp/output") if gcs.is_enabled() else Path("output")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

storage.init_db()

# In-memory stores for active pipeline state
sessions = {}
jobs = {}

USE_GCS = gcs.is_enabled()

# --- Google SSO (OAuth2 ID token verification) ---

GOOGLE_CLIENT_ID = os.environ.get("GOOGLE_CLIENT_ID", "")

ALLOWED_EMAILS = [
    e.strip() for e in os.environ.get("ALLOWED_EMAILS", "").split(",") if e.strip()
]


def _verify_token(request: Request) -> dict | None:
    """Verify Google ID token from Authorization header."""
    auth_header = request.headers.get("authorization", "")
    if not auth_header.startswith("Bearer "):
        return None
    token = auth_header[7:]
    try:
        return google_id_token.verify_oauth2_token(
            token, google_requests.Request(), GOOGLE_CLIENT_ID,
        )
    except Exception:
        return None


def _check_auth(request: Request) -> tuple[str | None, str | None]:
    """Returns (user_id, error_message). If error_message is set, reject the request."""
    if not GOOGLE_CLIENT_ID:
        # Auth not configured — allow all (local dev)
        return request.headers.get("x-user-id", "anonymous"), None
    decoded = _verify_token(request)
    if not decoded:
        return None, "Authentication required"
    email = decoded.get("email", "")
    if ALLOWED_EMAILS and email not in ALLOWED_EMAILS:
        return None, f"Access denied for {email}"
    return decoded["sub"], None


class GenerateRequest(BaseModel):
    session_id: str
    answers: dict = {}
    mode: str = "image"


def _ext(filename):
    return Path(filename).suffix if filename else ".webm"


def _upload_and_url(local_path, gcs_path):
    """Upload file to GCS if enabled, return the URL to use."""
    if USE_GCS:
        return gcs.upload(local_path, gcs_path)
    return f"/output/{Path(local_path).name}"


# ----- Endpoints -----

@app.get("/api/auth/config")
async def auth_config():
    """Return Google OAuth client ID for the frontend."""
    return {"client_id": GOOGLE_CLIENT_ID}


@app.post("/api/analyze")
async def analyze(request: Request, audio: UploadFile = File(...)):
    user_id, err = _check_auth(request)
    if err:
        return JSONResponse({"error": err}, status_code=401)
    session_id = uuid.uuid4().hex[:8]

    # Save uploaded audio and convert to WAV
    upload_path = OUTPUT_DIR / f"{session_id}_upload{_ext(audio.filename)}"
    upload_path.write_bytes(await audio.read())

    wav_path = OUTPUT_DIR / f"{session_id}_audio.wav"
    await asyncio.to_thread(pipeline.convert_to_wav, str(upload_path), str(wav_path))

    result = await asyncio.to_thread(pipeline.analyze_structured, str(wav_path))

    # Upload audio to GCS for later use in reveal
    audio_url = _upload_and_url(str(wav_path), f"{session_id}/audio.wav")

    sessions[session_id] = {
        "user_id": user_id,
        "audio_path": str(wav_path),
        "audio_url": audio_url,
        "analysis": result,
    }

    return {
        "session_id": session_id,
        "analysis": result.get("analysis", ""),
        "sound_tags": result.get("sound_tags", []),
        "confidence": result.get("confidence", "high"),
        "questions": result.get("questions", []),
    }


@app.post("/api/generate")
async def generate(request: Request, req: GenerateRequest):
    user_id, err = _check_auth(request)
    if err:
        return JSONResponse({"error": err}, status_code=401)

    session = sessions.get(req.session_id)
    if not session:
        return JSONResponse({"error": "Session not found"}, status_code=404)

    analysis_text = session["analysis"].get("analysis", "")
    sound_tags = session["analysis"].get("sound_tags", [])

    job_id = uuid.uuid4().hex[:8]
    jobs[job_id] = {
        "status": "exploring_worlds",
        "session_id": req.session_id,
        "user_id": user_id,
        "world": {},
        "mode": req.mode,
        "result": {},
        "branches": [],
        "evaluation": None,
        "verification": None,
    }

    storage.save_creation(job_id, user_id, req.mode)

    thread = threading.Thread(
        target=_explore_and_generate,
        args=(job_id, analysis_text, sound_tags, req.answers or None),
        daemon=True,
    )
    thread.start()

    return {"job_id": job_id}


@app.get("/api/status/{job_id}")
async def status(job_id: str):
    job = jobs.get(job_id)
    if not job:
        return JSONResponse({"error": "Job not found"}, status_code=404)

    resp = {
        "status": job["status"],
        "world_name": job["world"].get("world_name", ""),
        "concept": job["world"].get("concept", ""),
        "sound_mapping": job["world"].get("sound_mapping", {}),
        "mode": job.get("mode", "image"),
    }

    if job.get("branches"):
        resp["branches"] = [
            {"world_name": w.get("world_name", ""), "concept": w.get("concept", "")}
            for w in job["branches"]
        ]
    if job.get("evaluation"):
        resp["evaluation"] = job["evaluation"]
    if job.get("verification"):
        resp["verification"] = {
            "score": job["verification"].get("score", 0),
            "feedback": job["verification"].get("feedback", ""),
        }

    if job["status"] == "complete":
        resp["result"] = job["result"]
    elif job["status"] == "error":
        resp["error"] = job.get("error", "Unknown error")

    return resp


@app.get("/api/gallery")
async def gallery(request: Request):
    user_id, err = _check_auth(request)
    if err:
        return JSONResponse({"error": err}, status_code=401)
    return {"creations": storage.get_gallery(user_id)}


@app.get("/api/creation/{creation_id}")
async def get_creation(request: Request, creation_id: str):
    user_id, err = _check_auth(request)
    if err:
        return JSONResponse({"error": err}, status_code=401)
    creation = storage.get_creation(creation_id)
    if not creation:
        return JSONResponse({"error": "Creation not found"}, status_code=404)
    return creation


# ----- Background generation pipeline -----

def _explore_and_generate(job_id, analysis_text, sound_tags, answers):
    job = jobs[job_id]
    session = sessions[job["session_id"]]
    sid = job["session_id"]
    mode = job.get("mode", "image")

    try:
        # Step 1: World exploration + reality image in parallel
        job["status"] = "exploring_worlds"
        print(f"[{job_id}] Exploring 3 world branches + reality image in parallel...")

        reality_prompt = session["analysis"].get("reality_prompt", "")
        if not reality_prompt:
            reality_prompt = f"Photorealistic photo of: {analysis_text[:500]}"
        reality_path = str(OUTPUT_DIR / f"{sid}_reality.png")

        with ThreadPoolExecutor(max_workers=2) as pool:
            future_explore = pool.submit(
                pipeline.explore_worlds, analysis_text, sound_tags, answers, 3,
            )
            future_reality = pool.submit(
                pipeline.generate_image, reality_prompt, reality_path, "1:1",
            )
            result = future_explore.result()
            future_reality.result()

        job["branches"] = result["worlds"]
        job["evaluation"] = result["evaluation"]
        world = result["winner"]
        job["world"] = world

        reality_url = _upload_and_url(reality_path, f"{sid}/reality.png")
        job["result"]["reality_image"] = reality_url
        print(f"[{job_id}] Winner: {world.get('world_name')}")

        # Step 2: Generate world image with verification
        job["status"] = "generating_world"
        image_prompt = world.get("image_prompt", world.get("video_prompt", ""))
        world_path = str(OUTPUT_DIR / f"{sid}_fantasy.png")

        _, verification = pipeline.generate_with_verification(
            image_prompt=image_prompt,
            output_path=world_path,
            analysis_text=analysis_text,
            sound_tags=sound_tags,
            world_name=world.get("world_name", ""),
            concept=world.get("concept", ""),
            sound_mapping=world.get("sound_mapping", {}),
            aspect_ratio="1:1",
        )
        world_url = _upload_and_url(world_path, f"{sid}/fantasy.png")
        job["result"]["world_image"] = world_url
        job["verification"] = verification
        print(f"[{job_id}] Verification score: {verification.get('score', '?')}/10")

        # Step 3: Stitch side-by-side comparison
        job["status"] = "stitching"
        compare_path = str(OUTPUT_DIR / f"{sid}_compare.png")
        pipeline.stitch_side_by_side(reality_path, world_path, compare_path)
        compare_url = _upload_and_url(compare_path, f"{sid}/compare.png")
        job["result"]["compare_image"] = compare_url

        job["result"]["audio_url"] = session["audio_url"]

        # Step 4: Video generation (optional)
        if mode == "video":
            job["status"] = "generating_video"
            print(f"[{job_id}] Generating video with Veo 3.1...")
            video_prompt = world.get("video_prompt", world.get("image_prompt", ""))
            raw_video_path = str(OUTPUT_DIR / f"{sid}_raw.mp4")
            pipeline.generate_video(video_prompt, raw_video_path)

            job["status"] = "compositing_video"
            print(f"[{job_id}] Compositing audio onto video...")
            final_video_path = str(OUTPUT_DIR / f"{sid}_world.mp4")
            pipeline.composite(raw_video_path, session["audio_path"], final_video_path)
            video_url = _upload_and_url(final_video_path, f"{sid}/world.mp4")
            job["result"]["video_url"] = video_url
            print(f"[{job_id}] Video complete!")

        job["status"] = "complete"
        print(f"[{job_id}] Complete!")

        storage.update_creation(
            job_id,
            world_name=world.get("world_name", ""),
            concept=world.get("concept", ""),
            sound_tags=json.dumps(sound_tags) if not USE_GCS else sound_tags,
            sound_mapping=json.dumps(world.get("sound_mapping", {})) if not USE_GCS else world.get("sound_mapping", {}),
            reality_image=job["result"].get("reality_image", ""),
            world_image=job["result"].get("world_image", ""),
            compare_image=job["result"].get("compare_image", ""),
            video_url=job["result"].get("video_url", ""),
            audio_url=job["result"].get("audio_url", ""),
            status="complete",
            verification_score=verification.get("score", 0) if verification else 0,
        )

    except Exception as e:
        job["status"] = "error"
        job["error"] = str(e)
        storage.update_creation(job_id, status="error")
        traceback.print_exc()


# ----- Static files -----

# Serve local output files only in dev (in cloud, files are on GCS)
if not USE_GCS:
    app.mount("/output", StaticFiles(directory="output"), name="output")

app.mount("/", StaticFiles(directory="static", html=True), name="static")