"""EUMORA REST API -- FastAPI backend for the frontend."""

import os
import sys
import threading
import time as _time
from collections import OrderedDict
from concurrent.futures import ThreadPoolExecutor
from concurrent.futures import TimeoutError as FutureTimeout
from contextlib import asynccontextmanager, contextmanager
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from dotenv import load_dotenv

load_dotenv()
sys.path.insert(0, str(Path(__file__).parent))

# Global socket timeout covers ALL network calls including Spotify's OAuth token
# requests which bypass spotipy's requests_timeout parameter.
import socket as _socket
_socket.setdefaulttimeout(10)

# CORS: comma-separated origins, or "*" for development
# e.g. ALLOWED_ORIGINS=https://yourapp.vercel.app,https://yourapp.com
_raw_origins = os.environ.get("ALLOWED_ORIGINS", "*")
ALLOWED_ORIGINS = [o.strip() for o in _raw_origins.split(",")] if _raw_origins != "*" else ["*"]


# ---------------------------------------------------------------------------
# Request models
# ---------------------------------------------------------------------------

class PredictRequest(BaseModel):
    text: Optional[str] = None
    # raw=True: plain softmax as the HF router returns it, for the Vercel
    # emotion route's fallback when the router is unavailable.
    raw: bool = False
    reference_url: Optional[str] = None
    target_sarcasm_prior: float = Field(0.15, gt=0.0, lt=1.0)
    disable_prior_adjustment: bool = False

    model_config = {"str_max_length": 2000}  # prevent oversized payloads


class RecommendRequest(BaseModel):
    text: Optional[str] = None
    reference_url: Optional[str] = None
    limit: int = 10
    genre: Optional[str] = None
    blend: bool = True
    max_popularity: int = 65
    target_sarcasm_prior: float = Field(0.15, gt=0.0, lt=1.0)
    disable_prior_adjustment: bool = False


class ScoreRequest(BaseModel):
    title: str
    artist: str
    user_probabilities: dict
    # Optional: a preview the client already has (iTunes-sourced tracks carry
    # one). Saves a search; only allow-listed preview hosts are fetched.
    preview_url: Optional[str] = None


# ---------------------------------------------------------------------------
# App lifecycle -- load heavy models once at startup
# ---------------------------------------------------------------------------

_predictor   = None
_recommender = None

# Acoustic analysis (VGGish + valence/arousal head). Loaded in the background
# so a slow weight load can never delay startup or /health, and optional: with
# it unavailable, /api/score behaves exactly as it did before audio existed.
_acoustic = None
_acoustic_status = "disabled"
# Four-modality VAD fusion (Objective 8): a 32 KB JSON of linear models, loaded
# with the app. When present and approved by its ship rule, it produces the
# match score; otherwise scoring falls back to the 50/50 lyrics + audio mix.
_fusion = None
_audio_cache: "OrderedDict[tuple, dict]" = OrderedDict()
_audio_lock = threading.Lock()
# Light threads that wait on the acoustic worker processes (src/acoustic_worker)
# and fill the cache. The compute runs in niced processes, and only when no
# lyric work is in flight (below): on a Space capped at 2 CPUs, DeBERTa is the
# scoring bottleneck, and the sound must never slow it down.
_audio_pool = ThreadPoolExecutor(max_workers=4, thread_name_prefix="audio")
AUDIO_CACHE_SIZE = 2000
# Audio is waited for only until this long after the request arrived, so it
# can never make a response slower than the lyric path already is. Analysis
# that misses the deadline keeps running, is cached, and the response says
# `audio_pending` so the client can ask again once, cheaply.
AUDIO_BUDGET_S = 7.0
NO_PREVIEW_TTL_S = 6 * 3600    # the catalogue changes slowly
ERROR_TTL_S = 300              # rate limits and timeouts clear quickly

