main.py

import asyncio
import json
import os
import threading
import uuid
from datetime import datetime, timedelta, timezone
from functools import lru_cache
from pathlib import Path

import aiofiles
from fastapi import FastAPI, File, Form, HTTPException, Query, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse

from cleanup import start_cleanup_scheduler
from models import JobState, JobStatus, UploadResponse, YoutubeJobRequest, YoutubeJobStatus
from processor import enqueue_job, get_job_status, init_queue, start_worker

BASE_DIR = Path(__file__).parent
UPLOAD_DIR = BASE_DIR / "tmp" / "uploads"
OUTPUT_DIR = BASE_DIR / "tmp" / "outputs"

# Prefer /data (HF Spaces persistent storage) so stats survive redeployments.
_HF_PERSISTENT = Path("/data")
if _HF_PERSISTENT.exists():
    try:
        _HF_PERSISTENT.mkdir(parents=True, exist_ok=True)
        (_HF_PERSISTENT / ".write_test").write_text("ok")
        (_HF_PERSISTENT / ".write_test").unlink()
        DATA_DIR = _HF_PERSISTENT
    except OSError:
        DATA_DIR = BASE_DIR
else:
    DATA_DIR = BASE_DIR

STATS_FILE = DATA_DIR / "stats.json"
DOWNLOAD_LOG_FILE = DATA_DIR / "download_log.json"
_stats_lock = threading.Lock()
_log_lock = threading.Lock()


# Baseline stats that survive redeployments (HF Spaces ephemeral filesystem).
_BASELINE_STATS = {"total_visits": 100, "total_downloads": 5}


def _read_stats() -> dict:
    if not STATS_FILE.exists():
        return dict(_BASELINE_STATS)
    try:
        with open(STATS_FILE, "r") as f:
            return json.load(f)
    except (json.JSONDecodeError, OSError):
        return dict(_BASELINE_STATS)


def _write_stats(stats: dict) -> None:
    with open(STATS_FILE, "w") as f:
        json.dump(stats, f)


def _increment_stat(key: str) -> None:
    with _stats_lock:
        stats = _read_stats()
        stats[key] = stats.get(key, 0) + 1
        _write_stats(stats)


def _read_download_log() -> list[dict]:
    if not DOWNLOAD_LOG_FILE.exists():
        return []
    try:
        with open(DOWNLOAD_LOG_FILE, "r") as f:
            return json.load(f)
    except (json.JSONDecodeError, OSError):
        return []


def _append_download_record(job_id: str, original_filename: str, stem: str) -> None:
    with _log_lock:
        log = _read_download_log()
        log.append({
            "job_id": job_id,
            "original_filename": original_filename,
            "stem": stem,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        })
        with open(DOWNLOAD_LOG_FILE, "w") as f:
            json.dump(log, f)


ALLOWED_CONTENT_TYPES = {"audio/mpeg", "audio/mp3"}
MAX_FILE_SIZE = 500 * 1024 * 1024

# Pitch/tempo shifting limits. Kept conservative so artifacts stay acceptable with the
# stock ffmpeg filter chain (no librubberband available in the slim image).
PITCH_SEMITONES_MIN = -12.0
PITCH_SEMITONES_MAX = 12.0
TEMPO_PCT_MIN = 50.0
TEMPO_PCT_MAX = 150.0


def _pitch_tempo_suffix(pitch_semitones: float, tempo_pct: float) -> str:
    """Suffix used to cache a transformed stem on disk. Empty when no transform."""
    if abs(pitch_semitones) < 1e-3 and abs(tempo_pct - 100.0) < 1e-3:
        return ""
    # Round so near-identical slider positions share a cache file.
    p_int = int(round(pitch_semitones * 10))
    t_int = int(round(tempo_pct * 10))
    return f"__p{p_int}_t{t_int}"


