backend/api.py

"""
FastAPI Backend Presentation Layer - Async Optimized
"""
from fastapi import FastAPI, HTTPException, UploadFile, File, Form, Request
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import Optional
import base64
import os
import logging
import uuid
from io import BytesIO
from dotenv import load_dotenv

from backend.pipeline import run_pipeline, PipelineResult
from backend.sunbird_client import transcribe_audio_async

load_dotenv()

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

# Use a persistent, project-local directory so audio survives backend restarts.
# Override with SUNBIRD_AUDIO_CACHE_DIR env var for production deployments.
_AUDIO_CACHE_DIR = os.environ.get(
    "SUNBIRD_AUDIO_CACHE_DIR",
    os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "audio_cache"),
)
_AUDIO_CACHE_DIR = os.path.normpath(_AUDIO_CACHE_DIR)
os.makedirs(_AUDIO_CACHE_DIR, exist_ok=True)
MAX_CACHE_FILES = 100

# Simple in-memory job store: {request_id: {status: pending|running|completed|failed, result: dict or None, error: str or None, timestamp: float}}
_JOB_STORE = {}


def _detect_audio_media_type(data: bytes) -> str:
    """Return MIME type inferred from magic bytes; fall back to WAV."""
    if data[:4] == b"RIFF" and data[8:12] == b"WAVE":
        return "audio/wav"
    if data[:3] == b"ID3" or (len(data) > 1 and data[0] == 0xFF and (data[1] & 0xE0) == 0xE0):
        return "audio/mpeg"
    if data[:4] == b"OggS":
        return "audio/ogg"
    if data[4:8] in (b"ftyp", b"moov"):
        return "audio/mp4"
    return "audio/wav"


def _ext_for_media_type(media_type: str) -> str:
    return {
        "audio/mpeg": ".mp3",
        "audio/ogg": ".ogg",
        "audio/mp4": ".m4a",
    }.get(media_type, ".wav")


app = FastAPI(
    title="Sunbird AI Pipeline API",
    description="Transcribe audio, summarise and translate text, generate speech — all via Sunbird AI",
    version="1.0.0",
)

# Middleware to add request ID for tracing
@app.middleware("http")
async def add_request_id(request: Request, call_next):
    request.state.request_id = str(uuid.uuid4())
    logger.info(f"[{request.state.request_id}] {request.method} {request.url.path}")
    response = await call_next(request)
    return response

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],      # allow all origins for local dev
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── Centralized Exception Handlers ─────────────────────────────────────────────
@app.exception_handler(ValueError)
async def value_error_handler(request: Request, exc: ValueError):
    logger.warning(f"Validation error on {request.url.path}: {exc}")
    return JSONResponse(
        status_code=400,
        content={"error": True, "type": "ValidationError", "detail": str(exc)},
    )

@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    logger.error(f"Unexpected error on {request.url.path}: {exc}", exc_info=True)
    return JSONResponse(
        status_code=500,
        content={"error": True, "type": "InternalServerError", "detail": str(exc)},
    )

class PipelineResponse(BaseModel):
    transcript: Optional[str] = None
    summary: str
    translated_summary: str
    target_language: str
    audio_b64: Optional[str] = None    # Optional for backward compat; use /audio/{request_id} endpoint for large files
    word_count: int
    reading_time_seconds: int
    request_id: Optional[str] = None   # Tracing ID

def _pipeline_result_to_response(result: PipelineResult, request_id: str, include_audio_b64: bool = False) -> dict:
    """
    Convert PipelineResult to response dict.
    If include_audio_b64=False, omit audio from JSON to reduce payload size.
    Use GET /pipeline/audio/{request_id} to fetch audio separately.
    """
    response_data = {
        "transcript": result.transcript,
        "summary": result.summary,
        "translated_summary": result.translated_summary,
        "target_language": result.target_language,
        "word_count": result.word_count,
        "reading_time_seconds": result.reading_time_seconds,
        "request_id": request_id,
    }
    
    # Store audio in cache for streaming endpoint
    if result.audio_bytes:
        # Enforce max files limit
        cache_files = os.listdir(_AUDIO_CACHE_DIR)
        if len(cache_files) >= MAX_CACHE_FILES:
            # simple eviction: delete oldest file
            cache_files.sort(key=lambda f: os.path.getmtime(os.path.join(_AUDIO_CACHE_DIR, f)))
            try:
                os.remove(os.path.join(_AUDIO_CACHE_DIR, cache_files[0]))
            except OSError:
                pass
                
        # Detect real content type and write with correct extension
        media_type = _detect_audio_media_type(result.audio_bytes)
        ext = _ext_for_media_type(media_type)
        audio_path = os.path.join(_AUDIO_CACHE_DIR, f"{request_id}{ext}")
        with open(audio_path, "wb") as f:
            f.write(result.audio_bytes)

        # Store media type alongside audio so the GET endpoint serves it correctly
        meta_path = os.path.join(_AUDIO_CACHE_DIR, f"{request_id}.ct")
        with open(meta_path, "w") as f:
            f.write(media_type)
        
        # Include base64 only if small and explicitly requested (backward compat)
        if include_audio_b64 and len(result.audio_bytes) < 1_000_000:
            response_data["audio_b64"] = base64.b64encode(result.audio_bytes).decode("utf-8")
        else:
            response_data["audio_b64"] = None
    else:
        response_data["audio_b64"] = None
    
    return response_data