# Audio yields to lyrics. The Space runs under a 2-CPU cgroup quota
# (cpu.max 200000/100000), where process priorities inside the container do
# little, so the API schedules it itself: an analysis starts only when no
# request is in its lyric phase, or after AUDIO_GATE_MAX_S under continuous
# load. Measured with audio held back this way, a burst of 10 new tracks
# returns exactly as fast as lyrics alone. A lone request still gets its sound
# within the budget: its own lyric phase ends in a second or two.
_lyric_active = 0
_lyric_idle = threading.Condition()
AUDIO_GATE_MAX_S = 30.0


@contextmanager
def _lyric_phase():
    global _lyric_active
    with _lyric_idle:
        _lyric_active += 1
    try:
        yield
    finally:
        with _lyric_idle:
            _lyric_active -= 1
            if _lyric_active == 0:
                _lyric_idle.notify_all()


# Lyric distributions per track: up to five DeBERTa passes each, and the same
# for every listener, so a repeat request (or the client's audio retry) costs
# nothing. Tracks with no lyrics are remembered briefly to skip re-fetching.
_lyric_cache: "OrderedDict[tuple, dict]" = OrderedDict()
_lyric_lock = threading.Lock()
LYRIC_CACHE_SIZE = 4000
NO_LYRICS_TTL_S = 600


def _load_acoustic(attempts: int = 4, backoff_s: float = 30.0):
    """Load VGGish and the head, retrying: the weights may need downloading if
    the image was built without them, and a network blip should not leave the
    Space lyric-only until its next restart."""
    global _acoustic, _acoustic_status
    for attempt in range(1, attempts + 1):
        _acoustic_status = "loading"
        try:
            from src.acoustic_worker import ProcessAnalyzer
            _acoustic = ProcessAnalyzer(workers=2, nice=10)
            _acoustic_status = "ready"
            print("Acoustic analyzer ready.")
            return
        except Exception as e:  # never fatal: lyric scoring still works
            _acoustic_status = f"error: {type(e).__name__}: {str(e)[:160]}"
            print(f"Warning: acoustic load attempt {attempt}/{attempts} failed -- {_acoustic_status}")
            if attempt < attempts:
                _time.sleep(backoff_s * attempt)


def _track_key(title: str, artist: str) -> tuple:
    return (title.strip().lower(), artist.strip().lower())


def _lyric_probs_for(title: str, artist: str) -> tuple:
    """(lyric distribution or None, whether lyrics were found, whether they
    read as English), cached."""
    key = _track_key(title, artist)
    now = _time.time()
    with _lyric_lock:
        hit = _lyric_cache.get(key)
        if hit and (hit["probs"] is not None or now < hit["retry_after"]):
            _lyric_cache.move_to_end(key)
            return hit["probs"], hit["found"], hit.get("english", True)

    lyrics = None
    lyrics_fetcher = getattr(_recommender, "lyrics_fetcher", None) if _recommender else None
    if lyrics_fetcher:
        try:
            lyrics_raw = lyrics_fetcher.get_lyrics(title, artist)
            if lyrics_raw:
                lyrics = lyrics_raw[:3000].strip()
        except Exception as e:
            print(f"Lyrics fetch failed for '{title}' by '{artist}': {e}")

    probs, english = None, True
    if lyrics:
        english = _lyrics_language_ok(lyrics)
        try:
            probs = _chunked_lyric_probs(_predictor, lyrics)
        except Exception as e:
            print(f"Chunked scoring failed for '{title}': {e}")
            return None, True, english  # transient: do not cache
    with _lyric_lock:
        _lyric_cache[key] = {"probs": probs, "found": bool(lyrics), "english": english,
                             "retry_after": now + NO_LYRICS_TTL_S}
        while len(_lyric_cache) > LYRIC_CACHE_SIZE:
            _lyric_cache.popitem(last=False)
    return probs, bool(lyrics), english


