"""Where a document's files live.

Every file the playground keeps for a document is named by the document's
storage key and sits in the local data directory: the served PDF, the parse
and figure artifacts the ingest read, and the upload workdir. On the
developer host that directory is the only copy.

Hosted, the local directory is scratch that a task replacement erases, and
the durable copy is an ``s3://`` prefix from settings holding one
``<doc_key>/`` prefix per document with the same file names inside it. Two
movements keep the two in step: ``publish`` copies files up after every
stage of an upload, and ``local_pdf`` fetches the served PDF back down when
a task no longer has it. Nothing here ever deletes from the prefix; removing
a document from the library removes its database rows and local files only,
so a later upload of the same bytes finds every artifact still there.

Run as a module, it publishes the whole current library once, which is how
the host's documents were moved into the bucket:

    uv run python -m quber.playground.storage
"""

from __future__ import annotations

import json
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterable, Optional

from cloudpathlib import S3Path
from loguru import logger

from quber.settings import get_settings

DATA_DIR: Path = get_settings().playground.data_dir
UPLOADS: Path = DATA_DIR / "uploads"


def prefix(doc_key: str) -> Optional[S3Path]:
    """The document's prefix in the bucket, or None on the developer host."""
    root = get_settings().playground.artifacts_uri
    if not root:
        return None
    return S3Path(root.rstrip("/") + "/" + doc_key + "/")


def publish(doc_key: str, paths: Iterable[Path]) -> int:
    """Copy the given local files under the document's prefix. Returns how
    many were copied; zero, without touching anything, on the developer host
    or when none of the paths exist."""
    dest = prefix(doc_key)
    if dest is None:
        return 0
    copied = 0
    for path in paths:
        if not path.is_file():
            continue
        (dest / path.name).upload_from(path, force_overwrite_to_cloud=True)
        copied += 1
    if copied:
        logger.info("published {} file(s) for {} to {}", copied, doc_key, dest)
    return copied


def publish_workdir(doc_key: str, workdir: Path) -> int:
    """Copy every file the upload workdir holds for the document."""
    if not workdir.is_dir():
        return 0
    return publish(doc_key, sorted(p for p in workdir.iterdir() if p.name.startswith(doc_key)))


def write_source_manifest(
    doc_key: str, *, filename: str, content_hash: str, source_uri: Optional[str] = None
) -> Path:
    """Record where the document came from beside its PDF: the name it was
    uploaded under, its full hash, the S3 object it was taken from if any, and
    when. The key carries none of this, so the prefix says it instead."""
    manifest = {
        "doc_key": doc_key,
        "filename": filename,
        "sha256": content_hash,
        "source_uri": source_uri,
        "uploaded_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
    }
    path = UPLOADS / f"{doc_key}.source.json"
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(manifest, indent=2))
    return path


def local_pdf(doc_key: str) -> Optional[Path]:
    """The served PDF as a local path, fetched from the bucket when this host
    no longer has it. None when it exists nowhere."""
    for candidate in (DATA_DIR / f"{doc_key}.pdf", UPLOADS / f"{doc_key}.pdf"):
        if candidate.is_file():
            return candidate
    remote = prefix(doc_key)
    if remote is None:
        return None
    source = remote / f"{doc_key}.pdf"
    if not source.exists():
        return None
    target = DATA_DIR / f"{doc_key}.pdf"
    target.parent.mkdir(parents=True, exist_ok=True)
    logger.info("fetching {} from {}", target.name, source)
    source.download_to(target)
    return target


def publish_library() -> None:
    """Publish every document the database knows about from the local data
    directory: the served PDF and the artifacts beside it, plus the upload
    workdir when this host still has it."""
    from quber.playground import db

    if prefix("probe") is None:
        sys.exit("QUBER_PLAYGROUND_ARTIFACTS_URI is not set; nothing to publish to")
    with db.connect() as conn:
        rows = conn.execute(
            "SELECT doc_key, filename, content_hash FROM ade_playground.documents ORDER BY id"
        ).fetchall()
    total = 0
    for doc_key, filename, content_hash in rows:
        files = sorted(DATA_DIR.glob(f"{doc_key}.*"))
        if not any(f.name == f"{doc_key}.source.json" for f in files):
            files.append(write_source_manifest(doc_key, filename=filename, content_hash=content_hash))
        copied = publish(doc_key, files) + publish_workdir(doc_key, UPLOADS / f"{doc_key}-artifacts")
        total += copied
        logger.info("{}: {} file(s)", doc_key, copied)
    logger.success("published {} file(s) for {} document(s)", total, len(rows))


if __name__ == "__main__":
    publish_library()