async def _apply_pitch_tempo(
    input_file: Path,
    output_file: Path,
    pitch_semitones: float,
    tempo_pct: float,
) -> None:
    """Transcode input into output with pitch/tempo adjustments via ffmpeg.

    Pitch shift uses the asetrate + aresample trick so we don't need librubberband.
    Tempo scaling uses atempo; the chain stays inside atempo's [0.5, 2.0] sweet spot
    because PITCH_SEMITONES_MAX and TEMPO_PCT_MIN/MAX are capped accordingly.
    """
    filters: list[str] = []
    # Normalize input rate first so asetrate math doesn't depend on source sample rate.
    filters.append("aresample=44100")

    pitch_ratio = 2 ** (pitch_semitones / 12.0)
    if abs(pitch_semitones) > 1e-3:
        filters.append(f"asetrate=44100*{pitch_ratio:.6f}")
        filters.append("aresample=44100")
        filters.append(f"atempo={1 / pitch_ratio:.6f}")

    if abs(tempo_pct - 100.0) > 1e-3:
        tempo_ratio = tempo_pct / 100.0
        filters.append(f"atempo={tempo_ratio:.6f}")

    cmd = [
        "ffmpeg",
        "-y",
        "-hide_banner",
        "-loglevel",
        "error",
        "-i",
        str(input_file),
        "-filter:a",
        ",".join(filters),
        "-codec:a",
        "libmp3lame",
        "-b:a",
        "320k",
        str(output_file),
    ]
    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    _, stderr = await proc.communicate()
    if proc.returncode != 0:
        output_file.unlink(missing_ok=True)
        raise RuntimeError(stderr.decode(errors="replace")[:500] or "ffmpeg transform failed")

app = FastAPI(title="Karaoke Maker API")

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


@app.on_event("startup")
async def startup():
    UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    start_cleanup_scheduler(UPLOAD_DIR, OUTPUT_DIR)
    init_queue()
    await start_worker()


@app.on_event("shutdown")
async def shutdown():
    from cleanup import stop_cleanup_scheduler

    stop_cleanup_scheduler()


@app.get("/health")
async def health():
    return {"status": "ok"}


def _read_stats_locked() -> dict:
    with _stats_lock:
        return _read_stats()


@app.get("/stats")
async def get_stats():
    return await asyncio.to_thread(_read_stats_locked)


@app.post("/stats/visit", status_code=204)
async def record_visit():
    await asyncio.to_thread(_increment_stat, "total_visits")


@app.get("/status/{job_id}", response_model=JobStatus)
async def status(job_id: str):
    job = get_job_status(job_id)
    if job is None:
        raise HTTPException(status_code=404, detail="Job not found")
    return job


@app.post("/upload", response_model=UploadResponse)
async def upload(
    file: UploadFile = File(...),
    trim_start_sec: float | None = Form(default=None),
    trim_end_sec: float | None = Form(default=None),
):
    if file.content_type not in ALLOWED_CONTENT_TYPES:
        raise HTTPException(status_code=400, detail="Only MP3 files are accepted")

    if trim_start_sec is not None and trim_start_sec < 0:
        raise HTTPException(status_code=400, detail="trim_start_sec must be >= 0")
    if trim_end_sec is not None and trim_end_sec <= 0:
        raise HTTPException(status_code=400, detail="trim_end_sec must be > 0")
    if trim_start_sec is not None and trim_end_sec is not None and trim_end_sec <= trim_start_sec:
        raise HTTPException(status_code=400, detail="trim_end_sec must be greater than trim_start_sec")

    content = await file.read()
    if len(content) > MAX_FILE_SIZE:
        raise HTTPException(status_code=400, detail="File too large (max 500MB)")

    is_mpeg_sync = len(content) >= 2 and content[0] == 0xFF and (content[1] & 0xE0) == 0xE0
    is_id3 = len(content) >= 3 and content[:3] == b"ID3"
    if not is_mpeg_sync and not is_id3:
        raise HTTPException(status_code=400, detail="File does not appear to be a valid MP3")

    job_id = str(uuid.uuid4())
    input_path = UPLOAD_DIR / f"{job_id}.mp3"
    job_output_dir = OUTPUT_DIR / job_id
    original_output_path = job_output_dir / "original.mp3"

    try:
        job_output_dir.mkdir(parents=True, exist_ok=True)
        async with aiofiles.open(input_path, "wb") as f:
            await f.write(content)
        async with aiofiles.open(original_output_path, "wb") as f:
            await f.write(content)
    except OSError as e:
        input_path.unlink(missing_ok=True)
        original_output_path.unlink(missing_ok=True)
        raise HTTPException(status_code=503, detail="Storage unavailable, please try again") from e

    try:
        await enqueue_job(
            job_id,
            str(input_path),
            str(job_output_dir),
            trim_start_sec=trim_start_sec,
            trim_end_sec=trim_end_sec,
            original_filename=file.filename or "unknown.mp3",
        )
    except RuntimeError as e:
        input_path.unlink(missing_ok=True)
        original_output_path.unlink(missing_ok=True)
        raise HTTPException(status_code=503, detail=str(e)) from e

    return UploadResponse(job_id=job_id)


