"""Attestation logic for documents and manifests."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

from .hash import sha256_hex, canonicalize_bytes
from .store import ObjectStore
from ..manifest.codec import (
    load_manifest_bytes,
    dump_manifest_bytes,
    load_collection_bytes,
    dump_collection_bytes,
    load_manifest_object,
    canonical_json_bytes,
)
from ..manifest.schema import DocManifest
from ..index.local import LocalIndex
from ..chunking.registry import get_method
from ..chunking.api import ManifestChunk


class AttestationError(Exception):
    pass


@dataclass(frozen=True)
class AttestationResult:
    doc_id: str
    manifest_hash: Optional[str]
    ok: bool


@dataclass(frozen=True)
class CollectionAttestationResult:
    collection_id: str
    ok: bool


def attest_file_to_doc_id(file_bytes: bytes, expected_doc_id: str) -> AttestationResult:
    """Attest that file bytes (canonicalized) match the given doc_id."""
    canon = canonicalize_bytes(file_bytes)
    actual = sha256_hex(canon)
    if actual != expected_doc_id:
        raise AttestationError("file hash does not match doc_id")
    return AttestationResult(doc_id=expected_doc_id, manifest_hash=None, ok=True)


def attest_manifest_bytes(
    manifest_bytes: bytes,
    store: ObjectStore,
    expected_doc_id: Optional[str] = None,
    expected_manifest_hash: Optional[str] = None,
) -> AttestationResult:
    """Attest a doc-manifest given its serialized bytes."""
    manifest = load_manifest_bytes(manifest_bytes)
    if expected_manifest_hash is not None:
        canonical_bytes = dump_manifest_bytes(manifest)
        actual_manifest_hash = sha256_hex(canonical_bytes)
        if actual_manifest_hash != expected_manifest_hash:
            raise AttestationError("manifest hash does not match expected")
    return attest_manifest(manifest, store, expected_doc_id, expected_manifest_hash)


def attest_manifest_bytes_shallow(
    manifest_bytes: bytes,
    expected_manifest_hash: Optional[str] = None,
) -> str:
    """Attest manifest integrity only (hash of canonical JSON object).

    This performs no schema validation and no chunk reconstruction.
    """
    obj = load_manifest_object(manifest_bytes)
    canonical_bytes = canonical_json_bytes(obj)
    actual_manifest_hash = sha256_hex(canonical_bytes)
    if expected_manifest_hash is not None and actual_manifest_hash != expected_manifest_hash:
        raise AttestationError("manifest hash does not match expected")
    return actual_manifest_hash


def attest_manifest(
    manifest: DocManifest,
    store: ObjectStore,
    expected_doc_id: Optional[str] = None,
    expected_manifest_hash: Optional[str] = None,
) -> AttestationResult:
    """Attest a doc-manifest object against local chunks."""
    try:
        method = get_method(manifest.chunking.get("method"))
    except ValueError as e:
        raise AttestationError(str(e)) from e
    if expected_doc_id is not None and manifest.doc_id != expected_doc_id:
        raise AttestationError("manifest doc_id does not match expected")

    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 AttestationError("missing chunk in store")
        if sha256_hex(chunk) != chunk_hash:
            raise AttestationError("chunk hash mismatch")
        chunks.append(ManifestChunk(data=chunk, entry=entry))

    reconstructed_bytes = method.reconstruct(chunks, params)
    reconstructed_hash = sha256_hex(reconstructed_bytes)
    if reconstructed_hash != manifest.doc_id:
        raise AttestationError("reconstructed doc hash does not match manifest doc_id")

    return AttestationResult(
        doc_id=manifest.doc_id,
        manifest_hash=expected_manifest_hash,
        ok=True,
    )


def attest_any_manifest_bytes(
    manifest_bytes: bytes,
    store: ObjectStore,
    index: LocalIndex,
    expected_manifest_hash: Optional[str] = None,
    expected_doc_id: Optional[str] = None,
) -> AttestationResult | CollectionAttestationResult:
    """Attest a manifest by type (doc-manifest or collection-manifest)."""
    obj = load_manifest_object(manifest_bytes)
    manifest_type = obj.get("type")
    if manifest_type == "doc-manifest":
        return attest_manifest_bytes(
            manifest_bytes,
            store,
            expected_doc_id=expected_doc_id,
            expected_manifest_hash=expected_manifest_hash,
        )
    if manifest_type == "collection-manifest":
        if expected_doc_id is not None:
            raise AttestationError("collection manifests do not accept --doc-id")
        return attest_collection_bytes(
            manifest_bytes,
            store,
            index,
            expected_collection_id=expected_manifest_hash,
        )
    raise AttestationError("unsupported manifest type for full attestation")


def attest_doc_id(doc_id: str, store: ObjectStore, index: LocalIndex) -> AttestationResult:
    """Attest a document by doc_id using the local index."""
    data = index.load()
    documents = data.get("documents", {})
    manifests = data.get("manifests", {})

    doc = documents.get(doc_id)
    if doc is None:
        raise AttestationError("document not found in local index")

    manifest_hashes = doc.get("manifests", [])
    if not manifest_hashes:
        raise AttestationError("no manifests for doc_id")

    for manifest_hash in manifest_hashes:
        manifest_bytes = store.get(manifest_hash)
        if manifest_bytes is None:
            continue
        obj = load_manifest_object(manifest_bytes)
        if obj.get("type") != "doc-manifest":
            continue
        return attest_manifest_bytes(
            manifest_bytes,
            store,
            expected_doc_id=doc_id,
            expected_manifest_hash=manifest_hash,
        )

    raise AttestationError("no doc-manifest found for doc_id")


def attest_collection_bytes(
    collection_bytes: bytes,
    store: ObjectStore,
    index: LocalIndex,
    expected_collection_id: Optional[str] = None,
) -> CollectionAttestationResult:
    """Attest a collection by verifying all referenced doc_ids."""
    collection = load_collection_bytes(collection_bytes)
    canonical = dump_collection_bytes(collection)
    collection_id = sha256_hex(canonical)
    if expected_collection_id is not None and collection_id != expected_collection_id:
        raise AttestationError("collection hash does not match expected")

    for entry in collection.documents:
        attest_doc_id(entry.doc_id, store, index)

    return CollectionAttestationResult(collection_id=collection_id, ok=True)
