app/services/fst_client.py

"""
FST (Fusion Segment Transformer) external API client (Layer 3).

Calls the HuggingFace Space ``mippia/AI-Music-Detection-FST``
Gradio API for high-accuracy AI music detection.

FST uses MERT + beat-aware segmentation and reports 99.99%
accuracy on benchmark datasets.  We treat it as a strong
external signal in the score-fusion pipeline.

Gracefully returns unavailable result on timeout or error.
"""

from __future__ import annotations

import asyncio
import io
import os
import tempfile
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Union

from .logging_config import get_logger

logger = get_logger(__name__)

# HuggingFace Space endpoint
FST_SPACE_ID = "mippia/AI-Music-Detection-FST"
FST_API_URL = f"https://{FST_SPACE_ID.replace('/', '-')}.hf.space"

# Timeouts
FST_CONNECT_TIMEOUT = 10.0   # seconds
FST_PREDICT_TIMEOUT = float(os.getenv("FST_PREDICT_TIMEOUT", "60"))
# After a failed connect, try again this much later instead of never.
FST_RETRY_AFTER = 300.0
# Calls in flight at once; a timed-out call keeps its thread until the
# remote answers, so this also bounds how many threads can pile up.
FST_MAX_CONCURRENT = 2
FST_ENABLED = os.getenv("FST_ENABLED", "1") != "0"


@dataclass
class FSTResult:
    """Result from FST external service."""

    available: bool
    is_ai: bool = False
    confidence: float = 0.5
    label: str = "unknown"
    raw_scores: Optional[dict] = None
    error: Optional[str] = None


