"""The GPU stages of an upload, run on the RunPod worker.

The hosted playground has no GPU. Two stages of an upload need one: the
docling parse of the document, and the embedding of its chunks at
ingestion. Both go to the RunPod serverless worker under the contract the
inbound-drop Lambda already uses with it: the job payload carries presigned
S3 URLs for everything the worker reads and writes, the worker holds no
credentials, and a completion marker written strictly after the outputs says
the outputs are whole. The files land in the document's own prefix in the
extracts bucket, so a parse's artifacts and an embedding's vectors sit
beside everything else the document owns.

``configured`` says whether this host sends work there at all. It needs
RUNPOD_ENDPOINT_ID, RUNPOD_API_KEY and QUBER_PLAYGROUND_ARTIFACTS_URI all set.
Without them, the app runs the docling parse on the local GPU in a
`quber fuse` subprocess, and the chunks embed locally in the ingest process.
"""

from __future__ import annotations

import json
import time
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional
from urllib.parse import urlparse

import boto3
import httpx
from botocore.config import Config
from loguru import logger

from quber.settings import get_settings

URL_EXPIRY_SECONDS = 4 * 3600
POLL_SECONDS = 5
LOG_MIRROR_SECONDS = 10
#: A job on a worker freshly placed on a new host pays a ten-gigabyte image
#: pull before it runs: eleven minutes measured once, over twenty on another
#: night. Either task can be the one that meets a fresh host.
PARSE_DEADLINE_SECONDS = 35 * 60
EMBED_DEADLINE_SECONDS = 35 * 60

#: The worker's artifact names, each with the file the name lands in. The
#: worker keys its upload URLs by the first and writes the second.
PARSE_ARTIFACTS = {"document": "docling", "confidence": "confidence", "cells": "cells"}

Log = Callable[[str], None]


def configured() -> bool:
    p = get_settings().playground
    return bool(p.runpod_endpoint_id and p.runpod_api_key and p.artifacts_uri)


def _bucket_and_prefix(doc_key: str) -> tuple[str, str]:
    uri = get_settings().playground.artifacts_uri or ""
    parsed = urlparse(uri)
    return parsed.netloc, parsed.path.strip("/") + "/" + doc_key + "/"


def _s3():
    # Signature version 4, explicitly: the extracts bucket is KMS-encrypted
    # and S3 refuses a version 2 presigned URL against it with a bare 400,
    # which is what a client without a pinned region produces for us-east-1.
    return boto3.client(
        "s3",
        region_name=get_settings().s3_cache.region or "us-east-1",
        config=Config(signature_version="s3v4"),
    )


def _presign(method: str, bucket: str, key: str) -> str:
    return _s3().generate_presigned_url(
        method, Params={"Bucket": bucket, "Key": key}, ExpiresIn=URL_EXPIRY_SECONDS
    )


def _submit(payload: Dict[str, Any]) -> str:
    p = get_settings().playground
    response = httpx.post(
        f"https://api.runpod.ai/v2/{p.runpod_endpoint_id}/run",
        json={"input": payload},
        headers={"Authorization": f"Bearer {p.runpod_api_key}"},
        timeout=30,
    )
    response.raise_for_status()
    job = response.json()
    logger.info("runpod job {} ({})", job.get("id"), job.get("status"))
    return job["id"]


def _read_if_fresh(bucket: str, key: str, since: float) -> Optional[bytes]:
    """The object's bytes, or None when it does not exist or predates ``since``.

    Nothing under a document's prefix is ever deleted, so a previous run's
    marker and log stay where they are; a run tells its own apart by the
    object's write time against the moment it submitted the job."""
    s3 = _s3()
    try:
        response = s3.get_object(Bucket=bucket, Key=key)
    except s3.exceptions.NoSuchKey:
        return None
    if response["LastModified"].timestamp() < since:
        return None
    return response["Body"].read()


def _read_json_if_fresh(bucket: str, key: str, since: float) -> Optional[Dict[str, Any]]:
    body = _read_if_fresh(bucket, key, since)
    return json.loads(body) if body is not None else None


def _read_text_if_fresh(bucket: str, key: str, since: float) -> Optional[str]:
    body = _read_if_fresh(bucket, key, since)
    return body.decode("utf-8", "replace") if body is not None else None


def _job_status(job_id: str) -> Dict[str, Any]:
    p = get_settings().playground
    response = httpx.get(
        f"https://api.runpod.ai/v2/{p.runpod_endpoint_id}/status/{job_id}",
        headers={"Authorization": f"Bearer {p.runpod_api_key}"},
        timeout=30,
    )
    response.raise_for_status()
    return response.json()


