from __future__ import annotations

import json
import uuid
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Annotated, Optional

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import JSONResponse, StreamingResponse

from audio.asr import create_engine
from audio.config import Settings, get_settings
from audio.metrics import gpu_memory
from audio.pipeline import AudioPipeline
from audio.queue import GpuSerialQueue, QueueFullError
from audio.translation import TranslationClient

_queue: GpuSerialQueue | None = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    global _queue
    settings = get_settings()
    asr = create_engine(settings)
    pipeline = AudioPipeline(settings, asr, TranslationClient(settings))
    _queue = GpuSerialQueue(
        pipeline,
        max_queue=settings.max_queue,
        gpu_concurrency=asr.replicas,
    )
    app.state.settings = settings
    yield


app = FastAPI(title="Translation Audio Service", lifespan=lifespan)


def queue() -> GpuSerialQueue:
    if _queue is None:
        raise RuntimeError("service not started")
    return _queue


def _save_upload(upload_dir: str, filename: str | None, data: bytes) -> Path:
    suffix = Path(filename or "audio.bin").suffix or ".bin"
    dest = Path(upload_dir) / f"{uuid.uuid4().hex}{suffix}"
    dest.parent.mkdir(parents=True, exist_ok=True)
    dest.write_bytes(data)
    return dest


@app.get("/health")
def health() -> dict:
    snap = gpu_memory()
    settings: Settings = get_settings()
    engine = queue().pipeline.asr
    return {
        "ok": True,
        "layer": "audio",
        "asr_backend": getattr(engine, "backend", settings.asr_backend),
        "asr_model": getattr(engine, "model_name", settings.asr_model),
        "device": getattr(engine, "device", settings.asr_device),
        "compute_type": getattr(engine, "compute_type", settings.asr_compute_type),
        "batch_size": getattr(engine, "batch_size", settings.asr_batch_size),
        "replicas": getattr(engine, "replicas", settings.asr_replicas),
        "gpu": {
            "name": snap.name,
            "used_mb": snap.used_mb,
            "total_mb": snap.total_mb,
            "free_mb": snap.free_mb,
            "compute_capability": (
                f"{snap.compute_major}.{snap.compute_minor}"
                if snap.compute_major is not None
                else None
            ),
        },
    }


@app.post("/translate")
async def translate_media(
    file: Annotated[UploadFile, File()],
    target_language: Annotated[str, Form()] = "en",
    source_language: Annotated[Optional[str], Form()] = None,
    translate: Annotated[bool, Form()] = True,
) -> JSONResponse:
    settings: Settings = get_settings()
    dest = _save_upload(settings.upload_dir, file.filename, await file.read())
    try:
        result = await queue().run(
            source=dest,
            target_language=target_language,
            source_language=source_language or None,
            translate=translate,
        )
    except QueueFullError as exc:
        raise HTTPException(status_code=429, detail=str(exc)) from exc
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    return JSONResponse(result)


@app.post("/translate/stream")
async def translate_stream(
    file: Annotated[UploadFile, File()],
    source_language: Annotated[Optional[str], Form()] = None,
) -> StreamingResponse:
    settings: Settings = get_settings()
    dest = _save_upload(settings.upload_dir, file.filename, await file.read())

    async def events():
        try:
            async for item in queue().stream(
                source=dest,
                source_language=source_language or None,
            ):
                yield f"data: {json.dumps(item, ensure_ascii=False)}\n\n"
        except QueueFullError as exc:
            yield f"data: {json.dumps({'type': 'error', 'detail': str(exc)})}\n\n"
        except Exception as exc:
            yield f"data: {json.dumps({'type': 'error', 'detail': str(exc)})}\n\n"

    return StreamingResponse(events(), media_type="text/event-stream")