def _audio_for(title: str, artist: str, preview_url: Optional[str]) -> Optional[dict]:
    """Cached acoustic analysis of a track's preview, or None."""
    key = _track_key(title, artist)
    now = _time.time()
    with _audio_lock:
        hit = _audio_cache.get(key)
        if hit and (hit["value"] is not None or now < hit["retry_after"]):
            _audio_cache.move_to_end(key)
            return hit["value"]
    with _lyric_idle:
        _lyric_idle.wait_for(lambda: _lyric_active == 0, timeout=AUDIO_GATE_MAX_S)
    try:
        value = _acoustic.analyze_track(title, artist, preview_url)
        retry_after = _time.time() + NO_PREVIEW_TTL_S
    except Exception as e:
        print(f"Audio analysis failed for '{title}' by '{artist}': {type(e).__name__}: {e}")
        value, retry_after = None, _time.time() + ERROR_TTL_S
    with _audio_lock:
        _audio_cache[key] = {"value": value, "retry_after": retry_after}
        while len(_audio_cache) > AUDIO_CACHE_SIZE:
            _audio_cache.popitem(last=False)
    return value


@asynccontextmanager
async def lifespan(app: FastAPI):
    global _predictor, _recommender

    from src.predict import EmotionPredictor
    from src.spotify import SpotifyRecommender

    print("Loading emotion model...")
    _predictor = EmotionPredictor(enable_viz=False)
    print("Emotion model ready.")

    try:
        _recommender = SpotifyRecommender()
        print("Spotify recommender ready.")
    except (ValueError, ImportError) as e:
        print(f"Warning: Spotify disabled -- {e}")
        _recommender = None

    if os.environ.get("EUMORA_AUDIO", "1") != "0":
        threading.Thread(target=_load_acoustic, name="acoustic-load", daemon=True).start()

    global _fusion
    try:
        from src.fusion import FUSION_PATH, FusionModel
        if FUSION_PATH.exists() and os.environ.get("EUMORA_FUSION", "1") != "0":
            candidate = FusionModel.load()
            _fusion = candidate if candidate.spec.get("ship_as_score") else None
            print(f"VAD fusion {'active' if _fusion else 'present but not approved to score'}.")
    except Exception as e:
        print(f"Warning: VAD fusion disabled -- {type(e).__name__}: {e}")
        _fusion = None

    yield

    # Stop the acoustic worker processes with the server; left running, they
    # outlive it as orphans.
    pool = getattr(_acoustic, "pool", None)
    if pool is not None:
        pool.shutdown(wait=False, cancel_futures=True)
    _predictor   = None
    _recommender = None


app = FastAPI(title="EUMORA API", version="1.0.0", lifespan=lifespan)

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


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------

def _cpu_info() -> dict:
    """What the container can actually use; os.cpu_count() reports the host."""
    info = {"cpu_count": os.cpu_count()}
    try:
        info["affinity"] = len(os.sched_getaffinity(0))
    except AttributeError:
        pass
    try:
        info["cgroup_cpu_max"] = Path("/sys/fs/cgroup/cpu.max").read_text().strip()
    except OSError:
        pass
    try:
        import torch
        info["torch_threads"] = torch.get_num_threads()
    except Exception:
        pass
    return info


@app.get("/")
def root():
    """What this service is, for anyone who opens the Space page itself."""
    return {
        "service": "EUMORA API",
        "app": "https://eumora.vercel.app",
        "status": "ok" if _predictor else "starting",
        "endpoints": {
            "GET /health": "model, audio and fusion status",
            "POST /api/score": "match a track to a listener's emotion (lyrics + sound, calibrated VAD)",
            "POST /api/predict": "emotion distribution for a text (raw=true: plain softmax)",
        },
        "acoustic": _acoustic_status,
        "fusion": "active" if _fusion is not None else "off",
    }