def _resolve_download_path(job_id: str, stem_name: str | None) -> tuple[Path, str]:
    job = get_job_status(job_id)
    if job is None:
        raise HTTPException(status_code=404, detail="Job not found")
    if job.state != JobState.done:
        raise HTTPException(status_code=409, detail="Job not complete yet")

    # Provide source track for before/after comparison clips.
    if stem_name == "original":
        # Prefer output-dir copy so availability matches generated stems.
        original_path = OUTPUT_DIR / job_id / "original.mp3"
        if not original_path.exists():
            # Backward compatibility for jobs created before original.mp3 copy.
            original_path = UPLOAD_DIR / f"{job_id}.mp3"
        if not original_path.exists():
            raise HTTPException(status_code=404, detail="Original file not found")
        return original_path, "original"

    if not job.stems:
        raise HTTPException(status_code=404, detail="No stems available for this job")

    resolved_stem = stem_name or "instrumental"
    if resolved_stem not in job.stems:
        if stem_name is None:
            resolved_stem = next(iter(job.stems.keys()))
        else:
            raise HTTPException(status_code=404, detail="Stem not found")

    filename = job.stems[resolved_stem]
    file_path = OUTPUT_DIR / job_id / filename
    if not file_path.exists():
        raise HTTPException(status_code=404, detail="Output file not found")

    return file_path, resolved_stem


def _validate_pitch_tempo(pitch_semitones: float, tempo_pct: float) -> None:
    if not PITCH_SEMITONES_MIN <= pitch_semitones <= PITCH_SEMITONES_MAX:
        raise HTTPException(
            status_code=400,
            detail=f"pitch_semitones must be between {PITCH_SEMITONES_MIN} and {PITCH_SEMITONES_MAX}",
        )
    if not TEMPO_PCT_MIN <= tempo_pct <= TEMPO_PCT_MAX:
        raise HTTPException(
            status_code=400,
            detail=f"tempo_pct must be between {TEMPO_PCT_MIN} and {TEMPO_PCT_MAX}",
        )


async def _prepare_download_file(
    source_path: Path,
    job_id: str,
    resolved_stem: str,
    pitch_semitones: float,
    tempo_pct: float,
) -> Path:
    suffix = _pitch_tempo_suffix(pitch_semitones, tempo_pct)
    if not suffix:
        return source_path

    transformed_dir = OUTPUT_DIR / job_id
    transformed_dir.mkdir(parents=True, exist_ok=True)
    transformed_path = transformed_dir / f"{resolved_stem}{suffix}.mp3"
    if transformed_path.exists():
        return transformed_path

    try:
        await _apply_pitch_tempo(
            input_file=source_path,
            output_file=transformed_path,
            pitch_semitones=pitch_semitones,
            tempo_pct=tempo_pct,
        )
    except RuntimeError as e:
        raise HTTPException(status_code=500, detail=f"Audio transform failed: {e}") from e
    return transformed_path


def _download_filename(resolved_stem: str, pitch_semitones: float, tempo_pct: float) -> str:
    if not _pitch_tempo_suffix(pitch_semitones, tempo_pct):
        return f"{resolved_stem}.mp3"
    parts = [resolved_stem]
    if abs(pitch_semitones) >= 1e-3:
        sign = "+" if pitch_semitones > 0 else ""
        parts.append(f"{sign}{pitch_semitones:g}st")
    if abs(tempo_pct - 100.0) >= 1e-3:
        parts.append(f"{tempo_pct:g}pct")
    return "_".join(parts) + ".mp3"


@app.get("/download/{job_id}")
async def download_default(
    job_id: str,
    pitch_semitones: float = Query(default=0.0),
    tempo_pct: float = Query(default=100.0),
):
    _validate_pitch_tempo(pitch_semitones, tempo_pct)
    source_path, resolved_stem = _resolve_download_path(job_id, None)
    file_path = await _prepare_download_file(
        source_path, job_id, resolved_stem, pitch_semitones, tempo_pct
    )
    job = get_job_status(job_id)
    original_name = (job.original_filename if job else None) or "unknown.mp3"
    await asyncio.to_thread(_increment_stat, "total_downloads")
    await asyncio.to_thread(_append_download_record, job_id, original_name, resolved_stem)
    return FileResponse(
        path=file_path,
        media_type="audio/mpeg",
        filename=_download_filename(resolved_stem, pitch_semitones, tempo_pct),
    )


