backend/main.py

import os
import logging
import tempfile
import subprocess
import re
from contextlib import asynccontextmanager
from functools import lru_cache
from typing import Any

import requests
from dotenv import load_dotenv
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
from transformers import pipeline

load_dotenv()

MODEL_ID = os.getenv("MODEL_ID", "viswamaicoe/swecha-gonthuka-asr")
ASR_DEVICE = os.getenv("ASR_DEVICE", "cpu").lower()
SWECHA_API_BASE = os.getenv("SWECHA_API_BASE", "https://api.corpus.swecha.org")
SWECHA_UPLOAD_PATH = os.getenv("SWECHA_UPLOAD_PATH", "/api/v1/content")
SWECHA_AUTH_TOKEN = os.getenv("SWECHA_AUTH_TOKEN", "")
HF_TOKEN = os.getenv("HF_TOKEN", "")
HF_ASR_MODEL = os.getenv("HF_ASR_MODEL", "viswamaicoe/swecha-gonthuka-asr")
HF_API_BASE = os.getenv("HF_API_BASE", "https://router.huggingface.co/hf-inference/models")


@lru_cache(maxsize=1)
def get_asr_pipeline():
    device = 0 if ASR_DEVICE == "cuda" else -1
    return pipeline(
        task="automatic-speech-recognition",
        model=MODEL_ID,
        device=device,
        chunk_length_s=30,  # Process audio in 30s chunks
        batch_size=8,  # Batch processes for speed
    )


def clean_noisy_telugu(text: str) -> str:
    """Removes common phonetic marker noise in Telugu ASR outputs."""
    if not text:
        return ""
    # Remove repeated phonetic markers like '్' if they appear in excessive sequences
    cleaned = re.sub(r"్{2,}", "్", text)
    # Basic cleanup of extra spaces
    cleaned = re.sub(r"\s+", " ", cleaned).strip()
    return cleaned


def transcribe_via_hf_api(raw_bytes: bytes) -> dict[str, Any]:
    """Send audio bytes to HuggingFace Inference API (fast, no local model needed)."""
    url = f"{HF_API_BASE.rstrip('/')}/{HF_ASR_MODEL}"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    resp = requests.post(url, headers=headers, data=raw_bytes, timeout=120)
    if resp.status_code != 200:
        raise HTTPException(
            status_code=resp.status_code,
            detail=f"HF API error: {resp.text[:300]}",
        )
    data = resp.json()
    text = data.get("text") or ""
    return {"text": clean_noisy_telugu(text)}


def transcribe_file_bytes(raw_bytes: bytes, suffix: str) -> dict[str, Any]:
    try:
        asr = get_asr_pipeline()
    except Exception as exc:
        raise HTTPException(
            status_code=500, detail=f"ASR model load failed: {exc}"
        ) from exc

    # Save original uploaded file
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as input_file:
        input_file.write(raw_bytes)
        input_path = input_file.name

    # Create converted WAV file
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as output_file:
        output_path = output_file.name

    try:
        # Convert to 16kHz mono PCM WAV
        ffmpeg_command = [
            "ffmpeg",
            "-y",
            "-i",
            input_path,
            "-acodec",
            "pcm_s16le",
            "-ac",
            "1",
            "-ar",
            "16000",
            output_path,
        ]

        result_ffmpeg = subprocess.run(
            ffmpeg_command,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )

        if result_ffmpeg.returncode != 0:
            raise HTTPException(
                status_code=400,
                detail=f"FFmpeg conversion failed: {result_ffmpeg.stderr.decode()}",
            )

        # Run ASR on converted file
        result = asr(
            output_path, generate_kwargs={"task": "transcribe", "language": "telugu"}
        )

        text = ""
        if isinstance(result, dict) and "text" in result:
            text = clean_noisy_telugu(result["text"])
        elif isinstance(result, str):
            text = clean_noisy_telugu(result)

    except Exception as exc:
        raise HTTPException(
            status_code=400, detail=f"Transcription failed: {exc}"
        ) from exc

    finally:
        # Clean up temp files
        for path in (input_path, output_path):
            try:
                if os.path.exists(path):
                    os.remove(path)
            except OSError:
                pass

    return {"text": text}


def push_to_swecha(
    audio_bytes: bytes,
    filename: str,
    content_type: str,
    transcript: str,
    title: str,
    description: str,
) -> dict[str, Any]:
    if not SWECHA_AUTH_TOKEN:
        raise HTTPException(
            status_code=400, detail="SWECHA_AUTH_TOKEN is not configured"
        )

    url = f"{SWECHA_API_BASE.rstrip('/')}/{SWECHA_UPLOAD_PATH.lstrip('/')}"
    headers = {"Authorization": f"Bearer {SWECHA_AUTH_TOKEN}"}

    files = {
        "audio": (filename or "audio.webm", audio_bytes, content_type or "audio/webm")
    }

    data = {
        "title": title,
        "description": description,
        "transcript": transcript,
    }

    resp = requests.post(url, headers=headers, files=files, data=data, timeout=60)

    try:
        payload = resp.json()
    except Exception:
        payload = {"raw": resp.text}

    if resp.status_code >= 400:
        raise HTTPException(
            status_code=resp.status_code, detail={"swecha_error": payload}
        )

    return payload


