import logging
import os
import tempfile
import time
from typing import Optional

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import PlainTextResponse

logger = logging.getLogger("asr_cpu")
logging.basicConfig(level=logging.INFO)

ASR_ENGINE = os.getenv("ASR_ENGINE", "paraformer").strip().lower()
ASR_PORT = int(os.getenv("ASR_PORT", "9805"))

PARAFORMER_MODEL = os.getenv("PARAFORMER_MODEL", "/models/paraformer-mindie-common")
PARAFORMER_BATCH_SIZE_S = int(os.getenv("PARAFORMER_BATCH_SIZE_S", "300"))

WHISPER_MODEL = os.getenv(
    "WHISPER_MODEL", "/models/faster-whisper/Systran/faster-whisper-base"
)
WHISPER_DEVICE = os.getenv("WHISPER_DEVICE", "cpu")
WHISPER_COMPUTE_TYPE = os.getenv("WHISPER_COMPUTE_TYPE", "int8")
WHISPER_BEAM_SIZE = int(os.getenv("WHISPER_BEAM_SIZE", "5"))
WHISPER_VAD_FILTER = os.getenv("WHISPER_VAD_FILTER", "false").lower() == "true"
WHISPER_CPU_THREADS = int(os.getenv("WHISPER_CPU_THREADS", "4"))
WHISPER_FORCE_SIMPLIFIED = os.getenv("WHISPER_FORCE_SIMPLIFIED", "true").lower() == "true"

app = FastAPI(title="Unified CPU ASR API", version="1.0.0")
model = None
to_simplified = None


def _simplifier():
    global to_simplified
    if to_simplified is not None:
        return to_simplified
    if not WHISPER_FORCE_SIMPLIFIED:
        to_simplified = lambda text: text
        return to_simplified
    try:
        from opencc import OpenCC

        cc = OpenCC("t2s")
        to_simplified = lambda text: cc.convert(text) if text else text
    except Exception as exc:
        logger.warning("OpenCC unavailable, skip t2s conversion: %s", exc)
        to_simplified = lambda text: text
    return to_simplified


def load_model():
    global model
    if model is not None:
        return model

    started = time.perf_counter()
    if ASR_ENGINE in {"paraformer", "paraformer-cpu", "funasr"}:
        from funasr import AutoModel

        logger.info("Loading Paraformer CPU model: %s", PARAFORMER_MODEL)
        model = AutoModel(
            model=PARAFORMER_MODEL,
            device="cpu",
            disable_update=True,
            disable_pbar=True,
        )
    elif ASR_ENGINE in {"faster-whisper", "whisper", "whisper-base", "whisper-tiny"}:
        from faster_whisper import WhisperModel

        logger.info(
            "Loading faster-whisper model=%s device=%s compute=%s",
            WHISPER_MODEL,
            WHISPER_DEVICE,
            WHISPER_COMPUTE_TYPE,
        )
        model = WhisperModel(
            WHISPER_MODEL,
            device=WHISPER_DEVICE,
            compute_type=WHISPER_COMPUTE_TYPE,
            local_files_only=True,
            cpu_threads=WHISPER_CPU_THREADS,
        )
    else:
        raise RuntimeError(f"Unsupported ASR_ENGINE={ASR_ENGINE!r}")

    logger.info("Model loaded in %.3fs", time.perf_counter() - started)
    return model


@app.on_event("startup")
async def startup_event():
    load_model()


@app.get("/health")
async def health():
    return {
        "status": "healthy",
        "engine": ASR_ENGINE,
        "paraformer_model": PARAFORMER_MODEL,
        "whisper_model": WHISPER_MODEL,
    }


@app.get("/")
async def root():
    return {"status": "ok", "engine": ASR_ENGINE}


def _transcribe_paraformer(path: str):
    res = load_model().generate(input=path, batch_size_s=PARAFORMER_BATCH_SIZE_S)
    text = res[0].get("text", "") if res else ""
    return {
        "text": text.replace(" ", ""),
        "language": "zh",
        "segments": [],
    }


def _transcribe_whisper(path: str, language: Optional[str]):
    lang = language.strip() if language and language.strip() else None
    segments, info = load_model().transcribe(
        path,
        language=lang,
        beam_size=WHISPER_BEAM_SIZE,
        vad_filter=WHISPER_VAD_FILTER,
    )
    convert = _simplifier()
    seg_list = []
    texts = []
    for segment in segments:
        text = convert(segment.text)
        seg_list.append(
            {
                "id": len(seg_list),
                "start": round(segment.start, 3),
                "end": round(segment.end, 3),
                "text": text,
            }
        )
        texts.append(text)
    return {
        "text": "".join(texts).strip(),
        "language": info.language,
        "duration": round(info.duration, 3),
        "segments": seg_list,
    }


async def _handle_upload(
    file: UploadFile,
    language: Optional[str],
    response_format: str,
):
    suffix = os.path.splitext(file.filename or "")[1] or ".wav"
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
        content = await file.read()
        tmp.write(content)
        tmp_path = tmp.name

    started = time.perf_counter()
    try:
        if ASR_ENGINE in {"paraformer", "paraformer-cpu", "funasr"}:
            result = _transcribe_paraformer(tmp_path)
        else:
            result = _transcribe_whisper(tmp_path, language)
        result["elapsed"] = round(time.perf_counter() - started, 3)

        if response_format == "text":
            return PlainTextResponse(result["text"])
        if response_format == "verbose_json":
            return {
                "task": "transcribe",
                "language": result.get("language", "zh"),
                "duration": result.get("duration"),
                "text": result["text"],
                "segments": result.get("segments", []),
                "elapsed": result["elapsed"],
            }
        return {
            "text": result["text"],
            "language": result.get("language", "zh"),
            "elapsed": result["elapsed"],
        }
    finally:
        try:
            os.unlink(tmp_path)
        except FileNotFoundError:
            pass


@app.post("/v1/audio/transcriptions")
async def transcribe_v1(
    file: UploadFile = File(...),
    language: Optional[str] = Form(None),
    response_format: Optional[str] = Form("json"),
    model: Optional[str] = Form(None),
    temperature: Optional[float] = Form(None),
    prompt: Optional[str] = Form(None),
):
    try:
        return await _handle_upload(file, language, response_format or "json")
    except Exception as exc:
        logger.exception("ASR failed")
        raise HTTPException(status_code=500, detail=str(exc))


@app.post("/api/v1/audio/transcriptions")
async def transcribe_api(
    file: UploadFile = File(...),
    language: Optional[str] = Form(None),
    response_format: Optional[str] = Form("json"),
    model: Optional[str] = Form(None),
    temperature: Optional[float] = Form(None),
    prompt: Optional[str] = Form(None),
):
    return await transcribe_v1(file, language, response_format, model, temperature, prompt)