@app.get("/download/{job_id}/{stem_name}")
async def download_stem(
    job_id: str,
    stem_name: str,
    pitch_semitones: float = Query(default=0.0),
    tempo_pct: float = Query(default=100.0),
):
    _validate_pitch_tempo(pitch_semitones, tempo_pct)
    source_path, resolved_stem = _resolve_download_path(job_id, stem_name)
    file_path = await _prepare_download_file(
        source_path, job_id, resolved_stem, pitch_semitones, tempo_pct
    )
    job = get_job_status(job_id)
    original_name = (job.original_filename if job else None) or "unknown.mp3"
    await asyncio.to_thread(_increment_stat, "total_downloads")
    await asyncio.to_thread(_append_download_record, job_id, original_name, resolved_stem)
    return FileResponse(
        path=file_path,
        media_type="audio/mpeg",
        filename=_download_filename(resolved_stem, pitch_semitones, tempo_pct),
    )


@app.get("/admin/downloads")
async def get_download_log():
    log = await asyncio.to_thread(_read_download_log)
    return log


# ---------------------------------------------------------------------------
# YouTube karaoke pipeline endpoints
# ---------------------------------------------------------------------------

from supabase import create_client  # noqa: E402
from db import SongRepo  # noqa: E402
from storage import SupabaseStorage  # noqa: E402
from youtube import extract_youtube_id, InvalidYoutubeUrl, fetch_metadata, download_audio  # noqa: E402
from separate import separate_vocals_to_mp3  # noqa: E402
from youtube_pipeline import PipelineDeps, run_pipeline  # noqa: E402


@lru_cache(maxsize=1)
def _youtube_deps() -> PipelineDeps:
    sb = create_client(os.environ["SUPABASE_URL"], os.environ["SUPABASE_SERVICE_ROLE_KEY"])
    repo = SongRepo(sb)
    storage = SupabaseStorage(
        client=sb,
        bucket=os.environ.get("STORAGE_BUCKET", "karaoke"),
        public_base_url=os.environ["STORAGE_PUBLIC_BASE_URL"],
    )
    return PipelineDeps(
        repo=repo,
        storage=storage,
        fetch_metadata=fetch_metadata,
        download_audio=download_audio,
        separate_vocals=separate_vocals_to_mp3,
        work_dir=Path(BASE_DIR) / "tmp" / "youtube",
        max_duration_sec=int(os.environ.get("YT_MAX_DURATION_SEC", "600")),
    )


@app.post("/jobs/youtube", response_model=YoutubeJobStatus)
async def create_youtube_job(req: YoutubeJobRequest):
    try:
        youtube_id = extract_youtube_id(req.url)
    except InvalidYoutubeUrl:
        raise HTTPException(status_code=400, detail="invalid youtube url")

    deps = _youtube_deps()
    job_id = deps.repo.create_job(youtube_id)

    # Fire-and-forget via asyncio task; pipeline writes status into Supabase.
    asyncio.create_task(asyncio.to_thread(run_pipeline, req.url, job_id, deps))
    return YoutubeJobStatus(id=job_id, status="queued")


@app.get("/jobs/youtube/{job_id}", response_model=YoutubeJobStatus)
async def get_youtube_job(job_id: str):
    deps = _youtube_deps()
    row = deps.repo.get_job(job_id)
    if not row:
        raise HTTPException(status_code=404, detail="job not found")
    slug = None
    if row["status"] == "ready":
        song = deps.repo.get_song(row["youtube_id"])
        if song:
            slug = song.get("slug")
    return YoutubeJobStatus(
        id=row["id"], status=row["status"], error=row.get("error"), slug=slug,
    )


@app.get("/admin/youtube-stats")
def youtube_stats():
    deps = _youtube_deps()
    sb = deps.repo._c
    since = (datetime.now(timezone.utc) - timedelta(hours=24)).isoformat()
    res = sb.table("jobs").select("status").gte("created_at", since).execute()
    counts: dict[str, int] = {}
    for row in res.data or []:
        counts[row["status"]] = counts.get(row["status"], 0) + 1
    total = sum(counts.values()) or 1
    return {
        "window": "24h",
        "total": total,
        "counts": counts,
        "failed_fetch_rate": counts.get("failed_fetch", 0) / total,
    }