def _wait_for_marker(
    bucket: str, marker_key: str, log_key: Optional[str], deadline: float, log: Log, job_id: str, since: float
) -> Dict[str, Any]:
    """Poll until the marker exists, mirroring the worker's log into ``log``
    as it grows. The marker is the truth about the outputs; the job's own
    status is watched beside it so a job that ends without writing one, or
    that RunPod fails or cancels, is reported at once rather than at the
    deadline."""
    started = time.time()
    mirrored = 0
    last_mirror = 0.0
    ended_at: Optional[float] = None
    while True:
        marker = _read_json_if_fresh(bucket, marker_key, since)
        if marker is not None:
            if log_key:
                text = _read_text_if_fresh(bucket, log_key, since) or ""
                for line in text.splitlines()[mirrored:]:
                    log(line)
            if marker.get("status") != "complete":
                raise RuntimeError(
                    f"worker failed at {marker.get('stage', '?')}: {marker.get('error', 'no error recorded')}"
                )
            return marker
        elapsed = time.time() - started
        if elapsed > deadline:
            raise TimeoutError(f"no completion marker after {int(elapsed)}s ({marker_key})")
        status = _job_status(job_id)
        state = status.get("status")
        if state in ("FAILED", "CANCELLED", "TIMED_OUT"):
            raise RuntimeError(
                f"worker job {job_id} {state}: {json.dumps(status.get('error') or status.get('output'))[:800]}"
            )
        if state == "COMPLETED":
            # The marker is written before the handler returns, so a completed
            # job with no marker a moment later returned without doing the work.
            ended_at = ended_at or time.time()
            if time.time() - ended_at > 3 * POLL_SECONDS:
                raise RuntimeError(
                    f"worker job {job_id} completed without a completion marker; output: "
                    f"{json.dumps(status.get('output'))[:800]}"
                )
        if log_key and time.time() - last_mirror >= LOG_MIRROR_SECONDS:
            last_mirror = time.time()
            text = _read_text_if_fresh(bucket, log_key, since)
            if text:
                lines = text.splitlines()
                for line in lines[mirrored:]:
                    log(line)
                mirrored = len(lines)
        time.sleep(POLL_SECONDS)


def parse(doc_key: str, workdir: Path, log: Log) -> None:
    """Parse the document on the worker and bring the three parse artifacts
    into ``workdir`` for fusion. The PDF must already be in the document's
    prefix as ``<doc_key>.pdf``."""
    bucket, prefix = _bucket_and_prefix(doc_key)
    marker_key = f"{prefix}{doc_key}.parse.complete.json"
    log_key = f"{prefix}{doc_key}.parse.log"
    s3 = _s3()
    # Clock skew between this host and S3 is absorbed by a minute's grace.
    since = time.time() - 60
    payload = {
        "base": doc_key,
        "pdf_url": _presign("get_object", bucket, f"{prefix}{doc_key}.pdf"),
        "artifact_urls": {
            name: _presign("put_object", bucket, f"{prefix}{doc_key}.{stem}.json")
            for name, stem in PARSE_ARTIFACTS.items()
        },
        "marker_url": _presign("put_object", bucket, marker_key),
        "log_url": _presign("put_object", bucket, log_key),
    }
    job_id = _submit(payload)
    log(f"parse submitted to the GPU worker as job {job_id}")
    marker = _wait_for_marker(bucket, marker_key, log_key, PARSE_DEADLINE_SECONDS, log, job_id, since)
    log(
        f"parse complete: {marker.get('pages')} pages, {marker.get('tables')} tables, {marker.get('total_seconds')}s"
    )
    for stem in PARSE_ARTIFACTS.values():
        s3.download_file(bucket, f"{prefix}{doc_key}.{stem}.json", str(workdir / f"{doc_key}.{stem}.json"))


def embed(doc_key: str, texts: List[str], log: Log = logger.info) -> List[List[float]]:
    """Embed the chunks on the worker. The texts go up as one file, the
    vectors come back as one file, both kept in the document's prefix."""
    bucket, prefix = _bucket_and_prefix(doc_key)
    chunks_key = f"{prefix}{doc_key}.chunks.json"
    vectors_key = f"{prefix}{doc_key}.vectors.json"
    marker_key = f"{prefix}{doc_key}.embed.complete.json"
    log_key = f"{prefix}{doc_key}.embed.log"
    s3 = _s3()
    since = time.time() - 60
    s3.put_object(Bucket=bucket, Key=chunks_key, Body=json.dumps({"texts": texts}).encode("utf-8"))
    payload = {
        "task": "embed",
        "base": doc_key,
        "chunks_url": _presign("get_object", bucket, chunks_key),
        "vectors_url": _presign("put_object", bucket, vectors_key),
        "marker_url": _presign("put_object", bucket, marker_key),
        "log_url": _presign("put_object", bucket, log_key),
    }
    job_id = _submit(payload)
    log(f"embedding of {len(texts)} chunks submitted to the GPU worker as job {job_id}")
    marker = _wait_for_marker(bucket, marker_key, log_key, EMBED_DEADLINE_SECONDS, log, job_id, since)
    vectors = json.loads(s3.get_object(Bucket=bucket, Key=vectors_key)["Body"].read())["vectors"]
    if len(vectors) != len(texts):
        raise RuntimeError(f"worker returned {len(vectors)} vectors for {len(texts)} chunks")
    log(f"embedding complete: {len(vectors)} vectors in {marker.get('total_seconds')}s")
    return vectors