logger = logging.getLogger("asr")


@asynccontextmanager
async def lifespan(app):
    """Pre-warm the ASR model at startup so first transcription is fast."""
    import threading
    def _load():
        try:
            logger.info("Pre-warming ASR model '%s'...", MODEL_ID)
            get_asr_pipeline()
            logger.info("ASR model ready.")
        except Exception as exc:
            logger.warning("Model pre-warm failed (will retry on first request): %s", exc)
    threading.Thread(target=_load, daemon=True).start()
    yield


app = FastAPI(title="Swecha Telugu ASR Service", version="1.0.0", lifespan=lifespan)

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


@app.get("/health")
def health():
    return {
        "status": "ok",
        "model": MODEL_ID,
        "device": ASR_DEVICE,
        "model_loaded": get_asr_pipeline.cache_info().currsize > 0,
    }

# API routes will be defined below


@app.post("/transcribe")
async def transcribe(audio: UploadFile = File(...)):
    raw_bytes = await audio.read()
    if not raw_bytes:
        raise HTTPException(status_code=400, detail="Empty audio file")

    suffix = os.path.splitext(audio.filename or "audio.webm")[1] or ".webm"
    result = transcribe_file_bytes(raw_bytes, suffix)
    return result


class TranscribeUrlRequest(BaseModel):
    audio_url: str


@app.post("/transcribe-url")
async def transcribe_url(req: TranscribeUrlRequest):
    """Fetch audio from a URL server-side (avoids browser CORS) and transcribe it.

    Handles two URL types:
    - Swecha Corpus API URLs: need Bearer auth header
    - Pre-signed S3/CDN URLs: must NOT send Authorization (causes 400/Content-Length mismatch)
    """
    try:
        audio_headers = {}
        # Detect if URL is a direct Swecha Corpus API URL (not a pre-signed S3/CDN URL).
        # Pre-signed URLs contain query params like X-Amz-Signature or token= — skip auth.
        swecha_host = SWECHA_API_BASE.replace("https://", "").replace("http://", "").split("/")[0]
        is_swecha_host = swecha_host in req.audio_url
        is_presigned = any(p in req.audio_url for p in ("X-Amz-", "token=", "Signature=", "AWSAccessKeyId="))
        if SWECHA_AUTH_TOKEN and is_swecha_host and not is_presigned:
            audio_headers["Authorization"] = f"Bearer {SWECHA_AUTH_TOKEN}"

        # Use stream=True to avoid Content-Length mismatch on CDN/S3 responses.
        # This reads the body chunk-by-chunk regardless of what the headers say.
        with requests.get(
            req.audio_url,
            headers=audio_headers,
            timeout=60,
            stream=True,
            allow_redirects=True,
        ) as audio_resp:
            if audio_resp.status_code == 403:
                raise HTTPException(
                    status_code=400,
                    detail="Audio URL expired or access denied (403). Please reload standups to get a fresh URL.",
                )
            audio_resp.raise_for_status()
            raw_bytes = audio_resp.content  # stream=True makes this safe vs Content-Length
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(
            status_code=400, detail=f"Failed to fetch audio from URL: {exc}"
        ) from exc

    if not raw_bytes:
        raise HTTPException(status_code=400, detail="Empty audio file at URL")

    url_path = req.audio_url.split("?")[0]
    suffix = os.path.splitext(url_path)[1] or ".webm"
    return transcribe_file_bytes(raw_bytes, suffix)


@app.post("/transcribe-and-store")
async def transcribe_and_store(
    audio: UploadFile = File(...),
    title: str = Form(...),
    description: str = Form(...),
):
    raw_bytes = await audio.read()
    if not raw_bytes:
        raise HTTPException(status_code=400, detail="Empty audio file")

    suffix = os.path.splitext(audio.filename or "audio.webm")[1] or ".webm"
    result = transcribe_file_bytes(raw_bytes, suffix)
    text = result.get("text", "")

    swecha_response = push_to_swecha(
        audio_bytes=raw_bytes,
        filename=audio.filename or f"audio{suffix}",
        content_type=audio.content_type or "audio/webm",
        transcript=text,
        title=title,
        description=description,
    )

    return {
        "text": text,
        "raw": result,
        "swecha_response": swecha_response,
    }

# FINAL STEP: Mount static files from the frontend directory
# We MUST do this AFTER all other routes so it doesn't intercept them.
frontend_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "frontend"))

if os.path.exists(frontend_path):
    @app.get("/", include_in_schema=False)
    async def read_index():
        return FileResponse(os.path.join(frontend_path, "index.html"))

    # Mount remaining static files (js, css)
    app.mount("/", StaticFiles(directory=frontend_path), name="static")