"""
Speech2Text MBIP — Transcription Worker
FastAPI service wrapping faster-whisper for Bahasa Melayu transcription.
Internal-only service; not exposed to internet.
"""

import os
import tempfile
import logging
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, UploadFile, File, HTTPException
from pydantic import BaseModel
from faster_whisper import WhisperModel

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="Speech2Text Transcription Worker", docs_url=None, redoc_url=None)

# Config from environment
MODEL_SIZE = os.getenv("WHISPER_MODEL", "small")
LANGUAGE = os.getenv("WHISPER_LANGUAGE", "ms")
DEVICE = os.getenv("WHISPER_DEVICE", "cpu")
COMPUTE_TYPE = os.getenv("WHISPER_COMPUTE_TYPE", "int8")

logger.info(f"Loading Whisper model: {MODEL_SIZE}, device: {DEVICE}, compute: {COMPUTE_TYPE}")

# Load model once at startup
model = WhisperModel(MODEL_SIZE, device=DEVICE, compute_type=COMPUTE_TYPE)

logger.info("Whisper model loaded successfully.")


class TranscribeResponse(BaseModel):
    success: bool
    transcript: Optional[str] = None
    language: Optional[str] = None
    confidence: Optional[float] = None
    duration_seconds: Optional[float] = None
    error: Optional[str] = None


@app.get("/health")
def health_check():
    return {"status": "ok", "model": MODEL_SIZE, "language": LANGUAGE}


@app.post("/transcribe", response_model=TranscribeResponse)
async def transcribe_audio(audio: UploadFile = File(...)):
    """
    Terima fail audio, transkripkan ke teks Bahasa Melayu.
    Hanya accessible dari dalam Docker network sahaja.
    """
    # Validate content type
    allowed_types = {
        "audio/mpeg", "audio/mp3", "audio/wav", "audio/x-wav",
        "audio/mp4", "audio/x-m4a", "audio/aac", "video/mp4",
        "audio/ogg", "audio/flac", "audio/webm",
    }

    if audio.content_type and audio.content_type not in allowed_types:
        logger.warning(f"Rejected content type: {audio.content_type}")
        raise HTTPException(status_code=400, detail=f"Unsupported audio format: {audio.content_type}")

    # Save to temp file
    suffix = Path(audio.filename or "audio.mp3").suffix or ".mp3"

    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            content = await audio.read()
            tmp.write(content)
            tmp_path = tmp.name

        logger.info(f"Transcribing: {audio.filename}, size: {len(content)} bytes")

        # Run transcription
        segments, info = model.transcribe(
            tmp_path,
            language=LANGUAGE,
            beam_size=5,
            vad_filter=True,
        )

        # Collect segments
        all_text = []
        total_confidence = []

        for segment in segments:
            all_text.append(segment.text.strip())
            # faster-whisper segments have avg_logprob; convert to 0-1 range
            if hasattr(segment, "avg_logprob") and segment.avg_logprob is not None:
                import math
                confidence = min(1.0, max(0.0, math.exp(segment.avg_logprob)))
                total_confidence.append(confidence)

        transcript = " ".join(all_text).strip()
        avg_confidence = sum(total_confidence) / len(total_confidence) if total_confidence else None
        duration = info.duration if hasattr(info, "duration") else None

        logger.info(f"Transcription complete: {len(transcript)} chars, lang: {info.language}")

        return TranscribeResponse(
            success=True,
            transcript=transcript,
            language=info.language,
            confidence=round(avg_confidence, 4) if avg_confidence else None,
            duration_seconds=round(duration, 2) if duration else None,
        )

    except Exception as e:
        logger.error(f"Transcription failed: {e}", exc_info=True)
        return TranscribeResponse(
            success=False,
            error=str(e),
        )

    finally:
        # Always cleanup temp file
        if "tmp_path" in locals():
            try:
                os.unlink(tmp_path)
            except Exception:
                pass