from __future__ import annotations

from contextlib import asynccontextmanager
import asyncio
from pathlib import Path
import uuid

from aiortc import RTCPeerConnection, RTCSessionDescription
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import FileResponse, JSONResponse
from pydantic import BaseModel, Field

from dictation.config import DictationConfig
from dictation.session import DictationSession
from dictation.transcriber import FasterWhisperTranscriber, Transcriber


class Offer(BaseModel):
    sdp: str = Field(min_length=1, max_length=256_000)
    type: str


class DictationService:
    def __init__(
        self,
        config: DictationConfig,
        transcriber: Transcriber,
    ) -> None:
        self.config = config
        self.transcriber = transcriber
        self.sessions: dict[str, DictationSession] = {}
        self.ready = False
        self.startup_error: str | None = None
        self._lock = asyncio.Lock()

    async def start(self) -> None:
        try:
            await self.transcriber.warmup()
            self.ready = True
        except Exception as error:
            self.startup_error = str(error)
            raise

    async def accept_offer(self, offer: Offer) -> dict[str, str]:
        if offer.type != "offer":
            raise HTTPException(422, "SDP type must be offer")
        async with self._lock:
            if len(self.sessions) >= self.config.max_sessions:
                raise HTTPException(429, "Dictation session limit reached")
            peer = RTCPeerConnection()
            session_id = str(uuid.uuid4())
            session = DictationSession(
                session_id,
                peer,
                self.transcriber,
                self.config,
            )
            self.sessions[session_id] = session

        @peer.on("datachannel")
        def on_datachannel(channel) -> None:
            session.attach_channel(channel)

        @peer.on("track")
        def on_track(track) -> None:
            session.attach_track(track)

        @peer.on("connectionstatechange")
        async def on_connectionstatechange() -> None:
            if peer.connectionState in {"failed", "closed"}:
                await self.close_session(session_id)

        try:
            await peer.setRemoteDescription(
                RTCSessionDescription(sdp=offer.sdp, type=offer.type)
            )
            answer = await peer.createAnswer()
            await peer.setLocalDescription(answer)
            return {
                "sdp": peer.localDescription.sdp,
                "type": peer.localDescription.type,
                "sessionId": session_id,
            }
        except Exception:
            await self.close_session(session_id)
            raise

    async def close_session(self, session_id: str) -> bool:
        async with self._lock:
            session = self.sessions.pop(session_id, None)
        if not session:
            return False
        await session.close()
        return True

    async def close(self) -> None:
        for session_id in list(self.sessions):
            await self.close_session(session_id)
        close = getattr(self.transcriber, "close", None)
        if close:
            close()


def create_app(
    config: DictationConfig | None = None,
    transcriber: Transcriber | None = None,
) -> FastAPI:
    resolved_config = config or DictationConfig.from_environment()
    resolved_transcriber = transcriber or FasterWhisperTranscriber(
        resolved_config.model_dir,
        resolved_config.compute_type,
    )
    service = DictationService(resolved_config, resolved_transcriber)
    web_root = Path(__file__).parent / "web"

    @asynccontextmanager
    async def lifespan(app: FastAPI):
        app.state.dictation = service
        try:
            await service.start()
            yield
        finally:
            await service.close()

    app = FastAPI(
        title="Zenbu WebRTC Dictation",
        lifespan=lifespan,
    )

    @app.get("/")
    async def index():
        return FileResponse(web_root / "index.html")

    @app.get("/dictation.js")
    async def javascript():
        return FileResponse(web_root / "dictation.js")

    @app.get("/dictation.css")
    async def stylesheet():
        return FileResponse(web_root / "dictation.css")

    @app.get("/health")
    async def health(request: Request):
        current: DictationService = request.app.state.dictation
        status = 200 if current.ready else 503
        return JSONResponse(
            {
                "status": "ok" if current.ready else "unavailable",
                "model": str(current.config.model_dir),
                "computeType": current.config.compute_type,
                "activeSessions": len(current.sessions),
                "maxSessions": current.config.max_sessions,
                "error": current.startup_error,
            },
            status_code=status,
        )

    @app.post("/api/webrtc/offer")
    async def offer(body: Offer, request: Request):
        current: DictationService = request.app.state.dictation
        if not current.ready:
            raise HTTPException(503, "Dictation model is unavailable")
        return await current.accept_offer(body)

    @app.post("/api/webrtc/session/{session_id}/close")
    async def close_session(session_id: str, request: Request):
        current: DictationService = request.app.state.dictation
        await current.close_session(session_id)
        return {"closed": True}

    return app