"""
LMDB-backed index and state storage for Verisoo.

Provides fast lookups on top of the append-only binary log:
- record_id -> log offset
- attestor_fingerprint -> [record indices]
- image_sha256 -> [record indices]
- image_phash -> [record indices]  (for perceptual search)
- merkle tree state

LMDB is memory-mapped, so reads are zero-copy and very fast.
Single-writer model fits append-only perfectly.
"""

from __future__ import annotations

import json
import struct
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Iterator

import lmdb

from .exceptions import VerisooError
from .models import AttestationRecord, Identity, ImageHashes


class IndexError(VerisooError):
    """Error in index operations."""

    pass


# LMDB database names (sub-databases within the environment)
DB_RECORDS = b"records"  # index (uint64) -> record_id
DB_BY_ATTESTOR = b"by_attestor"  # fingerprint -> [indices] (JSON array)
DB_BY_SHA256 = b"by_sha256"  # sha256 -> [indices]
DB_BY_PHASH = b"by_phash"  # phash -> [indices]
DB_IDENTITIES = b"identities"  # fingerprint -> identity JSON
DB_META = b"meta"  # key -> value (merkle root, log size, etc.)


@dataclass
class IndexStats:
    """Statistics about the index."""

    record_count: int
    attestor_count: int
    unique_images_sha256: int
    unique_images_phash: int
    merkle_root: str | None
    log_file_size: int