@app.get("/health")
def health():
    import time as _t
    genius_ok = (
        _recommender is not None
        and _recommender.lyrics_fetcher is not None
        and _recommender.lyrics_fetcher.genius is not None
    )
    on_hf = bool(os.environ.get("SPACE_ID") or os.environ.get("SPACES_ZERO_GPU"))

    # Quick DeBERTa warmup timing (first call cold, this measures it)
    deberta_ms = None
    if _predictor:
        t0 = _t.monotonic()
        _predictor.predict("test", create_chart=False)
        deberta_ms = round((_t.monotonic() - t0) * 1000)

    return {
        "status": "ok",
        "version": os.environ.get("COMMIT_SHA", "dea7be4"),
        "on_hf": on_hf,
        "spotify": _recommender is not None,
        "genius": genius_ok,
        "deberta_warmup_ms": deberta_ms,
        "acoustic": _acoustic_status,
        "fusion": "active" if _fusion is not None else "off",
        "cpu": _cpu_info(),
    }


@app.get("/health/genius")
def health_genius():
    """Attempt a live Genius search to verify the token and connectivity."""
    import os
    token = os.environ.get("GENIUS_ACCESS_TOKEN")
    if not token:
        return {"status": "error", "reason": "GENIUS_ACCESS_TOKEN env var not set"}

    genius_client = _recommender.lyrics_fetcher.genius if _recommender and _recommender.lyrics_fetcher else None
    if not genius_client:
        init_err = getattr(getattr(_recommender, "lyrics_fetcher", None), "init_error", None)
        return {"status": "error", "reason": f"Genius client is None — init error: {init_err or 'unknown'}"}

    genius_result = {}
    try:
        song = genius_client.search_song("Happy", "Pharrell Williams")
        if song and song.lyrics:
            genius_result = {"status": "ok", "test_track": song.title, "lyrics_length": len(song.lyrics)}
        else:
            genius_result = {"status": "error", "reason": "search_song returned None or empty"}
    except Exception as e:
        genius_result = {"status": "error", "reason": f"{type(e).__name__}: {str(e)[:200]}"}

    # Test MusixMatch lyrics fallback
    import json, os, urllib.parse, urllib.request
    mm_key = os.environ.get("MUSIXMATCH_API_KEY")
    mm_result = {}
    if mm_key:
        try:
            params = urllib.parse.urlencode({"q_track": "Happy", "q_artist": "Pharrell Williams", "apikey": mm_key, "format": "json"})
            req = urllib.request.Request(f"https://api.musixmatch.com/ws/1.1/matcher.lyrics.get?{params}", headers={"User-Agent": "eumora/1.0"})
            with urllib.request.urlopen(req, timeout=8) as r:
                data = json.loads(r.read())
            hdr = data.get("message", {}).get("header", {})
            body = data.get("message", {}).get("body", {})
            lyrics_body = body.get("lyrics", {}).get("lyrics_body", "")
            mm_result = {"status_code": hdr.get("status_code"), "lyrics_preview": lyrics_body[:100] if lyrics_body else None}
        except Exception as e:
            mm_result = {"error": f"{type(e).__name__}: {e}"}
    else:
        mm_result = {"error": "MUSIXMATCH_API_KEY not set"}

    return {"genius": genius_result, "musixmatch": mm_result}


@app.get("/health/spotify")
def health_spotify():
    """Test that Spotify search actually returns tracks."""
    if not _recommender:
        return {"status": "error", "reason": "recommender not loaded"}
    try:
        res = _recommender.sp.search(q="happy pop upbeat", type="track", limit=5, market="US")
        items = res.get("tracks", {}).get("items", [])
        return {
            "status": "ok" if items else "empty",
            "track_count": len(items),
            "sample": [{"name": t["name"], "artist": t["artists"][0]["name"]} for t in items[:3]],
        }
    except Exception as e:
        return {"status": "error", "reason": f"{type(e).__name__}: {str(e)[:200]}"}


