"""RunPod serverless worker for the docling parse path.

One job in, one status out. The input payload carries three presigned S3
URLs; the worker holds no AWS credentials — the URLs carry the authorization
in all three directions:

- `pdf_url` — GET, the source PDF.
- `artifact_urls` — PUTs keyed `document` / `confidence` / `cells`, one per
  parse artifact, so the destination prefix holds the exact files
  `quber fuse --artifacts-dir` loads.
- `marker_url` — PUT, receives a small completion JSON written strictly
  after every artifact. Its existence guarantees the artifacts are
  complete; an S3 event notification on this key triggers the downstream
  workflow. On failure it carries `status: failed` and the failing stage,
  so failures ride the same channel as successes.
- `log_url` — PUT, optional. When present, the worker mirrors its log to
  this key every few seconds while the job runs and once more when it
  ends, so the log survives the worker and can be tailed near-live by
  anything that can read the key. Absent, logging behavior is unchanged.

The handler also returns the marker payload as the job output, and returns
it under `error` on failure so the job is marked FAILED for the caller's
deadline check.

A second task rides the same contract. A payload with `task: embed` carries
`chunks_url` (GET, a JSON object with a `texts` list), `vectors_url` (PUT,
the embeddings in the same order under `vectors`), `marker_url` and the
optional `log_url`. The embedding model is the one the playground queries
with, bge-large-en-v1.5, loaded once per worker and kept, so a document's
thousands of chunks embed on the GPU while a single question still embeds
on the app's CPU and both land in one vector space.

A startup fitness check makes a worker without a visible GPU exit before
it can accept jobs, so jobs queue for a healthy worker instead of failing
on a broken one.
"""

from __future__ import annotations

import inspect
import json
import logging
import tempfile
import threading
import time
import traceback
from pathlib import Path
from typing import Any
from urllib.parse import urlparse

import httpx
from loguru import logger

import runpod
from quber.core.parsers import parser_for_preset
from quber.core.parsers.result import ParseResult

TRANSFER_TIMEOUT_SECONDS = 300
LOG_PUSH_INTERVAL_SECONDS = 5
LOG_LINE_FORMAT = "{time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | {name}:{function}:{line} - {message}"


class LogShipper:
    """Mirror one job's log to a presigned S3 PUT URL, near-live.

    Registered as a loguru sink for the life of a job. A daemon thread PUTs
    the full accumulated text every LOG_PUSH_INTERVAL_SECONDS when new lines
    have arrived; close() pushes once more synchronously so the final lines
    land on success and failure alike. Every push carries the whole log and
    overwrites the same key, so a failed or missed push loses nothing — the
    next one carries everything.
    """

    def __init__(self, url: str) -> None:
        self.url = url
        self.lines: list[str] = []
        self.lock = threading.Lock()
        self.dirty = threading.Event()
        self.stopped = threading.Event()
        self.pump = threading.Thread(target=self.run, name="log-shipper", daemon=True)
        self.pump.start()

    def write(self, message: str) -> None:
        with self.lock:
            self.lines.append(message.rstrip("\n"))
        self.dirty.set()

    def close(self) -> None:
        self.stopped.set()
        self.pump.join(timeout=LOG_PUSH_INTERVAL_SECONDS)
        self.push()

    def run(self) -> None:
        while not self.stopped.wait(LOG_PUSH_INTERVAL_SECONDS):
            if self.dirty.is_set():
                self.dirty.clear()
                self.push()

    def push(self) -> None:
        with self.lock:
            body = "\n".join(self.lines).encode("utf-8")
        try:
            httpx.put(self.url, content=body, timeout=TRANSFER_TIMEOUT_SECONDS).raise_for_status()
        except Exception as exc:  # noqa: BLE001 — a log push must never fail the job
            self.dirty.set()
            with self.lock:
                self.lines.append(f"log push failed: {exc}")


class StdlibLogBridge(logging.Handler):
    """Route stdlib logging records into loguru so docling's output reaches every sink."""

    def emit(self, record: logging.LogRecord) -> None:
        try:
            level: str | int = logger.level(record.levelname).name
        except ValueError:
            level = record.levelno
        frame, depth = inspect.currentframe(), 0
        while frame and (depth == 0 or frame.f_code.co_filename == logging.__file__):
            frame = frame.f_back
            depth += 1
        logger.opt(depth=depth, exception=record.exc_info).log(level, record.getMessage())


