from __future__ import annotations

import asyncio
import logging
import os
import tempfile
import uuid
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from funasr import AutoModel
from pydantic import BaseModel

LOGGER = logging.getLogger("funasr-service")
MAX_AUDIO_BYTES = 15 * 1024 * 1024
CHUNK_BYTES = 1024 * 1024
SUPPORTED_FORMATS = {"mp3", "aac", "wav", "m4a", "amr"}


class Settings:
    def __init__(self) -> None:
        self.model = os.getenv("ASR_MODEL", "paraformer-zh")
        self.vad_model = os.getenv("ASR_VAD_MODEL", "fsmn-vad")
        self.punc_model = os.getenv("ASR_PUNC_MODEL", "ct-punc")
        self.device = os.getenv("ASR_DEVICE", "cpu")
        self.max_concurrency = positive_int("ASR_MAX_CONCURRENCY", 1)
        self.batch_size_seconds = positive_int("ASR_BATCH_SIZE_SECONDS", 300)


class HealthResponse(BaseModel):
    status: str
    model: str
    device: str


class TranscriptionResponse(BaseModel):
    text: str
    duration_ms: int
    request_id: str


def positive_int(name: str, default: int) -> int:
    raw = os.getenv(name, str(default))
    try:
        value = int(raw)
    except ValueError as error:
        raise RuntimeError(f"{name} must be an integer") from error
    if value <= 0:
        raise RuntimeError(f"{name} must be greater than 0")
    return value


def load_model(settings: Settings) -> Any:
    LOGGER.info(
        "loading FunASR model model=%s vad=%s punc=%s device=%s",
        settings.model,
        settings.vad_model,
        settings.punc_model,
        settings.device,
    )
    model = AutoModel(
        model=settings.model,
        vad_model=settings.vad_model,
        punc_model=settings.punc_model,
        device=settings.device,
        disable_update=True,
    )
    LOGGER.info("FunASR model is ready")
    return model


@asynccontextmanager
async def lifespan(app: FastAPI):
    settings = Settings()
    app.state.settings = settings
    app.state.model = await asyncio.to_thread(load_model, settings)
    app.state.inference_slots = asyncio.Semaphore(settings.max_concurrency)
    yield


app = FastAPI(
    title="Teaching Feedback FunASR Service",
    version="1.0.0",
    lifespan=lifespan,
)


@app.get("/health", response_model=HealthResponse)
async def health() -> HealthResponse:
    settings: Settings = app.state.settings
    return HealthResponse(status="ok", model=settings.model, device=settings.device)


@app.post("/v1/transcriptions", response_model=TranscriptionResponse)
async def create_transcription(
    file: UploadFile = File(...),
    language: str = Form("zh"),
    duration_ms: int = Form(0),
    hotword: str = Form(""),
) -> TranscriptionResponse:
    if language not in {"zh", "zh-cn", "cmn"}:
        raise HTTPException(status_code=422, detail="this service is configured for Chinese audio")
    audio_format = file_extension(file.filename)
    temporary_path = await save_upload(file, audio_format)
    request_id = str(uuid.uuid4())
    try:
        async with app.state.inference_slots:
            text = await asyncio.to_thread(
                run_inference,
                app.state.model,
                temporary_path,
                app.state.settings.batch_size_seconds,
                hotword.strip(),
            )
    except Exception as error:
        LOGGER.exception("FunASR inference failed request_id=%s", request_id)
        raise HTTPException(status_code=500, detail=f"speech recognition failed: {error}") from error
    finally:
        temporary_path.unlink(missing_ok=True)

    if not text:
        raise HTTPException(status_code=422, detail="no clear speech was recognized")
    return TranscriptionResponse(
        text=text,
        duration_ms=max(0, duration_ms),
        request_id=request_id,
    )


def file_extension(filename: str | None) -> str:
    suffix = Path(filename or "segment.mp3").suffix.lower().lstrip(".")
    if suffix not in SUPPORTED_FORMATS:
        raise HTTPException(status_code=415, detail="unsupported audio format")
    return suffix


async def save_upload(file: UploadFile, audio_format: str) -> Path:
    handle = tempfile.NamedTemporaryFile(suffix=f".{audio_format}", delete=False)
    path = Path(handle.name)
    total_bytes = 0
    try:
        with handle:
            while chunk := await file.read(CHUNK_BYTES):
                total_bytes += len(chunk)
                if total_bytes > MAX_AUDIO_BYTES:
                    raise HTTPException(status_code=413, detail="audio file exceeds 15 MB")
                handle.write(chunk)
    except Exception:
        path.unlink(missing_ok=True)
        raise
    finally:
        await file.close()
    if total_bytes == 0:
        path.unlink(missing_ok=True)
        raise HTTPException(status_code=422, detail="audio file is empty")
    return path


def run_inference(model: Any, path: Path, batch_size_seconds: int, hotword: str) -> str:
    options: dict[str, Any] = {
        "input": str(path),
        "batch_size_s": batch_size_seconds,
        "use_itn": True,
    }
    if hotword:
        options["hotword"] = hotword
    result = model.generate(**options)
    if not isinstance(result, list):
        return ""
    texts = [str(item.get("text", "")).strip() for item in result if isinstance(item, dict)]
    return "\n".join(text for text in texts if text).strip()