class FSTClientService:
    """
    Client for FST AI Music Detection HuggingFace Space.

    Uses the Gradio Client API to submit audio and receive
    predictions. Falls back gracefully if the space is
    sleeping, overloaded, or unreachable.
    """

    def __init__(self) -> None:
        self._client = None
        self._available: Optional[bool] = None
        self._failed_at = 0.0
        self._gate: Optional[asyncio.Semaphore] = None

    def _ensure_client(self) -> bool:
        """Lazy-initialize Gradio client (blocking: call from a thread)."""
        if self._available is False and time.monotonic() - self._failed_at > FST_RETRY_AFTER:
            self._available = None
        if self._available is not None:
            return self._available

        try:
            from gradio_client import Client
            self._client = Client(
                FST_SPACE_ID,
                hf_token=None,  # Public space
            )
            self._available = True
            logger.info(f"FST client connected: {FST_SPACE_ID}")
            return True
        except ImportError:
            logger.warning(
                "gradio_client not installed — FST layer disabled"
            )
            self._available = False
            self._failed_at = float("inf")
            return False
        except Exception as e:
            logger.warning(f"FST client init failed: {e}")
            self._available = False
            self._failed_at = time.monotonic()
            return False

    async def predict(
        self,
        source: Union[Path, bytes, io.BytesIO],
        suffix: str = ".wav",
    ) -> FSTResult:
        """
        Submit audio to FST Space for AI detection.

        Args:
            source: Audio file path, raw bytes, or BytesIO.

        Returns:
            FSTResult with detection outcome.
        """
        if not FST_ENABLED:
            return FSTResult(available=False, error="fst_disabled")
        if self._gate is None:
            self._gate = asyncio.Semaphore(FST_MAX_CONCURRENT)

        audio_path: Optional[Path] = None
        try:
            async with self._gate:
                # Connecting fetches the remote config over the network.
                if not await asyncio.to_thread(self._ensure_client):
                    return FSTResult(
                        available=False,
                        error="fst_client_unavailable",
                    )
                audio_path = self._to_file_path(source, suffix)
                return await self._call_api(audio_path)
        except Exception as e:
            logger.warning(f"FST prediction failed: {e}")
            return FSTResult(
                available=False,
                error=str(e),
            )
        finally:
            if audio_path is not None and not isinstance(source, Path):
                audio_path.unlink(missing_ok=True)

    async def _call_api(self, audio_path: Path) -> FSTResult:
        """Call the FST Gradio API."""
        try:
            # Run synchronous Gradio client in executor
            loop = asyncio.get_running_loop()
            result = await asyncio.wait_for(
                loop.run_in_executor(
                    None,
                    self._sync_predict,
                    audio_path,
                ),
                timeout=FST_PREDICT_TIMEOUT,
            )
            return result

        except asyncio.TimeoutError:
            logger.warning(
                f"FST prediction timed out after "
                f"{FST_PREDICT_TIMEOUT}s"
            )
            return FSTResult(
                available=False,
                error="fst_timeout",
            )

    def _sync_predict(self, audio_path: Path) -> FSTResult:
        """Synchronous Gradio predict call."""
        try:
            result = self._client.predict(
                str(audio_path),
                api_name="/predict",
            )

            # Parse Gradio response
            # FST typically returns label + confidence dict
            return self._parse_response(result)

        except Exception as e:
            logger.warning(f"FST sync predict error: {e}")
            return FSTResult(
                available=False,
                error=str(e),
            )

    def _parse_response(self, response: object) -> FSTResult:
        """
        Parse FST Gradio API response.

        FST response format varies — handle multiple formats:
        1. Dict with 'label' and 'confidences'
        2. String label with confidence
        3. Raw dict with scores
        """
        try:
            if isinstance(response, dict):
                return self._parse_dict_response(response)
            elif isinstance(response, str):
                return self._parse_string_response(response)
            elif isinstance(response, (list, tuple)):
                # First element is usually the classification
                if len(response) > 0:
                    return self._parse_response(response[0])
            else:
                logger.warning(
                    f"Unexpected FST response type: "
                    f"{type(response)}"
                )
                return FSTResult(
                    available=True,
                    label="parse_error",
                    error=f"unexpected_type: {type(response).__name__}",
                )
        except Exception as e:
            logger.warning(f"FST response parse error: {e}")
            return FSTResult(
                available=False,
                error=f"parse_error: {e}",
            )

    def _parse_dict_response(self, data: dict) -> FSTResult:
        """Parse dict-style response."""
        # Format: {"label": "AI", "confidences": [{"label": "AI", "confidence": 0.99}, ...]}
        label = data.get("label", "unknown")
        confidences = data.get("confidences", [])

        is_ai = "ai" in label.lower() or "fake" in label.lower()

        confidence = 0.5
        if confidences and isinstance(confidences, list):
            for item in confidences:
                if isinstance(item, dict):
                    item_label = item.get("label", "")
                    if "ai" in item_label.lower() or "fake" in item_label.lower():
                        confidence = float(
                            item.get("confidence", 0.5)
                        )
                        break

        # If no AI confidence found, use first confidence
        if confidence == 0.5 and confidences:
            first = confidences[0]
            if isinstance(first, dict):
                confidence = float(
                    first.get("confidence", 0.5)
                )
                if not is_ai:
                    confidence = 1.0 - confidence

        return FSTResult(
            available=True,
            is_ai=is_ai,
            confidence=round(
                max(0.01, min(0.99, confidence)), 4
            ),
            label=label,
            raw_scores=data,
        )

    def _parse_string_response(self, text: str) -> FSTResult:
        """Parse string-style response."""
        lower = text.lower().strip()
        is_ai = any(
            kw in lower
            for kw in ("ai", "fake", "generated", "synthetic")
        )
        # Conservative confidence for string-only responses
        confidence = 0.75 if is_ai else 0.25

        return FSTResult(
            available=True,
            is_ai=is_ai,
            confidence=confidence,
            label=text.strip(),
        )

    @staticmethod
    def _to_file_path(
        source: Union[Path, bytes, io.BytesIO],
        suffix: str = ".wav",
    ) -> Path:
        """Convert source to a file path for Gradio upload (caller deletes it)."""
        if isinstance(source, Path):
            return source
        if isinstance(source, bytes):
            source = io.BytesIO(source)
        tmp = tempfile.NamedTemporaryFile(
            suffix=suffix or ".wav", delete=False,
        )
        tmp.write(source.read())
        tmp.flush()
        tmp.close()
        return Path(tmp.name)