api.py

"""Minimal API for nearby discovery, PixArt images, and the POI map."""

from __future__ import annotations

from datetime import datetime, timezone
import json
import os
from pathlib import Path
import threading
from typing import Any
from uuid import uuid4
from zoneinfo import ZoneInfo

from benchmark_pixart import benchmark_pixart
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import RedirectResponse, Response
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field, model_validator

from discovery import DiscoveryError, discover_nearby, fetch_ground_map_image, recommendation_message
from map_writer import (
    GROUND_MAP_PATH,
    MAP_CONFIG_PATH,
    MAP_IMAGES_ROOT,
    MAP_ROOT,
    image_path_for,
    map_image_url,
    write_active_map,
    write_ground_map_image,
)
from pixart import RocmUnavailableError, generate_hero_image, runtime_info
from user_context_schema import schema_document


REPO_ROOT = Path(__file__).resolve().parent
ARTIFACTS_ROOT = REPO_ROOT / "artifacts"
BENCHMARKS_ROOT = ARTIFACTS_ROOT / "benchmarks"


def _load_local_env() -> None:
    """Load simple KEY=VALUE API credentials without adding a runtime dependency."""

    env_path = REPO_ROOT / ".env"
    if not env_path.exists():
        return
    for raw_line in env_path.read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        key = key.strip()
        if key:
            os.environ.setdefault(key, value.strip().strip('"').strip("'"))


_load_local_env()

app = FastAPI(title="Nearby Mood Discovery")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


class RecommendationRequest(BaseModel):
    mood: str = Field(min_length=1, max_length=500)
    latitude: float = Field(ge=-90, le=90)
    longitude: float = Field(ge=-180, le=180)
    radius_m: int = Field(default=3000, ge=100, le=40_000)
    max_results: int = Field(default=8, ge=1, le=20)
    image_limit: int = Field(default=8, ge=0, le=20)
    generate_images: bool = True
    inference_steps: int = Field(default=20, ge=1, le=50)
    guidance_scale: float = Field(default=4.5, ge=1.0, le=20.0)
    seed: int = Field(default=42, ge=0, le=2_147_483_647)
    timezone_name: str = Field(default="America/Los_Angeles")
    compile_transformer: bool = False

    @model_validator(mode="after")
    def _default_image_limit_to_max_results(self) -> "RecommendationRequest":
        """Generate a hero image for every returned POI unless the client caps it explicitly."""

        if "image_limit" not in self.model_fields_set:
            self.image_limit = self.max_results
        return self


class PixArtBenchmarkRequest(BaseModel):
    prompt: str = Field(default="An editorial hero image of a bowl of ramen, warm evening light, no text or logos.", min_length=1)
    warmups: int = Field(default=1, ge=0, le=5)
    runs: int = Field(default=3, ge=1, le=20)
    inference_steps: int = Field(default=20, ge=1, le=50)
    guidance_scale: float = Field(default=4.5, ge=1.0, le=20.0)
    seed: int = Field(default=42, ge=0, le=2_147_483_647)
    compile_transformer: bool = False


def _color_for(source: str) -> str:
    return {"geoapify": "#D85A30", "ticketmaster": "#6E56CF"}.get(source, "#378ADD")


def _hero_prompt(place: dict[str, Any], mood: str) -> str:
    descriptor = place.get("description") or place.get("category") or "local destination"
    return (
        f"A polished editorial hero image inspired by a nearby {descriptor} for someone feeling {mood}. "
        "Warm, inviting composition; no people identifiable as real people; no venue signage, logos, brands, or readable text. "
        "This is a representative illustration, not a factual depiction of a venue or event."
    )


def _map_poi(place: dict[str, Any], image_url: str | None, image_metadata: dict[str, Any] | None) -> dict[str, Any]:
    return {
        "id": place["id"],
        "name": place["name"],
        "category": place["category"],
        "item": place.get("description") or place["category"],
        "lat": place["lat"],
        "lon": place["lon"],
        "distance_m": place["distance_m"],
        "rating": place.get("rating"),
        "review_count": place.get("review_count"),
        "source": place["source"],
        "source_url": place.get("source_url"),
        "directions_url": place["directions_url"],
        "open_now": place.get("open_now"),
        "closes_at": place.get("closes_at"),
        "event_date": place.get("event_date"),
        "event_time": place.get("event_time"),
        "imageUrl": image_url,
        "image_label": image_metadata.get("image_label") if image_metadata else None,
        "color": _color_for(place["source"]),
    }


