app.py

from __future__ import annotations

import json
import os
import secrets
import subprocess
import sys
import threading
import uuid
from datetime import UTC, datetime
from pathlib import Path
from typing import Annotated, Any

import memray
from fastapi import Depends, FastAPI, Header, HTTPException, status
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field

from workload import build_customer_ltv_report

OUTPUT_DIR = Path(os.getenv("MEMRAY_OUTPUT_DIR", "./memray-results")).resolve()
# Memray only allows one active tracker per process, so request-level captures need a gate.
PROFILE_LOCK = threading.Lock()
MAX_ROWS = int(os.getenv("MEMRAY_MAX_ROWS", "750000"))


class ProfileRequest(BaseModel):
    rows: int = Field(default=250_000, ge=50_000, le=2_000_000)
    customers: int = Field(default=25_000, ge=1_000, le=250_000)
    top_n: int = Field(default=10, ge=3, le=25)
    seed: int = Field(default=7, ge=0, le=10_000)


class ProfileResult(BaseModel):
    run_id: str
    created_at: str
    duration_seconds: float
    capture_file: str
    flamegraph_file: str | None = None
    flamegraph_error: str | None = None
    summary: dict[str, Any]


app = FastAPI(
    title="Render Memray Example",
    description=(
        "Profile a NumPy/pandas-heavy workload with Memray native traces and "
        "persist the results on Render."
    ),
)


def ensure_profile_token(
    x_profile_token: Annotated[str | None, Header()] = None,
) -> None:
    expected = os.getenv("MEMRAY_PROFILE_TOKEN")
    if not expected:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="Set MEMRAY_PROFILE_TOKEN before enabling profiling endpoints.",
        )
    if not x_profile_token or not secrets.compare_digest(x_profile_token, expected):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Missing or invalid X-Profile-Token header.",
        )


def profile_output_paths(run_id: str) -> tuple[Path, Path, Path]:
    return (
        OUTPUT_DIR / f"{run_id}.bin",
        OUTPUT_DIR / f"{run_id}.html",
        OUTPUT_DIR / f"{run_id}.json",
    )


def write_metadata(metadata_path: Path, payload: dict[str, Any]) -> None:
    metadata_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")


def build_flamegraph(capture_path: Path, flamegraph_path: Path) -> str | None:
    try:
        subprocess.run(
            [sys.executable, "-m", "memray", "flamegraph", "-o", str(flamegraph_path), str(capture_path)],
            check=True,
            capture_output=True,
            text=True,
        )
    except FileNotFoundError:
        return "Unable to invoke memray for report generation, so only the capture file was saved."
    except subprocess.CalledProcessError as exc:
        error_output = exc.stderr.strip() or exc.stdout.strip() or str(exc)
        return f"Failed to render flamegraph: {error_output}"
    return None


def latest_profiles(limit: int = 10) -> list[dict[str, Any]]:
    profiles: list[dict[str, Any]] = []
    for metadata_path in sorted(OUTPUT_DIR.glob("*.json"), reverse=True)[:limit]:
        profiles.append(json.loads(metadata_path.read_text(encoding="utf-8")))
    return profiles


def resolve_profile_asset(run_id: str, suffix: str) -> Path:
    if not run_id.replace("-", "").isalnum():
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Unknown profile run.")

    candidate = OUTPUT_DIR / f"{run_id}.{suffix}"
    if not candidate.exists():
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Profile artifact not found.")
    return candidate


@app.on_event("startup")
def create_output_dir() -> None:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)


@app.get("/")
def root() -> dict[str, Any]:
    return {
        "service": "render-memray-example",
        "output_dir": str(OUTPUT_DIR),
        "max_rows": MAX_ROWS,
        "recent_profiles": latest_profiles(limit=3),
    }


@app.post("/profile/native", response_model=ProfileResult, dependencies=[Depends(ensure_profile_token)])
def capture_native_profile(request: ProfileRequest) -> ProfileResult:
    if request.rows > MAX_ROWS:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail=f"rows must be <= MEMRAY_MAX_ROWS ({MAX_ROWS}).",
        )

    if not PROFILE_LOCK.acquire(blocking=False):
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail="A profile is already running. Retry once the current capture finishes.",
        )

    run_id = f"{datetime.now(UTC):%Y%m%dT%H%M%SZ}-{uuid.uuid4().hex[:8]}"
    capture_path, flamegraph_path, metadata_path = profile_output_paths(run_id)
    started_at = datetime.now(UTC)

    try:
        with memray.Tracker(str(capture_path), native_traces=True):
            summary = build_customer_ltv_report(
                rows=request.rows,
                customer_count=request.customers,
                top_n=request.top_n,
                seed=request.seed,
            ).as_dict()
    finally:
        PROFILE_LOCK.release()

    duration = round((datetime.now(UTC) - started_at).total_seconds(), 3)
    flamegraph_error = build_flamegraph(capture_path, flamegraph_path)

    payload = {
        "run_id": run_id,
        "created_at": started_at.isoformat(),
        "duration_seconds": duration,
        "capture_file": str(capture_path),
        "flamegraph_file": None if flamegraph_error else str(flamegraph_path),
        "flamegraph_error": flamegraph_error,
        "summary": summary,
    }
    write_metadata(metadata_path, payload)
    return ProfileResult(**payload)


@app.get("/profiles", dependencies=[Depends(ensure_profile_token)])
def list_profiles() -> list[dict[str, Any]]:
    return latest_profiles(limit=20)


@app.get("/profiles/{run_id}/capture", dependencies=[Depends(ensure_profile_token)])
def download_capture(run_id: str) -> FileResponse:
    asset = resolve_profile_asset(run_id, "bin")
    return FileResponse(asset, media_type="application/octet-stream", filename=asset.name)


@app.get("/profiles/{run_id}/flamegraph", dependencies=[Depends(ensure_profile_token)])
def download_flamegraph(run_id: str) -> FileResponse:
    asset = resolve_profile_asset(run_id, "html")
    return FileResponse(asset, media_type="text/html", filename=asset.name)