"""Write extraction artifacts to a local directory or an ``s3://`` prefix.

`output_sink` is a context manager that hands back a local directory to write
into. When the destination is an ``s3://`` URI, everything written there is
uploaded to the bucket on exit; local destinations are written in place and
never copied.

The write-local-then-upload shape exists because the artifact writers (pymupdf
for the annotated PDF, the HTML renderer, ``json.dump``) all need a real local
filesystem path -- none can stream to S3 directly. Uploads go through
cloudpathlib's default client, which resolves AWS credentials from the boto3
chain (env / ``~/.aws`` / IAM role), so no keys are passed explicitly.
"""

from __future__ import annotations

import json
import tempfile
from contextlib import contextmanager
from pathlib import Path
from typing import Any, Generator

from cloudpathlib import AnyPath, S3Path
from loguru import logger


def is_s3(output_dir: str) -> bool:
    """True when ``output_dir`` is an ``s3://`` URI rather than a local path."""
    return isinstance(AnyPath(output_dir), S3Path)


@contextmanager
def output_sink(output_dir: str) -> Generator[Path, None, None]:
    """Yield a local directory to write artifacts into.

    For a local ``output_dir`` the directory is created and yielded directly,
    so writes land in their final location. For an ``s3://`` ``output_dir`` a
    temporary directory is yielded; on clean exit every file written under it
    is uploaded to the bucket under the same relative path, then the temporary
    directory is removed. An exception inside the ``with`` block skips the
    upload and propagates, so a failed run never leaves partial artifacts in
    the bucket.
    """
    dest = AnyPath(output_dir)
    if isinstance(dest, S3Path):
        with tempfile.TemporaryDirectory(prefix="quber-out-") as tmp:
            work = Path(tmp)
            yield work
            uploaded = 0
            for f in sorted(work.rglob("*")):
                if not f.is_file():
                    continue
                target = dest / f.relative_to(work).as_posix()
                target.upload_from(f, force_overwrite_to_cloud=True)
                logger.info("Uploaded {} -> {}", f.name, target)
                uploaded += 1
            logger.info("Uploaded {} artifact(s) to {}", uploaded, dest)
    else:
        local = Path(output_dir)
        local.mkdir(parents=True, exist_ok=True)
        yield local


def write_completion_marker(output_dir: str, filename: str, payload: dict[str, Any]) -> None:
    """Write a job's completion marker as its own, final upload.

    The marker's existence tells downstream (the join step that launches
    fusion) that this job's artifacts are fully in place, so it must be the
    last object to land. It goes through its own `output_sink` after the
    artifact sink has closed: a single sink uploads files in sorted order,
    and a marker name can sort before the artifacts it vouches for.
    """
    with output_sink(output_dir) as out:
        (out / filename).write_text(json.dumps(payload, indent=2, default=str), encoding="utf-8")