app/main.py

import os
import uuid
from collections import Counter
from pathlib import Path

# Load .env from backend root if present
_env_path = Path(__file__).parent.parent / ".env"
if _env_path.exists():
    for _line in _env_path.read_text().splitlines():
        if _line.strip() and not _line.startswith("#") and "=" in _line:
            _k, _v = _line.split("=", 1)
            os.environ.setdefault(_k.strip(), _v.strip())

import librosa
import numpy as np
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response
from pydantic import BaseModel

from app.schemas import AnalysisResponse, FretboardPosition, KeySegment
from app.services.audio_loader import chunk_audio, load_audio_bytes
from app.services.boundary_detector import find_boundaries
from app.services.chromagram import extract_chroma_mean
from app.services.key_detector import detect_key
from app.services.mert_encoder import MERTEncoder
from app.services.music_theory import get_fretboard_positions, get_pentatonic_notes, get_scale_notes
from app.services.smoother import anchor_to_global_key, merge_segments, smooth_transitions

app = FastAPI(title="KeyShift API")
_origins = ["http://localhost:3000", "http://localhost:3001"]
if _frontend := os.environ.get("FRONTEND_URL"):
    _origins.append(_frontend)
app.add_middleware(
    CORSMiddleware,
    allow_origins=_origins,
    allow_methods=["POST", "GET"],
    allow_headers=["*"],
)

_encoder: MERTEncoder | None = None
_audio_store: dict[str, tuple[bytes, str]] = {}  # token -> (audio_bytes, content_type)

AUDIO_MIME_TYPES = {
    "audio/wav","audio/wave","audio/mpeg","audio/mp3",
    "audio/mp4","audio/ogg","audio/flac","audio/x-wav",
}

CHUNK_DURATION    = 10.0
OVERLAP           = 0.5
HOP_DURATION      = CHUNK_DURATION * (1 - OVERLAP)  # 5.0 s
MAX_UPLOAD_BYTES  = 50 * 1024 * 1024  # 50 MB
MAX_SECTION_SECS  = 30.0
MAX_MERGE_SECS    = 45.0


def get_encoder() -> MERTEncoder:
    global _encoder
    if _encoder is None:
        _encoder = MERTEncoder()
    return _encoder