def init_stdlib_log_bridge() -> None:
    logging.basicConfig(handlers=[StdlibLogBridge()], level=logging.INFO, force=True)
    # httpx logs every request at INFO, and each log push is itself an httpx
    # request: routed into the shipper, that is one new line per push, which
    # re-marks the log dirty and pushes again forever. The logged URL also
    # carries the presigned signature — credential material that must not be
    # written into the log object it authorizes.
    logging.getLogger("httpx").setLevel(logging.WARNING)


def download_pdf(url: str, destination: Path) -> None:
    with httpx.stream("GET", url, timeout=TRANSFER_TIMEOUT_SECONDS) as response:
        response.raise_for_status()
        with destination.open("wb") as sink:
            for chunk in response.iter_bytes():
                sink.write(chunk)


def upload_json(url: str, payload: dict[str, Any]) -> None:
    # No Content-Type header: a presigned PUT only accepts headers that were
    # signed, and the caller signs none.
    response = httpx.put(
        url,
        content=json.dumps(payload, default=str).encode("utf-8"),
        timeout=TRANSFER_TIMEOUT_SECONDS,
    )
    response.raise_for_status()


def document_base(pdf_url: str, provided: str | None) -> str:
    if provided:
        return provided
    stem = Path(urlparse(pdf_url).path).stem
    return stem or "document"


def upload_artifacts(result: ParseResult, workdir: Path, base: str, artifact_urls: dict[str, str]) -> None:
    """Upload the document/confidence/cells artifacts, each to its own URL.

    The three files land exactly as `ParseResult.save` writes them, so the
    S3 prefix matches what `quber fuse --artifacts-dir` loads — no
    re-bundling anywhere downstream.
    """
    artifacts = result.save(workdir, base)
    missing = sorted(set(artifacts) - set(artifact_urls))
    if missing:
        raise KeyError(f"artifact_urls missing {missing}")
    for name, path in artifacts.items():
        response = httpx.put(
            artifact_urls[name],
            content=path.read_bytes(),
            timeout=TRANSFER_TIMEOUT_SECONDS,
        )
        response.raise_for_status()


def summarize(result: ParseResult, base: str, started: float, parse_seconds: float) -> dict[str, Any]:
    document = result.document
    return {
        "document": base,
        "pages": len(document.pages),
        "tables": len(document.tables),
        "table_cells": sum(len(table.data.table_cells) for table in document.tables),
        "texts": len(document.texts),
        "parse_seconds": round(parse_seconds, 1),
        "total_seconds": round(time.time() - started, 1),
    }


_embedder: Any = None


def embedder() -> Any:
    """The embedding model, loaded on first use and kept for the worker's life."""
    global _embedder
    if _embedder is None:
        from quber.db.embeddings import EmbeddingService

        _embedder = EmbeddingService(provider="local", device="cuda")
    return _embedder


def download_json(url: str) -> Any:
    response = httpx.get(url, timeout=TRANSFER_TIMEOUT_SECONDS)
    response.raise_for_status()
    return response.json()


def embed_handler(payload: dict[str, Any]) -> dict[str, Any]:
    marker_url = payload.get("marker_url")
    log_url = payload.get("log_url")
    shipper = LogShipper(log_url) if log_url else None
    sink_id = logger.add(shipper.write, level="INFO", format=LOG_LINE_FORMAT) if shipper else None
    started = time.time()
    stage = "validate"
    try:
        chunks_url = payload["chunks_url"]
        vectors_url = payload["vectors_url"]
        if not marker_url:
            raise KeyError("marker_url")
        base = payload.get("base", "document")

        stage = "gpu_check"
        import torch

        if not torch.cuda.is_available():
            raise RuntimeError("no CUDA device visible on this worker")
        logger.info("GPU: {}", torch.cuda.get_device_name(0))

        stage = "download"
        texts = download_json(chunks_url)["texts"]
        logger.info("Embedding {} chunks for {}", len(texts), base)

        stage = "embed"
        embed_started = time.time()
        vectors = embedder().embed(texts, batch_size=64) if texts else []
        embed_seconds = time.time() - embed_started

        stage = "upload_results"
        upload_json(vectors_url, {"vectors": [list(map(float, v)) for v in vectors]})

        marker = {
            "status": "complete",
            "document": base,
            "chunks": len(texts),
            "embed_seconds": round(embed_seconds, 1),
            "total_seconds": round(time.time() - started, 1),
        }
        stage = "upload_marker"
        upload_json(marker_url, marker)
        logger.info("Embedded {}: {}", base, marker)
        return marker
    except Exception as exc:  # noqa: BLE001 -- every failure must reach the marker
        failure = {
            "status": "failed",
            "stage": stage,
            "error": f"{type(exc).__name__}: {exc}",
            "traceback": "".join(traceback.format_exception(exc))[-4000:],
            "total_seconds": round(time.time() - started, 1),
        }
        logger.exception("Embed job failed at {}", stage)
        if marker_url and stage != "validate":
            try:
                upload_json(marker_url, failure)
            except Exception as marker_exc:  # noqa: BLE001
                logger.error("Marker upload failed too: {}", marker_exc)
        # The platform accepts `error` as a string only; a dict is refused at
        # job-done with a 400 and the job then shows as completed with no output.
        job_output: dict[str, Any] = {"error": json.dumps(failure)}
        if stage == "gpu_check":
            job_output["refresh_worker"] = True
        return job_output
    finally:
        if shipper is not None:
            logger.remove(sink_id)
            shipper.close()


