"""MPack local object store — pack file read/write for the Muse VCS.

Layout
------
Pack files live under ``.muse/objects/pack/sha256/``, keeping the algorithm
canonical in the path, mirroring the loose object store convention::

    .muse/objects/sha256/<prefix>/<remainder>   ← loose objects
    .muse/objects/pack/sha256/<64hex>.mpack     ← MPack data file
    .muse/objects/pack/sha256/<64hex>.idx       ← MPack seek index

Pack file format (``.mpack``)
------------------------------
::

    [4 bytes]  magic:        b"MUSE"
    [1 byte]   version:      1
    [8 bytes]  object_count: uint64 little-endian
    --- object records (object_count entries) ---
    [71 bytes] object_id:    b"sha256:" + 64 lowercase hex chars
    [8 bytes]  length:       uint64 little-endian
    [N bytes]  content:      raw bytes
    --- footer ---
    [32 bytes] integrity:    SHA-256 of every byte above

Index file format (``.idx``)
-----------------------------
::

    [4 bytes]  magic:          b"MUSI"
    [1 byte]   version:        1
    [8 bytes]  entry_count:    uint64 little-endian
    --- entries, sorted by object_id (enables binary search) ---
    [71 bytes] object_id
    [8 bytes]  content_offset: uint64 little-endian
    [8 bytes]  content_length: uint64 little-endian
    --- footer ---
    [32 bytes] integrity:      SHA-256 of every byte above

``content_offset`` is the byte offset to the first content byte in the
``.mpack`` file (past the object_id and length fields for that entry).
A read is: ``seek(content_offset); read(content_length)`` — O(1), no decode.

Pack ID
-------
``pack_id = "sha256:" + sha256(<entire .mpack file bytes>).hexdigest()``

The pack is content-addressed like every other Muse object.  Writing the same
objects twice produces the same pack file and the same pack_id — idempotent.
"""

from __future__ import annotations

import bisect
import hashlib
import os
import pathlib
import struct
import tempfile
from typing import TYPE_CHECKING

from muse.core.paths import packs_dir
from muse.core.types import DEFAULT_HASH_ALGO, split_id
from muse.core.validation import validate_object_id

if TYPE_CHECKING:
    pass

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

_PACK_MAGIC = b"MUSE"
_IDX_MAGIC = b"MUSI"
_VERSION = 1
_OID_BYTES = 71        # len("sha256:") + 64 hex chars
_FOOTER_BYTES = 32     # SHA-256 digest
_PACK_HEADER_BYTES = 4 + 1 + 8          # magic + version + object_count
_IDX_HEADER_BYTES = 4 + 1 + 8           # magic + version + entry_count
_IDX_ENTRY_BYTES = _OID_BYTES + 8 + 8   # object_id + content_offset + content_length


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _sha256_file(path: pathlib.Path) -> bytes:
    h = hashlib.sha256()
    with path.open("rb") as fh:
        for chunk in iter(lambda: fh.read(65536), b""):
            h.update(chunk)
    return h.digest()


def _atomic_write(dest: pathlib.Path, data: bytes) -> None:
    """Write *data* to *dest* atomically via a temp file + rename."""
    dest.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp_str = tempfile.mkstemp(dir=dest.parent, prefix=".pack-tmp-")
    tmp = pathlib.Path(tmp_str)
    try:
        with os.fdopen(fd, "wb") as fh:
            fh.write(data)
            fh.flush()
            os.fsync(fh.fileno())
        os.replace(tmp, dest)
    except Exception:
        tmp.unlink(missing_ok=True)
        raise


def _pack_path(repo_root: pathlib.Path, hex_id: str) -> pathlib.Path:
    return packs_dir(repo_root) / f"{hex_id}.mpack"


def _idx_path(repo_root: pathlib.Path, hex_id: str) -> pathlib.Path:
    return packs_dir(repo_root) / f"{hex_id}.idx"


# ---------------------------------------------------------------------------
# Build
# ---------------------------------------------------------------------------


def _build_pack(objects: list[tuple[str, bytes]]) -> bytes:
    """Encode *objects* as an MPack binary blob. Returns raw bytes."""
    h = hashlib.sha256()

    def _emit(chunk: bytes) -> bytes:
        h.update(chunk)
        return chunk

    parts: list[bytes] = []
    parts.append(_emit(_PACK_MAGIC))
    parts.append(_emit(struct.pack("<BQ", _VERSION, len(objects))))

    for oid, content in objects:
        oid_bytes = oid.encode()
        assert len(oid_bytes) == _OID_BYTES, f"bad oid length: {oid!r}"
        parts.append(_emit(oid_bytes))
        parts.append(_emit(struct.pack("<Q", len(content))))
        parts.append(_emit(content))

    parts.append(h.digest())  # footer — NOT fed back into h
    return b"".join(parts)


