from __future__ import annotations

import asyncio
import json
import re
from contextlib import asynccontextmanager, suppress

import httpx
from fastapi import FastAPI, Header, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse

from .downstream import DownstreamError, LocalMediaClient
from .models import JobAccepted, MediaJobRequest, MediaJobView
from .settings import MediaSettings
from .store import (
    TERMINAL,
    IdempotencyConflict,
    JobCapacityExceeded,
    JobNotFound,
    JobStore,
    ReplayGap,
)


def _bounded_header(value: str | None, label: str) -> str:
    normalized = (value or "").strip()
    if not normalized or len(normalized) > 240:
        raise HTTPException(status_code=400, detail=f"A bounded {label} is required")
    return normalized


_SAFE_RESOURCE_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.:-]{0,239}$")


def _lease_id(value: str | None) -> str:
    normalized = _bounded_header(value, "X-Model-Lease")
    if not _SAFE_RESOURCE_ID.fullmatch(normalized):
        raise HTTPException(status_code=400, detail="X-Model-Lease has an invalid format")
    return normalized


class MediaRuntime:
    def __init__(
        self,
        settings: MediaSettings,
        *,
        client: httpx.AsyncClient | None = None,
    ) -> None:
        self.settings = settings
        self.store = JobStore(
            max_active=settings.max_active_jobs,
            max_records=settings.max_records,
            max_events=settings.max_events_per_job,
            max_subscribers=settings.max_subscribers,
        )
        self.downstream = LocalMediaClient(settings, client=client)

    async def close(self) -> None:
        await self.store.cancel_all()
        await self.downstream.close()

    async def create(
        self,
        request: MediaJobRequest,
        *,
        lease_id: str,
        idempotency_key: str,
    ) -> tuple[MediaJobView, bool]:
        endpoint = self.settings.endpoint(request.kind)
        if not endpoint:
            raise HTTPException(
                status_code=503,
                detail=f"No explicit local {request.kind} endpoint is configured",
            )
        options_size = len(
            json.dumps(request.options, sort_keys=True, separators=(",", ":")).encode("utf-8")
        )
        if options_size > self.settings.max_options_bytes:
            raise HTTPException(status_code=413, detail="Media options are too large")
        try:
            replay = await self.store.replay(
                request, lease_id=lease_id, idempotency_key=idempotency_key
            )
        except IdempotencyConflict as exc:
            raise HTTPException(status_code=409, detail=str(exc)) from None
        if replay is not None:
            return replay, False
        try:
            await self.downstream.verify_lease(lease_id, request.kind)
        except DownstreamError as exc:
            raise HTTPException(status_code=409, detail=str(exc)) from None
        try:
            view, created = await self.store.create(
                request,
                lease_id=lease_id,
                endpoint=endpoint,
                idempotency_key=idempotency_key,
            )
        except IdempotencyConflict as exc:
            raise HTTPException(status_code=409, detail=str(exc)) from None
        except JobCapacityExceeded as exc:
            raise HTTPException(status_code=503, detail=str(exc)) from None
        if created:
            task = asyncio.create_task(self._run(view.job_id), name=f"media-{view.job_id}")
            await self.store.set_task(view.job_id, task)
        return view, created

    async def _run(self, job_id: str) -> None:
        try:
            await asyncio.wait_for(self._run_inner(job_id), timeout=self.settings.job_timeout_s)
        except asyncio.CancelledError:
            await self.store.terminal(job_id, "cancelled", error_code="cancelled")
            raise
        except TimeoutError:
            await self.store.terminal(job_id, "failed", error_code="job_timeout")
        except DownstreamError:
            await self.store.terminal(job_id, "failed", error_code="downstream_error")
        except Exception:
            await self.store.terminal(job_id, "failed", error_code="internal_error")

    async def _run_inner(self, job_id: str) -> None:
        request, lease_id, endpoint, _downstream_id, _task = await self.store.runtime(job_id)
        await self.store.progress(job_id, "job.started", stage="submitted")
        result = await self.downstream.submit(
            request,
            endpoint=endpoint,
            lease_id=lease_id,
            idempotency_key=f"media:{job_id}",
        )
        if result.downstream_id:
            await self.store.set_downstream(job_id, result.downstream_id)
        if result.artifacts is not None:
            await self.store.terminal(
                job_id,
                "completed",
                artifacts=[item.model_dump(mode="json") for item in result.artifacts],
                summary=result.summary,
            )
            return
        if not result.downstream_id:
            raise DownstreamError("downstream job has no identifier")

        async for event_type, payload in self.downstream.events(
            endpoint=endpoint, downstream_id=result.downstream_id
        ):
            normalized = event_type.casefold()
            if normalized in {"job.completed", "media.job.completed", "completed"}:
                artifacts = self.downstream.artifacts(payload.get("artifacts"))
                summary = payload.get("summary") or payload.get("output")
                await self.store.terminal(
                    job_id,
                    "completed",
                    artifacts=[item.model_dump(mode="json") for item in artifacts],
                    summary=summary if isinstance(summary, str) else None,
                )
                return
            if normalized in {"job.failed", "media.job.failed", "failed", "error"}:
                await self.store.terminal(job_id, "failed", error_code="downstream_failed")
                return
            if normalized in {"job.cancelled", "media.job.cancelled", "cancelled"}:
                await self.store.terminal(job_id, "cancelled", error_code="downstream_cancelled")
                return
            progress = payload.get("progress")
            if not isinstance(progress, (int, float)) or isinstance(progress, bool):
                progress = None
            stage = payload.get("stage")
            await self.store.progress(
                job_id,
                "job.progress",
                progress=float(progress) if progress is not None else None,
                stage=stage if isinstance(stage, str) else normalized,
            )
        await self.store.terminal(
            job_id, "failed", error_code="stream_ended_without_terminal_event"
        )

    async def cancel(self, job_id: str) -> MediaJobView:
        view = await self.store.view(job_id)
        if view.status in TERMINAL:
            raise HTTPException(status_code=409, detail="Media job is already terminal")
        await self.store.progress(job_id, "job.cancelling", stage="cancel_requested")
        _request, _lease_id, endpoint, downstream_id, task = await self.store.runtime(job_id)
        cancel_error = False
        if downstream_id:
            try:
                await self.downstream.cancel(endpoint=endpoint, downstream_id=downstream_id)
            except DownstreamError:
                cancel_error = True
        if task and task is not asyncio.current_task() and not task.done():
            task.cancel()
            with suppress(asyncio.CancelledError):
                await task
        await self.store.terminal(
            job_id,
            "cancelled",
            error_code="downstream_cancel_unavailable" if cancel_error else "cancelled",
        )
        return await self.store.view(job_id)


