"""Bounded HTTP job service for local checkpoint manifests."""

from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
import fcntl
import json
import logging
from pathlib import Path
import secrets
from threading import BoundedSemaphore
from uuid import UUID, uuid4

from fastapi import Depends, FastAPI, HTTPException, Request
from fastapi.responses import FileResponse
from fastapi.security import APIKeyHeader, HTTPAuthorizationCredentials, HTTPBearer
from pydantic import BaseModel, ConfigDict, Field

from .core import _kernel
from .files import prune_manifest


class JobRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")
    manifest: str = Field(min_length=1, max_length=1024)
    seed: int = Field(default=0, ge=0, lt=1 << 64)


class CheckpointJobRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")
    checkpoint: str = Field(default=".", min_length=1, max_length=1024)
    keep: int = Field(strict=True, ge=0, lt=1 << 63)
    moments_manifest: str | None = Field(default=None, min_length=1, max_length=1024)
    seed: int = Field(default=0, ge=0, lt=1 << 64)


def _write_state(directory, run_id, state, **fields):
    path = directory / "state.json"
    tmp = directory / ".state.json.tmp"
    tmp.write_text(json.dumps(dict(state=state, run_id=run_id, **fields)))
    tmp.replace(path)


def create_app(input_root, job_root, *, workers=1, queue_size=2, threads_per_job=1,
               max_weights=100_000_000, token=None, require_token=False,
               checkpoint_root=None):
    """Create a single-process HTTP service with configurable CPU and queue limits."""
    if min(workers, threads_per_job, max_weights) < 1 or queue_size < 0:
        raise ValueError("invalid worker, queue, thread, or weight limit")
    inputs = Path(input_root).resolve(strict=True)
    jobs = Path(job_root).resolve()
    checkpoints = Path(checkpoint_root).resolve(strict=True) if checkpoint_root is not None else None
    if jobs.is_relative_to(inputs) or inputs.is_relative_to(jobs):
        raise ValueError("input and job roots must be separate directories")

    @asynccontextmanager
    async def lifespan(app):
        jobs.mkdir(parents=True, exist_ok=True)
        with (jobs / ".server.lock").open("a+b") as lock:
            try:
                fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
            except BlockingIOError as exc:
                raise RuntimeError("another service already owns this job directory") from exc
            _kernel()
            app.state.run_id = uuid4().hex
            app.state.slots = BoundedSemaphore(workers + queue_size)
            with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="prune") as pool:
                app.state.pool = pool
                yield

    app = FastAPI(title="Tensor Pruning", version="1.0.0", lifespan=lifespan)

    bearer = HTTPBearer(auto_error=False)
    api_key = APIKeyHeader(name="X-Tensor-Pruning-Token", auto_error=False)

    def authenticate(credentials: HTTPAuthorizationCredentials | None = Depends(bearer),
                     space_token: str | None = Depends(api_key)):
        if not token:
            raise HTTPException(503, "The operator must configure the API token")
        supplied = space_token if space_token is not None else (credentials.credentials if credentials else "")
        if not secrets.compare_digest(supplied.encode(), token.encode()):
            raise HTTPException(401, "A valid bearer token is required",
                                headers={"WWW-Authenticate": "Bearer"})

    protected = [Depends(authenticate)] if token or require_token else []

    @app.get("/health")
    def health():
        return {"status": "ok"}

    def read_job(job_id, request):
        directory = jobs / str(job_id)
        try:
            state = json.loads((directory / "state.json").read_text())
        except FileNotFoundError as exc:
            raise HTTPException(404, "Unknown job") from exc
        run_id = state.pop("run_id")
        if run_id != request.app.state.run_id and state["state"] in {"queued", "running"}:
            state["state"] = "interrupted"
        return directory, state

    def enqueue(request, operation):
        state = request.app.state
        if not state.slots.acquire(blocking=False):
            raise HTTPException(429, "Pruning queue is full", headers={"Retry-After": "1"})
        job_id = uuid4()
        directory = jobs / str(job_id)

        def run():
            try:
                _write_state(directory, state.run_id, "running")
                result = operation(directory)
                _write_state(directory, state.run_id, "completed", result=result)
            except (ValueError, TypeError, OSError, MemoryError) as exc:
                _write_state(directory, state.run_id, "failed", error=str(exc))
            except Exception:
                logging.getLogger(__name__).exception("Pruning job failed")
                _write_state(directory, state.run_id, "failed", error="Internal pruning failure")
            finally:
                state.slots.release()

        try:
            directory.mkdir()
            _write_state(directory, state.run_id, "queued")
            state.pool.submit(run)
        except Exception:
            state.slots.release()
            raise
        return {"id": str(job_id), "state": "queued"}

    @app.post("/v1/jobs", status_code=202, dependencies=protected)
    def submit(body: JobRequest, request: Request):
        manifest = Path(body.manifest)
        if manifest.is_absolute() or not (inputs / manifest).resolve().is_relative_to(inputs):
            raise HTTPException(400, "Manifest must be inside the input root")
        return enqueue(request, lambda directory: prune_manifest(
            inputs, body.manifest, directory / "output", threads=threads_per_job,
            seed=body.seed, max_weights=max_weights))

    if checkpoints is not None:
        from .checkpoint import prune_checkpoint

        @app.post("/v1/checkpoint-jobs", status_code=202, dependencies=protected)
        def submit_checkpoint(body: CheckpointJobRequest, request: Request):
            relative = Path(body.checkpoint)
            checkpoint = (checkpoints / relative).resolve()
            if relative.is_absolute() or not checkpoint.is_relative_to(checkpoints):
                raise HTTPException(400, "Checkpoint must be inside the model mount")
            if not checkpoint.is_dir():
                raise HTTPException(404, "Mounted checkpoint directory not found")
            return enqueue(request, lambda directory: prune_checkpoint(
                checkpoint, directory / "output", keep=body.keep, workspace=directory,
                max_weights=max_weights, input_root=inputs,
                moments_manifest=body.moments_manifest, threads=threads_per_job, seed=body.seed))

    @app.get("/v1/jobs/{job_id}", dependencies=protected)
    def status(job_id: UUID, request: Request):
        _, state = read_job(job_id, request)
        return {"id": str(job_id), **state}

    @app.get("/v1/jobs/{job_id}/masks/{index}", dependencies=protected,
             response_class=FileResponse, responses={
                 200: {"content": {"application/octet-stream": {
                     "schema": {"type": "string", "format": "binary"}
                 }}}
             })
    def mask(job_id: UUID, index: int, request: Request):
        directory, state = read_job(job_id, request)
        if state["state"] != "completed":
            raise HTTPException(409, "The job has no completed masks")
        projections = state["result"]["projections"]
        if not 0 <= index < len(projections):
            raise HTTPException(404, "Unknown projection index")
        filename = projections[index]["mask"]
        return FileResponse(directory / "output" / filename,
                            filename=filename, media_type="application/octet-stream")

    return app