@app.get("/", include_in_schema=False)
def root() -> RedirectResponse:
    return RedirectResponse(url="/map/")


@app.get("/map/config.js", include_in_schema=False)
def map_config_js() -> Response:
    """Serve the active map payload uncached; it is rewritten after every recommendation."""

    return Response(
        content=MAP_CONFIG_PATH.read_text(encoding="utf-8"),
        media_type="application/javascript",
        headers={"Cache-Control": "no-store, no-cache, must-revalidate"},
    )


@app.get("/api/health")
def health() -> dict[str, Any]:
    return {
        "geoapify_configured": bool(os.getenv("GEOAPIFY_API_KEY")),
        "ticketmaster_configured": bool(os.getenv("TICKETMASTER_API_KEY")),
        "pixart": runtime_info(),
    }


@app.get("/api/schema/user-context")
@app.get("/schema/user_context", include_in_schema=False)
def user_context_schema() -> dict[str, Any]:
    """Expose the stable mood/location/preferences JSON Schema."""

    return schema_document()


# Recommendation jobs run in a background thread (generate_hero_image is a long,
# blocking, per-POI GPU call) so the client can poll for progress -- e.g. "image 3 of
# 8" -- instead of one opaque multi-minute request with no feedback. In-memory registry
# is fine for this single-process demo; jobs are never pruned, which is an acceptable
# tradeoff for a hackathon MVP rather than a long-running service.
_jobs: dict[str, dict[str, Any]] = {}
_jobs_lock = threading.Lock()


def _set_job(job_id: str, **fields: Any) -> None:
    with _jobs_lock:
        _jobs[job_id].update(fields)


def _run_recommendation_job(job_id: str, request: RecommendationRequest) -> None:
    try:
        _set_job(job_id, phase="discovering")
        mood, discovered, provider_notes = discover_nearby(
            mood_text=request.mood,
            latitude=request.latitude,
            longitude=request.longitude,
            radius_m=request.radius_m,
            per_source_limit=max(request.max_results * 3, 12),
        )
    except DiscoveryError as exc:
        _set_job(job_id, status="error", detail=str(exc))
        return

    selected = discovered[: request.max_results]
    generated_images: dict[str, tuple[str, dict[str, Any]]] = {}
    if request.generate_images:
        to_generate = selected[: request.image_limit]
        _set_job(job_id, phase="generating_images", current=0, total=len(to_generate))
        try:
            for index, place in enumerate(to_generate):
                image_path = image_path_for(place["id"])
                metadata = generate_hero_image(
                    prompt=_hero_prompt(place, mood["label"]),
                    output_path=image_path,
                    seed=request.seed + index,
                    inference_steps=request.inference_steps,
                    guidance_scale=request.guidance_scale,
                    compile_transformer=request.compile_transformer,
                )
                generated_images[place["id"]] = (map_image_url(image_path), metadata)
                _set_job(job_id, current=index + 1)
        except (RocmUnavailableError, ValueError) as exc:
            _set_job(job_id, status="error", detail=f"PixArt generation did not run: {exc}")
            return
        except Exception as exc:
            _set_job(job_id, status="error", detail=f"PixArt generation failed: {exc}")
            return

    map_pois: list[dict[str, Any]] = []
    for place in selected:
        generated = generated_images.get(place["id"])
        map_pois.append(
            _map_poi(
                place,
                generated[0] if generated else place.get("provider_image_url"),
                generated[1] if generated else None,
            )
        )

    _set_job(job_id, phase="ground_map")
    ground_map_url: str | None = None
    ground_map_error: str | None = None
    # One retry: this is a single external network call for a demo-critical visual, and
    # a transient failure here previously fell back to the procedural ground texture
    # with zero indication anything went wrong (see provider_notes surfacing below).
    for attempt in range(2):
        try:
            ground_size_m = max(200, request.radius_m * 2)
            ground_map_bytes = fetch_ground_map_image(
                latitude=request.latitude,
                longitude=request.longitude,
                extent_m=ground_size_m,
            )
            write_ground_map_image(ground_map_bytes)
            cache_bust = int(datetime.now(timezone.utc).timestamp())
            ground_map_url = f"{map_image_url(GROUND_MAP_PATH)}?t={cache_bust}"
            ground_map_error = None
            break
        except DiscoveryError as exc:
            ground_map_error = str(exc)
    if ground_map_error:
        # Best-effort: the billboard map still works with its procedural ground texture.
        provider_notes.append(f"Ground map image: {ground_map_error}")

    _set_job(job_id, phase="finalizing")
    primary = selected[0]
    message = recommendation_message(primary, mood["label"])
    write_active_map(
        latitude=request.latitude,
        longitude=request.longitude,
        radius_m=request.radius_m,
        mood=mood["label"],
        mood_text=request.mood,
        recommendation=message,
        pois=map_pois,
        ground_map_url=ground_map_url,
    )
    result = {
        "mood": mood["label"],
        "origin": {"lat": request.latitude, "lon": request.longitude},
        "radius_m": request.radius_m,
        "provider_notes": provider_notes,
        "recommendation": message,
        "primary_poi": map_pois[0],
        "pois": map_pois,
        "map_url": f"/map/?updated={int(datetime.now(timezone.utc).timestamp())}",
        "generated_images": {poi_id: metadata for poi_id, (_, metadata) in generated_images.items()},
    }
    _set_job(job_id, status="done", result=result)