@app.post("/api/predict")
def predict(req: PredictRequest):
    if not _predictor:
        raise HTTPException(503, "Model not loaded")
    if not req.text and not req.reference_url:
        raise HTTPException(400, "Must provide text or reference_url")

    if req.raw:
        if not req.text:
            raise HTTPException(400, "raw prediction needs text")
        probs = _predictor.predict_raw(req.text)
        labels = sorted(({"label": k, "score": round(v, 6)} for k, v in probs.items()),
                        key=lambda x: x["score"], reverse=True)
        return {"labels": labels, "probabilities": {k: round(v, 6) for k, v in probs.items()}}

    prior = None if req.disable_prior_adjustment else req.target_sarcasm_prior

    text_result = None
    if req.text:
        text_result = _predictor.predict(req.text, target_sarcasm_prior=prior)

    ref_probs = None
    if req.reference_url and _recommender:
        try:
            ref_probs = _recommender.get_reference_emotion(req.reference_url, _predictor)
        except Exception as e:
            print(f"Reference URL failed in /api/predict: {e}")

    if text_result and ref_probs:
        t_probs = text_result.get("probabilities", {})
        merged = {k: (t_probs.get(k, 0.0) + ref_probs.get(k, 0.0)) / 2 for k in set(t_probs) | set(ref_probs)}
        top = max(merged, key=merged.get)
        return {
            "emotion":       top,
            "confidence":    merged[top],
            "probabilities": merged,
            "music_context": text_result.get("music_context", {}),
            "explanation":   text_result.get("explanation", ""),
        }
    elif ref_probs:
        top = max(ref_probs, key=ref_probs.get)
        return {
            "emotion":       top,
            "confidence":    ref_probs[top],
            "probabilities": ref_probs,
            "music_context": _predictor._build_music_context(top) if hasattr(_predictor, "_build_music_context") else {},
            "explanation":   "Emotion derived from the reference track's lyrics.",
        }
    elif text_result:
        return {
            "emotion":       text_result["emotion"],
            "confidence":    text_result["confidence"],
            "probabilities": text_result["probabilities"],
            "music_context": text_result["music_context"],
            "explanation":   text_result["explanation"],
        }
    else:
        raise HTTPException(503, "Could not extract emotion from text or reference track")


@app.post("/api/recommend")
def recommend(req: RecommendRequest):
    if not _predictor:
        raise HTTPException(503, "Model not loaded")

    if not req.text and not req.reference_url:
        raise HTTPException(400, "Must provide text or reference_url")

    prior = None if req.disable_prior_adjustment else req.target_sarcasm_prior

    predict_result = None
    if req.text:
        predict_result = _predictor.predict(req.text, target_sarcasm_prior=prior)
        
    ref_probs = None
    if req.reference_url and _recommender:
        ref_probs = _recommender.get_reference_emotion(req.reference_url, _predictor)

    if not predict_result and not ref_probs:
        raise HTTPException(400, "Could not extract emotion from text or reference track.")

    # Merge or assign
    if predict_result and ref_probs:
        t_probs = predict_result.get("probabilities", {})
        merged = {k: (t_probs.get(k, 0.0) + ref_probs.get(k, 0.0)) / 2 for k in set(t_probs) | set(ref_probs)}
        top = max(merged, key=merged.get)
        
        predict_result["emotion"] = top
        predict_result["confidence"] = merged[top]
        predict_result["probabilities"] = merged
        
    elif ref_probs:
        top = max(ref_probs, key=ref_probs.get)
        predict_result = {
            "emotion": top,
            "confidence": ref_probs[top],
            "probabilities": ref_probs,
            "explanation": "Derived entirely from the reference track's lyrics.",
            "music_context": _predictor._build_music_context(top) if hasattr(_predictor, '_build_music_context') else {}
        }
        
    # Standard recommendation flow
    if not _recommender:
        return {
            "emotion":            predict_result["emotion"],
            "confidence":         predict_result["confidence"],
            "probabilities":      predict_result.get("probabilities", {}),
            "explanation":        predict_result.get("explanation", ""),
            "music_context":      predict_result.get("music_context", {}),
            "targets_used":       None,
            "tracks":             [],
            "spotify_unavailable": True,
        }

    try:
        import time as _t
        _t0 = _t.monotonic()
        result = _recommender.recommend(
            predict_result,
            limit=req.limit,
            genre_override=req.genre,
            blend=req.blend,
            raw_text=req.text,
            predictor=_predictor,
        )
        result["_timing_s"] = round(_t.monotonic() - _t0, 2)
        return result
    except RuntimeError as exc:
        err_str = str(exc)
        print(f"Spotify recommend error: {err_str}")
        return {
            "emotion":             predict_result["emotion"],
            "confidence":          predict_result["confidence"],
            "probabilities":       predict_result.get("probabilities", {}),
            "explanation":         predict_result.get("explanation", ""),
            "music_context":       predict_result.get("music_context", {}),
            "targets_used":        None,
            "tracks":              [],
            "spotify_unavailable": True,
        }