def handler(job: dict[str, Any]) -> dict[str, Any]:
    payload = job.get("input") or {}
    if payload.get("task") == "embed":
        return embed_handler(payload)
    marker_url = payload.get("marker_url")
    log_url = payload.get("log_url")
    shipper = LogShipper(log_url) if log_url else None
    sink_id = logger.add(shipper.write, level="INFO", format=LOG_LINE_FORMAT) if shipper else None
    started = time.time()
    stage = "validate"
    try:
        pdf_url = payload["pdf_url"]
        artifact_urls = payload["artifact_urls"]
        if not marker_url:
            raise KeyError("marker_url")
        base = document_base(pdf_url, payload.get("base"))

        stage = "gpu_check"
        import torch

        if not torch.cuda.is_available():
            # Some workers come up with no visible CUDA device (observed
            # live on an A5000 worker). docling's AUTO device would
            # silently fall back to CPU and crawl until it dies — fail
            # fast instead. The startup fitness check should have caught
            # this before any job; this is the in-job backstop.
            raise RuntimeError("no CUDA device visible on this worker")
        logger.info("GPU: {}", torch.cuda.get_device_name(0))

        with tempfile.TemporaryDirectory() as tmp:
            workdir = Path(tmp)
            pdf_path = workdir / f"{base}.pdf"

            stage = "download"
            download_pdf(pdf_url, pdf_path)
            logger.info("Downloaded {} ({} bytes)", base, pdf_path.stat().st_size)

            stage = "parse"
            parser = parser_for_preset("tuned-financial")
            parse_started = time.time()
            result = parser.parse(pdf_path)
            parse_seconds = time.time() - parse_started

            stage = "upload_results"
            upload_artifacts(result, workdir, base, artifact_urls)

        marker = {"status": "complete", **summarize(result, base, started, parse_seconds)}
        stage = "upload_marker"
        upload_json(marker_url, marker)
        logger.info("Completed {}: {}", base, marker)
        return marker

    except Exception as exc:  # noqa: BLE001 — every failure must reach the marker
        failure = {
            "status": "failed",
            "stage": stage,
            "error": f"{type(exc).__name__}: {exc}",
            # The full cause chain, tail-truncated. docling wraps the real
            # exception in a generic "Pipeline ... failed" RuntimeError, and
            # the marker is often the only diagnostic that leaves the worker.
            "traceback": "".join(traceback.format_exception(exc))[-4000:],
            "total_seconds": round(time.time() - started, 1),
        }
        logger.exception("Job failed at {}", stage)
        if marker_url and stage != "validate":
            try:
                upload_json(marker_url, failure)
            except Exception as marker_exc:  # noqa: BLE001
                logger.error("Marker upload failed too: {}", marker_exc)
        # `error` marks the job FAILED so the caller's deadline check sees a
        # loud status even when the marker never landed. A worker whose GPU
        # is missing is broken hardware-side: refresh_worker retires it so
        # the retry lands elsewhere.
        # The platform accepts `error` as a string only; a dict is refused at
        # job-done with a 400 and the job then shows as completed with no output.
        job_output: dict[str, Any] = {"error": json.dumps(failure)}
        if stage == "gpu_check":
            job_output["refresh_worker"] = True
        return job_output
    finally:
        if shipper is not None:
            logger.remove(sink_id)
            shipper.close()


def gpu_fitness() -> None:
    """Startup health gate: a worker with no visible GPU must never serve.

    Registered as an SDK fitness check, which runs before the worker starts
    accepting jobs; on failure the SDK exits the process (exit code 1) and
    the worker is marked unhealthy, so jobs queue for a healthy worker
    instead of failing on a broken one.
    """
    import torch

    if not torch.cuda.is_available():
        raise RuntimeError("no CUDA device visible on this worker")


if __name__ == "__main__":
    init_stdlib_log_bridge()
    runpod.serverless.register_fitness_check(gpu_fitness)
    runpod.serverless.start({"handler": handler})