def _build_idx(objects: list[tuple[str, bytes]], pack_bytes: bytes) -> bytes:
    """Build a sorted seek index for *objects* referencing *pack_bytes*."""
    # Compute content offsets within pack_bytes.
    # Pack layout: header(13) then for each object: oid(71) + length(8) + content(N)
    offsets: dict[str, tuple[int, int]] = {}
    cursor = _PACK_HEADER_BYTES
    for oid, content in objects:
        cursor += _OID_BYTES + 8          # skip object_id + length field
        offsets[oid] = (cursor, len(content))
        cursor += len(content)

    # Sort entries by object_id for binary search.
    sorted_oids = sorted(offsets)

    h = hashlib.sha256()

    def _emit(chunk: bytes) -> bytes:
        h.update(chunk)
        return chunk

    parts: list[bytes] = []
    parts.append(_emit(_IDX_MAGIC))
    parts.append(_emit(struct.pack("<BQ", _VERSION, len(sorted_oids))))

    for oid in sorted_oids:
        content_offset, content_length = offsets[oid]
        parts.append(_emit(oid.encode()))
        parts.append(_emit(struct.pack("<QQ", content_offset, content_length)))

    parts.append(h.digest())
    return b"".join(parts)


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def _parse_pack_bytes(pack_bytes: bytes) -> list[tuple[str, bytes]]:
    """Parse raw pack bytes (as produced by _build_pack) into (object_id, content) pairs.

    Verifies the SHA-256 footer before parsing. Returns [] for empty input.
    """
    if not pack_bytes:
        return []
    if pack_bytes[:4] != _PACK_MAGIC:
        raise ValueError(f"Bad pack magic: {pack_bytes[:4]!r}")
    body = pack_bytes[:-_FOOTER_BYTES]
    stored = pack_bytes[-_FOOTER_BYTES:]
    if hashlib.sha256(body).digest() != stored:
        raise OSError("Pack bytes failed SHA-256 integrity check")
    count = struct.unpack_from("<Q", pack_bytes, 5)[0]
    cursor = _PACK_HEADER_BYTES
    objects: list[tuple[str, bytes]] = []
    for _ in range(count):
        oid = pack_bytes[cursor: cursor + _OID_BYTES].decode()
        cursor += _OID_BYTES
        length = struct.unpack_from("<Q", pack_bytes, cursor)[0]
        cursor += 8
        content = pack_bytes[cursor: cursor + length]
        cursor += length
        objects.append((oid, content))
    return objects


def write_pack(repo_root: pathlib.Path, objects: list[tuple[str, bytes]]) -> str | None:
    """Write *objects* as an MPack file and its seek index.

    Returns the ``pack_id`` (``sha256:<64hex>``) on success, or ``None`` when
    *objects* is empty (no files written).

    Idempotent: if a pack with the same content already exists the existing
    files are left untouched and the same pack_id is returned.

    Args:
        repo_root: Root of the Muse repository.
        objects:   List of ``(object_id, content)`` pairs to pack.

    Returns:
        ``"sha256:<64hex>"`` pack_id, or ``None`` for an empty list.
    """
    if not objects:
        return None

    pack_bytes = _build_pack(objects)
    hex_id = hashlib.sha256(pack_bytes).hexdigest()
    pack_id = f"{DEFAULT_HASH_ALGO}:{hex_id}"

    mpack = _pack_path(repo_root, hex_id)
    idx = _idx_path(repo_root, hex_id)

    if mpack.exists() and idx.exists():
        return pack_id

    idx_bytes = _build_idx(objects, pack_bytes)
    _atomic_write(mpack, pack_bytes)
    _atomic_write(idx, idx_bytes)
    return pack_id


# ---------------------------------------------------------------------------
# Index lookup helpers
# ---------------------------------------------------------------------------


