"""Materialize documents/collections from stored manifests and chunks."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
import shutil
import uuid

from .store import ObjectStore
from .hash import sha256_hex
from ..manifest.codec import load_manifest_bytes, load_collection_bytes
from ..manifest.schema import DocManifest
from ..index.local import LocalIndex
from .attest import AttestationError, attest_manifest
from ..chunking.registry import get_method
from ..chunking.api import ManifestChunk


class MaterializeError(Exception):
    pass


@dataclass(frozen=True)
class MaterializeResult:
    doc_id: str
    output_path: Path


@dataclass(frozen=True)
class MaterializeCollectionResult:
    collection_id: str
    output_dir: Path


def _reconstruct_bytes(manifest: DocManifest, store: ObjectStore) -> bytes:
    """Reconstruct document bytes from a manifest and local chunks."""
    method = get_method(manifest.chunking.get("method"))
    params = manifest.chunking.get("params", {})
    chunks: list[ManifestChunk] = []
    for entry in manifest.chunks:
        chunk_hash = entry.get("hash")
        chunk = store.get(chunk_hash)
        if chunk is None:
            raise MaterializeError("missing chunk in store")
        if sha256_hex(chunk) != chunk_hash:
            raise MaterializeError("chunk hash mismatch")
        chunks.append(ManifestChunk(data=chunk, entry=entry))
    return method.reconstruct(chunks, params)


def materialize_doc_id(
    doc_id: str,
    store: ObjectStore,
    index: LocalIndex,
    output_path: Path,
) -> MaterializeResult:
    """Materialize a document identified by doc_id to output_path."""
    to_stdout = output_path == Path("-")
    data = index.load()
    documents = data.get("documents", {})
    doc = documents.get(doc_id)
    if doc is None:
        raise MaterializeError("document not found in local index")
    manifest_hashes = doc.get("manifests", [])
    if not manifest_hashes:
        raise MaterializeError("no manifests for doc_id")

    manifest_hash = manifest_hashes[0]
    manifest_bytes = store.get(manifest_hash)
    if manifest_bytes is None:
        raise MaterializeError("manifest not found in local store")

    manifest = load_manifest_bytes(manifest_bytes)
    # Verify manifest coherence before materializing.
    try:
        attest_manifest(manifest, store, expected_doc_id=doc_id, expected_manifest_hash=manifest_hash)
    except AttestationError as e:
        raise MaterializeError(str(e)) from e

    data_bytes = _reconstruct_bytes(manifest, store)
    if to_stdout:
        import sys

        sys.stdout.buffer.write(data_bytes)
        return MaterializeResult(doc_id=doc_id, output_path=output_path)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_bytes(data_bytes)
    return MaterializeResult(doc_id=doc_id, output_path=output_path)


def materialize_collection(
    collection_id: str,
    store: ObjectStore,
    index: LocalIndex,
    output_dir: Path,
) -> MaterializeCollectionResult:
    """Materialize all documents in a collection into output_dir."""
    collection_bytes = store.get(collection_id)
    if collection_bytes is None:
        raise MaterializeError("collection not found in local store")
    collection = load_collection_bytes(collection_bytes)

    if output_dir.exists():
        if any(output_dir.iterdir()):
            raise MaterializeError("output directory is not empty")
        output_dir.rmdir()

    tmp_dir = output_dir.with_name(f"{output_dir.name}.tmp-{uuid.uuid4().hex}")
    tmp_dir.mkdir(parents=True, exist_ok=False)

    try:
        for entry in collection.documents:
            if entry.path is None:
                raise MaterializeError("collection entry missing path")
            out_path = tmp_dir / entry.path
            materialize_doc_id(entry.doc_id, store, index, out_path)
        tmp_dir.replace(output_dir)
    except Exception:
        shutil.rmtree(tmp_dir, ignore_errors=True)
        raise

    return MaterializeCollectionResult(collection_id=collection_id, output_dir=output_dir)