class LMDBIndex:
    """
    LMDB-backed index for attestation records.

    Works alongside BinaryLog - the log stores raw records,
    this index provides fast lookups by various keys.
    """

    # 1GB default max size, grows as needed up to this
    DEFAULT_MAP_SIZE = 1 * 1024 * 1024 * 1024

    def __init__(
        self,
        path: Path,
        *,
        map_size: int = DEFAULT_MAP_SIZE,
        readonly: bool = False,
    ) -> None:
        self.path = path
        path.mkdir(parents=True, exist_ok=True)

        self._env = lmdb.open(
            str(path),
            map_size=map_size,
            max_dbs=10,
            readonly=readonly,
            # Allow multiple readers
            max_readers=126,
            # Don't sync on every write (we have CRC in binlog for durability)
            sync=False,
            # Memory-map is the whole point
            writemap=True,
        )

        # Open sub-databases
        with self._env.begin(write=True) as txn:
            self._db_records = self._env.open_db(DB_RECORDS, txn=txn)
            self._db_by_attestor = self._env.open_db(DB_BY_ATTESTOR, txn=txn)
            self._db_by_sha256 = self._env.open_db(DB_BY_SHA256, txn=txn)
            self._db_by_phash = self._env.open_db(DB_BY_PHASH, txn=txn)
            self._db_identities = self._env.open_db(DB_IDENTITIES, txn=txn)
            self._db_meta = self._env.open_db(DB_META, txn=txn)

    def close(self) -> None:
        """Close the LMDB environment."""
        self._env.close()

    def __enter__(self) -> LMDBIndex:
        return self

    def __exit__(self, *args) -> None:
        self.close()

    # --- Record indexing ---

    def index_record(
        self,
        index: int,
        record: AttestationRecord,
        record_id: str,
    ) -> None:
        """
        Add a record to the index.

        Called after appending to the binary log.
        """
        with self._env.begin(write=True) as txn:
            # Store index -> record_id mapping
            txn.put(
                struct.pack("<Q", index),
                record_id.encode(),
                db=self._db_records,
            )

            # Index by attestor
            self._append_to_list(
                txn,
                self._db_by_attestor,
                record.attestor_fingerprint.encode(),
                index,
            )

            # Index by image SHA-256
            self._append_to_list(
                txn,
                self._db_by_sha256,
                record.image_hashes.sha256.encode(),
                index,
            )

            # Index by perceptual hash
            self._append_to_list(
                txn,
                self._db_by_phash,
                record.image_hashes.phash.encode(),
                index,
            )

    def get_record_id(self, index: int) -> str | None:
        """Get record ID for index."""
        with self._env.begin() as txn:
            data = txn.get(struct.pack("<Q", index), db=self._db_records)
            return data.decode() if data else None

    def get_indices_by_attestor(self, fingerprint: str) -> list[int]:
        """Get all record indices for an attestor."""
        with self._env.begin() as txn:
            return self._get_list(txn, self._db_by_attestor, fingerprint.encode())

    def get_indices_by_sha256(self, sha256: str) -> list[int]:
        """Get all record indices for an image by SHA-256."""
        with self._env.begin() as txn:
            return self._get_list(txn, self._db_by_sha256, sha256.encode())

    def get_indices_by_phash(self, phash: str) -> list[int]:
        """Get all record indices for an image by perceptual hash."""
        with self._env.begin() as txn:
            return self._get_list(txn, self._db_by_phash, phash.encode())

    def find_similar_phash(self, phash: str, max_distance: int = 10) -> list[tuple[str, int]]:
        """
        Find perceptual hashes within hamming distance.

        Returns list of (phash, distance) tuples.
        This is O(n) over unique phashes - for large scale, would need LSH.
        """
        results = []
        target = int(phash, 16)

        with self._env.begin() as txn:
            cursor = txn.cursor(db=self._db_by_phash)
            for key, _ in cursor:
                candidate = key.decode()
                try:
                    candidate_int = int(candidate, 16)
                    distance = bin(target ^ candidate_int).count("1")
                    if distance <= max_distance:
                        results.append((candidate, distance))
                except ValueError:
                    continue

        return sorted(results, key=lambda x: x[1])

    # --- Identity storage ---

    def save_identity(self, identity: Identity) -> None:
        """Save an identity (public key + metadata)."""
        data = {
            "public_key": identity.public_key.hex(),
            "fingerprint": identity.fingerprint,
            "created_at": identity.created_at.isoformat(),
            "metadata": identity.metadata,
        }
        with self._env.begin(write=True) as txn:
            txn.put(
                identity.fingerprint.encode(),
                json.dumps(data).encode(),
                db=self._db_identities,
            )

    def get_identity(self, fingerprint: str) -> Identity | None:
        """Get identity by fingerprint."""
        with self._env.begin() as txn:
            data = txn.get(fingerprint.encode(), db=self._db_identities)
            if not data:
                return None

            obj = json.loads(data.decode())
            return Identity(
                public_key=bytes.fromhex(obj["public_key"]),
                fingerprint=obj["fingerprint"],
                created_at=datetime.fromisoformat(obj["created_at"]),
                metadata=obj.get("metadata", {}),
            )

    def list_identities(self) -> list[str]:
        """List all identity fingerprints."""
        fingerprints = []
        with self._env.begin() as txn:
            cursor = txn.cursor(db=self._db_identities)
            for key, _ in cursor:
                fingerprints.append(key.decode())
        return fingerprints

    # --- Metadata storage ---

    def set_meta(self, key: str, value: str) -> None:
        """Set a metadata value."""
        with self._env.begin(write=True) as txn:
            txn.put(key.encode(), value.encode(), db=self._db_meta)

    def get_meta(self, key: str) -> str | None:
        """Get a metadata value."""
        with self._env.begin() as txn:
            data = txn.get(key.encode(), db=self._db_meta)
            return data.decode() if data else None

    def set_merkle_root(self, root: str) -> None:
        """Store current merkle root."""
        self.set_meta("merkle_root", root)

    def get_merkle_root(self) -> str | None:
        """Get stored merkle root."""
        return self.get_meta("merkle_root")

    def set_record_count(self, count: int) -> None:
        """Store record count (for consistency checks)."""
        self.set_meta("record_count", str(count))

    def get_record_count(self) -> int:
        """Get stored record count."""
        val = self.get_meta("record_count")
        return int(val) if val else 0

    # --- Stats ---

    def get_stats(self) -> IndexStats:
        """Get index statistics."""
        with self._env.begin() as txn:
            record_count = txn.stat(self._db_records)["entries"]
            attestor_count = txn.stat(self._db_by_attestor)["entries"]
            sha256_count = txn.stat(self._db_by_sha256)["entries"]
            phash_count = txn.stat(self._db_by_phash)["entries"]

        return IndexStats(
            record_count=record_count,
            attestor_count=attestor_count,
            unique_images_sha256=sha256_count,
            unique_images_phash=phash_count,
            merkle_root=self.get_merkle_root(),
            log_file_size=0,  # Filled by storage layer
        )

    # --- Internal helpers ---

    def _append_to_list(self, txn, db, key: bytes, value: int) -> None:
        """Append an integer to a JSON list stored at key."""
        existing = txn.get(key, db=db)
        if existing:
            items = json.loads(existing.decode())
        else:
            items = []

        items.append(value)
        txn.put(key, json.dumps(items).encode(), db=db)

    def _get_list(self, txn, db, key: bytes) -> list[int]:
        """Get JSON list of integers from key."""
        data = txn.get(key, db=db)
        if not data:
            return []
        return json.loads(data.decode())

    def sync(self) -> None:
        """Force sync to disk."""
        self._env.sync()

    def compact(self) -> None:
        """Compact the database (copy to remove free pages)."""
        compact_path = self.path.parent / f"{self.path.name}.compact"
        self._env.copy(str(compact_path), compact=True)
        # Could swap files here, but leave that to caller