def _load_idx(idx_path: pathlib.Path) -> list[tuple[str, int, int]]:
    """Parse an index file into a sorted list of (object_id, offset, length)."""
    data = idx_path.read_bytes()

    # Verify footer integrity.
    body = data[:-_FOOTER_BYTES]
    stored_digest = data[-_FOOTER_BYTES:]
    actual_digest = hashlib.sha256(body).digest()
    if actual_digest != stored_digest:
        raise OSError(
            f"MPack index {idx_path.name} failed integrity check — store may be corrupt."
        )

    if not data.startswith(_IDX_MAGIC):
        raise OSError(f"MPack index {idx_path.name} has wrong magic bytes.")

    entry_count = struct.unpack_from("<Q", data, 5)[0]
    entries: list[tuple[str, int, int]] = []
    cursor = _IDX_HEADER_BYTES
    for _ in range(entry_count):
        oid = data[cursor: cursor + _OID_BYTES].decode()
        cursor += _OID_BYTES
        offset, length = struct.unpack_from("<QQ", data, cursor)
        cursor += 16
        entries.append((oid, offset, length))

    return entries  # already sorted (written sorted)


def _binary_search(entries: list[tuple[str, int, int]], oid: str) -> tuple[int, int] | None:
    """Binary-search sorted *entries* for *oid*. Returns (offset, length) or None."""
    keys = [e[0] for e in entries]
    i = bisect.bisect_left(keys, oid)
    if i < len(entries) and entries[i][0] == oid:
        return entries[i][1], entries[i][2]
    return None


def _all_idx_paths(repo_root: pathlib.Path) -> list[pathlib.Path]:
    """Return all .idx files in the pack store, or [] if none exist."""
    d = packs_dir(repo_root)
    if not d.exists():
        return []
    return sorted(d.glob("*.idx"))


# ---------------------------------------------------------------------------
# Read / has / list / verify
# ---------------------------------------------------------------------------


def has_object_in_packs(repo_root: pathlib.Path, object_id: str) -> bool:
    """Return ``True`` if *object_id* is present in any local pack."""
    for idx_path in _all_idx_paths(repo_root):
        try:
            entries = _load_idx(idx_path)
        except OSError:
            continue
        if _binary_search(entries, object_id) is not None:
            return True
    return False


def read_object_from_packs(repo_root: pathlib.Path, object_id: str) -> bytes | None:
    """Seek-read *object_id* from the first pack that contains it.

    Returns raw content bytes, or ``None`` if the object is absent from all
    local packs.

    Every read verifies the SHA-256 of the returned bytes against *object_id*
    before returning — silent corruption raises ``OSError``.
    """
    for idx_path in _all_idx_paths(repo_root):
        try:
            entries = _load_idx(idx_path)
        except OSError:
            continue
        result = _binary_search(entries, object_id)
        if result is None:
            continue
        content_offset, content_length = result
        hex_id = idx_path.stem
        mpack = _pack_path(repo_root, hex_id)
        with mpack.open("rb") as fh:
            fh.seek(content_offset)
            content = fh.read(content_length)
        # Integrity check — object_id is the muse-format hash (hash_blob).
        from muse.core.ids import hash_blob
        actual = hash_blob(content)
        if actual != object_id:
            raise OSError(
                f"Object {object_id} failed integrity check reading from pack "
                f"{mpack.name} — store may be corrupt."
            )
        return content
    return None


def list_packs(repo_root: pathlib.Path) -> list[str]:
    """Return the pack_id for every pack present in the local store."""
    result = []
    for idx_path in _all_idx_paths(repo_root):
        result.append(f"{DEFAULT_HASH_ALGO}:{idx_path.stem}")
    return result


def verify_pack(repo_root: pathlib.Path, pack_id: str) -> bool:
    """Re-hash *pack_id*'s ``.mpack`` and ``.idx`` footers.

    Returns ``True`` when both files pass. Raises ``OSError`` on any
    integrity failure so callers can distinguish corrupt from absent.

    Args:
        repo_root: Root of the Muse repository.
        pack_id:   ``sha256:<64hex>`` pack identifier.

    Raises:
        OSError:  If either file is missing, has a wrong magic, or fails its
                  SHA-256 footer check.
    """
    _, hex_id = split_id(pack_id)
    mpack = _pack_path(repo_root, hex_id)
    idx = _idx_path(repo_root, hex_id)

    for path, magic in [(mpack, _PACK_MAGIC), (idx, _IDX_MAGIC)]:
        if not path.exists():
            raise OSError(f"MPack file not found: {path}")
        data = path.read_bytes()
        if not data.startswith(magic):
            raise OSError(f"{path.name} has wrong magic bytes.")
        body = data[:-_FOOTER_BYTES]
        stored = data[-_FOOTER_BYTES:]
        actual = hashlib.sha256(body).digest()
        if actual != stored:
            raise OSError(
                f"{path.name} failed SHA-256 integrity check — store may be corrupt."
            )

    return True