# The scoring helpers live in src/scoring.py so the local benchmark can run the
# exact code this endpoint runs. Old names kept for anything importing them.
from src.scoring import (  # noqa: E402
    chunked_lyric_probs as _chunked_lyric_probs,
    combine as _combine,
    cosine as _cosine,
    emotion_match as _emotion_match,
    fused_response as _fused_response,
    lyrics_language_ok as _lyrics_language_ok,
    score_lyrics as _score_lyrics,
    score_probs as _score_probs,
)


@app.post("/api/score")
def score_track(req: ScoreRequest):
    """Score a track against the user's emotion by its lyrics and its sound.

    Lyrics: fetched, scored by stanza with DeBERTa, compared with the user's
    emotion vector. Audio: the 30 s preview, embedded with VGGish and placed in
    valence/arousal space. The two run in parallel and are fused 50/50; either
    alone is used when the other is unavailable (see scoring.combine).
    """
    if not _predictor:
        raise HTTPException(503, "Model not loaded")

    started = _time.monotonic()
    # Enter the lyric phase before submitting the audio, so the audio sees
    # this request's lyric work and waits for it rather than racing it.
    with _lyric_phase():
        audio_future = (_audio_pool.submit(_audio_for, req.title, req.artist, req.preview_url)
                        if _acoustic is not None else None)
        probs, lyrics_found, lyrics_english = _lyric_probs_for(req.title, req.artist)
    lyric_result = _score_probs(probs, req.user_probabilities) if probs else None

    audio, audio_match, pending = None, None, False
    if audio_future is not None:
        remaining = max(0.0, AUDIO_BUDGET_S - (_time.monotonic() - started))
        try:
            audio = audio_future.result(timeout=remaining)
        except FutureTimeout:
            pending = True  # still running; it will be cached when it finishes
        except Exception as e:
            print(f"Audio analysis failed for '{req.title}': {type(e).__name__}")
        if audio:
            audio_match = _acoustic.head.match(req.user_probabilities,
                                               (audio["valence"], audio["arousal"]))

    result = None
    if _fusion is not None:
        try:
            result = _fused_response(_fusion, req.user_probabilities, probs, audio, lyrics_english)
        except Exception as e:  # never let the new path take scoring down
            print(f"VAD fusion failed for '{req.title}': {type(e).__name__}: {e}")
            result = None
    if result is None:
        view = None
        if audio:
            view = {k: audio.get(k) for k in ("valence", "arousal", "sound", "descriptors", "preview_url")}
        result = _combine(lyric_result, view, audio_match)
    if result is None:
        result = {"match_score": None, "reason": "predict_failed" if lyrics_found else "no_lyrics"}
    if pending:
        result["audio_pending"] = True
    return result


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("api:app", host="0.0.0.0", port=8000, reload=True)