def create_app(
    settings: MediaSettings | None = None,
    *,
    client: httpx.AsyncClient | None = None,
) -> FastAPI:
    settings = (settings or MediaSettings.from_env()).validated()
    runtime = MediaRuntime(settings, client=client)

    @asynccontextmanager
    async def lifespan(_: FastAPI):
        yield
        await runtime.close()

    app = FastAPI(
        title="OMNI local media job gateway",
        version="1.0",
        docs_url=None,
        redoc_url=None,
        openapi_url=None,
        lifespan=lifespan,
    )
    app.state.runtime = runtime

    @app.middleware("http")
    async def bound_requests(request: Request, call_next):
        if request.headers.get("transfer-encoding"):
            return JSONResponse(status_code=411, content={"detail": "Chunked requests are denied"})
        raw_length = request.headers.get("content-length")
        if raw_length:
            try:
                length = int(raw_length)
            except ValueError:
                return JSONResponse(status_code=400, content={"detail": "Invalid Content-Length"})
            if length < 0 or length > settings.max_request_bytes:
                return JSONResponse(status_code=413, content={"detail": "Request body is too large"})
        return await call_next(request)

    @app.get("/health")
    @app.get("/healthz")
    async def health() -> dict[str, object]:
        return {
            "status": "ready",
            "service": "media-api",
            "model_downloads": "disabled",
            "auto_start": settings.auto_start,
            "configured": {
                "video": bool(settings.video_endpoint),
                "image": bool(settings.image_endpoint),
            },
        }

    @app.post("/v1/jobs", response_model=JobAccepted, status_code=202)
    async def create_job(
        request: MediaJobRequest,
        idempotency_key: str | None = Header(default=None, alias="Idempotency-Key"),
        lease_id: str | None = Header(default=None, alias="X-Model-Lease"),
    ) -> JobAccepted:
        key = _bounded_header(idempotency_key, "Idempotency-Key")
        lease = _lease_id(lease_id)
        view, created = await runtime.create(
            request, lease_id=lease, idempotency_key=key
        )
        return JobAccepted(job_id=view.job_id, replayed=not created)

    @app.get("/v1/jobs/{job_id}", response_model=MediaJobView)
    async def get_job(job_id: str) -> MediaJobView:
        try:
            return await runtime.store.view(job_id)
        except JobNotFound:
            raise HTTPException(status_code=404, detail="Media job not found") from None

    @app.get("/v1/jobs/{job_id}/events")
    async def events(
        job_id: str,
        last_event_id: str | None = Header(default=None, alias="Last-Event-ID"),
    ) -> StreamingResponse:
        try:
            after = int(last_event_id or "0")
        except ValueError:
            raise HTTPException(status_code=400, detail="Last-Event-ID must be an integer") from None
        if after < 0:
            raise HTTPException(status_code=400, detail="Last-Event-ID must be non-negative")
        try:
            await runtime.store.validate_after(job_id, after)
        except JobNotFound:
            raise HTTPException(status_code=404, detail="Media job not found") from None
        except ReplayGap as exc:
            raise HTTPException(status_code=409, detail=str(exc)) from None

        async def generate():
            async for event in runtime.store.stream(
                job_id, after, settings.subscriber_heartbeat_s
            ):
                yield event.sse() if event is not None else ": keepalive\n\n"

        return StreamingResponse(
            generate(),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
        )

    @app.post("/v1/jobs/{job_id}/cancel", response_model=MediaJobView, status_code=202)
    async def cancel_job(job_id: str) -> MediaJobView:
        try:
            return await runtime.cancel(job_id)
        except JobNotFound:
            raise HTTPException(status_code=404, detail="Media job not found") from None

    return app