@app.get("/", tags=["Health"])
async def health_check():
    logger.info("Health check endpoint accessed")
    return {
        "status": "healthy", 
        "service": "Sunbird AI Pipeline API", 
        "version": "1.0.0"
    }

class TextRequest(BaseModel):
    text: str = Field(..., min_length=1, max_length=50000, description="The input text to process")
    target_language: str = Field("Luganda", description="The target language for translation")

@app.post("/pipeline/text", response_model=PipelineResponse)
async def pipeline_text(body: TextRequest, request: Request):
    request_id = request.state.request_id
    logger.info(f"[{request_id}] Processing text pipeline request to {body.target_language}")
    try:
        result = await run_pipeline(text=body.text, target_language=body.target_language)
        logger.info(f"[{request_id}] Pipeline completed successfully")
        return _pipeline_result_to_response(result, request_id, include_audio_b64=False)
    except Exception as e:
        logger.error(f"[{request_id}] Pipeline failed: {e}", exc_info=True)
        raise


# -------------------- Async job endpoints --------------------
@app.post("/pipeline/submit")
async def pipeline_submit(body: TextRequest, request: Request):
    """
    Submit a text pipeline job. Returns immediately with a `request_id`.
    The job runs in the background and can be polled via `/pipeline/status/{request_id}`.
    """
    request_id = request.state.request_id
    logger.info(f"[{request_id}] Submitting text job to background")

    # initialize job store and prune old jobs
    import time
    current_time = time.time()
    
    # Prune jobs older than 1 hour (3600 seconds)
    keys_to_delete = [k for k, v in _JOB_STORE.items() if current_time - v.get("timestamp", current_time) > 3600]
    for k in keys_to_delete:
        del _JOB_STORE[k]
        
    _JOB_STORE[request_id] = {"status": "pending", "result": None, "error": None, "timestamp": current_time}

    async def _background():
        _JOB_STORE[request_id]["status"] = "running"
        try:
            res = await run_pipeline(text=body.text, target_language=body.target_language)
            # convert and store response; this will also cache audio bytes
            resp = _pipeline_result_to_response(res, request_id, include_audio_b64=False)
            _JOB_STORE[request_id]["result"] = resp
            _JOB_STORE[request_id]["status"] = "completed"
            logger.info(f"[{request_id}] Background job completed")
        except Exception as e:
            logger.error(f"[{request_id}] Background job failed: {e}", exc_info=True)
            _JOB_STORE[request_id]["error"] = str(e)
            _JOB_STORE[request_id]["status"] = "failed"

    # schedule background task
    import asyncio
    asyncio.create_task(_background())

    return {"request_id": request_id, "status": "submitted"}