def _run_pipeline(y: np.ndarray, sr: int, duration: float) -> AnalysisResponse:
    chunks           = chunk_audio(y, sr, CHUNK_DURATION, OVERLAP)
    embeddings       = get_encoder().encode_batch(chunks, sr)
    boundary_indices = find_boundaries(embeddings, hop_duration=HOP_DURATION)

    section_breaks = [0] + boundary_indices + [len(chunks)]
    coarse_sections = [
        chunks[section_breaks[i]: section_breaks[i + 1]]
        for i in range(len(section_breaks) - 1)
        if section_breaks[i] < section_breaks[i + 1]
    ]

    sections = []
    for sec in coarse_sections:
        dur = sec[-1]["end"] - sec[0]["start"]
        if dur <= MAX_SECTION_SECS:
            sections.append(sec)
        else:
            n = max(2, round(dur / MAX_SECTION_SECS))
            size = max(1, len(sec) // n)
            for i in range(0, len(sec), size):
                sub = sec[i: i + size]
                if sub:
                    sections.append(sub)

    raw_detections = []
    for section in sections:
        mean_chroma = np.stack([extract_chroma_mean(c["y"], sr) for c in section]).mean(axis=0)
        root, mode, confidence = detect_key(mean_chroma)
        raw_detections.append({
            "start": section[0]["start"],
            "end": section[-1]["end"],
            "key": root, "mode": mode, "confidence": confidence,
        })

    # Global key from full-song chroma — chord noise averages out, more reliable
    global_chroma = extract_chroma_mean(y, sr)
    global_key, global_mode, global_conf = detect_key(global_chroma)

    final_segments = merge_segments(smooth_transitions(raw_detections, window_size=5), max_duration=MAX_MERGE_SECS)

    # Anchor ambiguous segments: prevents C major songs showing G/F major on V/IV sections
    if global_conf > 0.72:
        final_segments = anchor_to_global_key(final_segments, global_key, global_mode)

    dominant = [global_key, global_mode]

    result = []
    for seg in final_segments:
        scale_notes      = get_scale_notes(seg["key"], seg["mode"])
        pentatonic_notes = get_pentatonic_notes(seg["key"], seg["mode"])
        fret_positions   = get_fretboard_positions(scale_notes)
        result.append(KeySegment(
            **{k: seg[k] for k in ("start","end","key","mode","confidence")},
            scale_notes=scale_notes,
            pentatonic_notes=pentatonic_notes,
            fretboard_positions=[FretboardPosition(**p) for p in fret_positions],
        ))

    # BPM detection
    tempo_arr, _ = librosa.beat.beat_track(y=y, sr=sr)
    bpm = float(np.atleast_1d(tempo_arr)[0])

    # Waveform RMS (~150 points, normalized 0–1)
    N_WF = 150
    hop_w = max(1, len(y) // N_WF)
    raw_rms = [
        float(np.sqrt(np.mean(y[i: i + hop_w] ** 2)))
        for i in range(0, len(y) - hop_w + 1, hop_w)
    ][:N_WF]
    max_rms = max(raw_rms) if raw_rms else 1.0
    waveform = [v / max_rms for v in raw_rms] if max_rms > 1e-8 else [0.0] * len(raw_rms)

    return AnalysisResponse(
        duration=duration, segments=result,
        dominant_key=dominant[0], dominant_mode=dominant[1],
        bpm=bpm, waveform=waveform,
    )


@app.get("/audio/{token}")
async def serve_audio(token: str) -> Response:
    if token not in _audio_store:
        raise HTTPException(status_code=404, detail="Audio not found")
    data, content_type = _audio_store[token]
    return Response(content=data, media_type=content_type, headers={
        "Accept-Ranges": "bytes",
        "Cache-Control": "no-store",
    })


@app.post("/analyze", response_model=AnalysisResponse)
async def analyze(file: UploadFile = File(...)) -> AnalysisResponse:
    if file.content_type not in AUDIO_MIME_TYPES:
        raise HTTPException(status_code=422, detail=f"Unsupported type: {file.content_type}")

    raw = await file.read()
    if len(raw) > MAX_UPLOAD_BYTES:
        raise HTTPException(status_code=413, detail="File too large. Maximum size is 50 MB.")

    token = str(uuid.uuid4())
    _audio_store[token] = (raw, file.content_type or "audio/mpeg")

    y, sr, duration = load_audio_bytes(raw)
    result = _run_pipeline(y, sr, duration)
    return AnalysisResponse(
        duration=result.duration, segments=result.segments,
        dominant_key=result.dominant_key, dominant_mode=result.dominant_mode,
        audio_token=token, bpm=result.bpm, waveform=result.waveform,
    )


class UrlRequest(BaseModel):

    url: str


@app.post("/analyze-url", response_model=AnalysisResponse)
async def analyze_url(body: UrlRequest) -> AnalysisResponse:
    import tempfile
    import shutil
    import warnings
    import yt_dlp

    with tempfile.TemporaryDirectory() as tmpdir:
        def _progress(d: dict) -> None:
            status = d.get("status", "")
            if status == "downloading":
                print(
                    f"[ydl] {d.get('_percent_str','?'):>7} of {d.get('_total_bytes_str','?')} "
                    f"at {d.get('_speed_str','?')}  ETA {d.get('_eta_str','?')}",
                    flush=True,
                )
            elif status == "finished":
                print(f"[ydl] finished → {d.get('filename','?')}", flush=True)
            elif status == "error":
                print(f"[ydl] error: {d}", flush=True)

        use_ffmpeg = shutil.which("ffmpeg") is not None
        ydl_opts: dict = {
            "format": "bestaudio[ext=webm]/bestaudio/best",
            "outtmpl": os.path.join(tmpdir, "audio.%(ext)s"),
            "quiet": True,
            "no_warnings": False,
            "verbose": True,
            "noplaylist": True,
            "progress_hooks": [_progress],
            "extractor_args": {"youtube": {"player_client": ["ios", "web"]}},
        }
        if use_ffmpeg:
            ydl_opts["postprocessors"] = [{
                "key": "FFmpegExtractAudio",
                "preferredcodec": "wav",
            }]
        title: str | None = None
        try:
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                info = ydl.extract_info(body.url, download=True)
                title = info.get("title") if info else None
        except Exception as e:
            raise HTTPException(status_code=422, detail=f"Download failed: {e}")

        files = os.listdir(tmpdir)
        if not files:
            raise HTTPException(status_code=422, detail="No audio downloaded")

        audio_path = os.path.join(tmpdir, files[0])
        ext = os.path.splitext(audio_path)[1].lower()
        content_type = {
            ".wav": "audio/wav", ".mp3": "audio/mpeg", ".ogg": "audio/ogg",
            ".webm": "audio/webm", ".m4a": "audio/mp4", ".flac": "audio/flac",
        }.get(ext, "audio/webm")

        try:
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                y, sr = librosa.load(audio_path, sr=22050, mono=True)
        except Exception as e:
            raise HTTPException(status_code=422, detail=f"Could not decode audio: {e}")
        duration = float(len(y) / sr)

        with open(audio_path, "rb") as f:
            audio_bytes = f.read()

    token = str(uuid.uuid4())
    _audio_store[token] = (audio_bytes, content_type)

    result = _run_pipeline(y, sr, duration)
    return AnalysisResponse(
        duration=result.duration, segments=result.segments,
        dominant_key=result.dominant_key, dominant_mode=result.dominant_mode,
        audio_token=token, title=title, bpm=result.bpm, waveform=result.waveform,
    )