@app.post("/api/recommendations", status_code=202)
def create_recommendations(request: RecommendationRequest) -> dict[str, Any]:
    try:
        ZoneInfo(request.timezone_name)
    except Exception as exc:
        raise HTTPException(status_code=422, detail="timezone_name must be a valid IANA time zone.") from exc

    job_id = uuid4().hex
    with _jobs_lock:
        _jobs[job_id] = {"status": "running", "phase": "queued", "current": 0, "total": 0}
    threading.Thread(target=_run_recommendation_job, args=(job_id, request), daemon=True).start()
    return {"job_id": job_id}


@app.get("/api/recommendations/{job_id}")
def get_recommendation_job(job_id: str) -> dict[str, Any]:
    with _jobs_lock:
        job = _jobs.get(job_id)
        if job is None:
            raise HTTPException(status_code=404, detail="Unknown job_id.")
        return dict(job)


@app.post("/api/benchmarks/pixart")
def run_pixart_benchmark(request: PixArtBenchmarkRequest) -> dict[str, Any]:
    try:
        report = benchmark_pixart(
            prompt=request.prompt,
            warmups=request.warmups,
            runs=request.runs,
            seed=request.seed,
            inference_steps=request.inference_steps,
            guidance_scale=request.guidance_scale,
            compile_transformer=request.compile_transformer,
        )
    except (RocmUnavailableError, ValueError) as exc:
        raise HTTPException(status_code=503, detail=str(exc)) from exc

    BENCHMARKS_ROOT.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    report_path = BENCHMARKS_ROOT / f"pixart-rocm-{timestamp}.json"
    report_path.write_text(json.dumps(report, indent=2), encoding="utf-8")
    report["report_path"] = str(report_path.relative_to(REPO_ROOT))
    return report


@app.middleware("http")
async def _no_cache_map_scripts(request: Request, call_next):
    """Force revalidation for /map's HTML/JS/CSS so an open browser tab can't keep running
    a stale app.js against a changed API response shape (this caused a POST-then-navigate
    to `/map/undefined` when a cached pre-background-job app.js read a field that no
    longer exists on the response). Images and config.js (which sets its own no-store
    header) are excluded -- they're either content-addressed or already handled.
    """

    response = await call_next(request)
    path = request.url.path
    if path.startswith("/map/") and not path.startswith("/map/images/") and path != "/map/config.js":
        response.headers["Cache-Control"] = "no-cache"
    return response


MAP_ROOT.mkdir(parents=True, exist_ok=True)
MAP_IMAGES_ROOT.mkdir(parents=True, exist_ok=True)
ARTIFACTS_ROOT.mkdir(parents=True, exist_ok=True)
app.mount("/map", StaticFiles(directory=MAP_ROOT, html=True), name="map")
app.mount("/artifacts", StaticFiles(directory=ARTIFACTS_ROOT), name="artifacts")