@app.post("/pipeline/submit_audio")
async def pipeline_submit_audio(
    file: UploadFile = File(...),
    target_language: str = Form("Luganda"),
    source_language: Optional[str] = Form(None),
    request: Request = None,
):
    request_id = request.state.request_id if request else str(uuid.uuid4())
    logger.info(f"[{request_id}] Submitting audio job to background for file {file.filename}")

    audio_bytes = await file.read()

    import time
    current_time = time.time()
    
    # Prune jobs older than 1 hour (3600 seconds)
    keys_to_delete = [k for k, v in _JOB_STORE.items() if current_time - v.get("timestamp", current_time) > 3600]
    for k in keys_to_delete:
        del _JOB_STORE[k]
        
    _JOB_STORE[request_id] = {"status": "pending", "result": None, "error": None, "timestamp": current_time}

    async def _background_audio():
        _JOB_STORE[request_id]["status"] = "running"
        try:
            res = await run_pipeline(
                audio_bytes=audio_bytes,
                audio_filename=file.filename,
                target_language=target_language,
                source_language=source_language,
            )
            resp = _pipeline_result_to_response(res, request_id, include_audio_b64=False)
            _JOB_STORE[request_id]["result"] = resp
            _JOB_STORE[request_id]["status"] = "completed"
            logger.info(f"[{request_id}] Background audio job completed")
        except Exception as e:
            logger.error(f"[{request_id}] Background audio job failed: {e}", exc_info=True)
            _JOB_STORE[request_id]["error"] = str(e)
            _JOB_STORE[request_id]["status"] = "failed"

    import asyncio
    asyncio.create_task(_background_audio())

    return {"request_id": request_id, "status": "submitted"}


@app.get("/pipeline/status/{request_id}")
async def pipeline_status(request_id: str):
    """Return job status and result when available."""
    if request_id not in _JOB_STORE:
        raise HTTPException(status_code=404, detail="Job not found")

    job = _JOB_STORE[request_id]
    resp = {"request_id": request_id, "status": job["status"]}
    if job["status"] == "completed":
        resp["result"] = job["result"]
    elif job["status"] == "failed":
        resp["error"] = job["error"]

    return resp

@app.post("/pipeline/audio", response_model=PipelineResponse)
async def pipeline_audio(
    file: UploadFile = File(...),
    target_language: str = Form("Luganda"),
    source_language: Optional[str] = Form(None),
    request: Request = None,
):
    request_id = request.state.request_id if request else str(uuid.uuid4())
    logger.info(f"[{request_id}] Processing audio pipeline for file {file.filename}")
    try:
        audio_bytes = await file.read()
        logger.info(f"[{request_id}] Audio file read: {len(audio_bytes)} bytes")
        result = await run_pipeline(
            audio_bytes=audio_bytes,
            audio_filename=file.filename,
            target_language=target_language,
            source_language=source_language,
        )
        logger.info(f"[{request_id}] Pipeline completed successfully")
        return _pipeline_result_to_response(result, request_id, include_audio_b64=False)
    except Exception as e:
        logger.error(f"[{request_id}] Pipeline failed: {e}", exc_info=True)
        raise

@app.get("/pipeline/audio/{request_id}")
async def get_pipeline_audio(request_id: str):
    """
    Fetch generated audio for a completed pipeline request.
    Serves the correct content type inferred at write time.
    """
    # Resolve filename: check each supported extension
    audio_path = None
    media_type = "audio/wav"
    for ext in (".wav", ".mp3", ".ogg", ".m4a"):
        candidate = os.path.join(_AUDIO_CACHE_DIR, f"{request_id}{ext}")
        if os.path.exists(candidate):
            audio_path = candidate
            break

    if audio_path is None:
        logger.warning(f"Audio not found for request {request_id}")
        raise HTTPException(status_code=404, detail="Audio not found or has expired")

    # Read stored content type if available
    meta_path = os.path.join(_AUDIO_CACHE_DIR, f"{request_id}.ct")
    if os.path.exists(meta_path):
        with open(meta_path) as f:
            media_type = f.read().strip() or media_type

    logger.info(f"Streaming audio for request {request_id} ({media_type}) from {audio_path}")
    ext = os.path.splitext(audio_path)[1] or ".wav"

    def iterfile():
        with open(audio_path, mode="rb") as fh:
            yield from fh

    return StreamingResponse(
        iterfile(),
        media_type=media_type,
        headers={"Content-Disposition": f"attachment; filename=speech{ext}"},
    )


@app.post("/stt/test")
async def stt_test(
    file: UploadFile = File(...),
    source_language: Optional[str] = Form(None),
    request: Request = None,
):
    """
    Focused STT integration endpoint for UI testing.
    Accepts an audio file and optional source_language hint, returns the transcription only.
    """
    request_id = request.state.request_id if request else str(uuid.uuid4())
    logger.info(f"[{request_id}] STT test endpoint called for file {file.filename}")
    audio_bytes = await file.read()
    try:
        transcript = await transcribe_audio_async(audio_bytes, file.filename, language=source_language)
        return {"request_id": request_id, "transcript": transcript}
    except Exception as e:
        logger.error(f"[{request_id}] STT test failed: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))

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