"""Signature Registry - central thread-safe store of all known signatures.

The registry is the canonical source of truth for signature data.  It maps
unique string identifiers (sig_id) to NaturalSignature objects and provides
O(1) lookup.  It also maintains a secondary path-based index for file-based
signatures.

Persistence is provided via an append-only log that is replayed on startup:
    Each record is:
        [1 byte op] [2 bytes id_len] [id_len bytes sig_id] [64 bytes signature | 0]
    Op codes:
        0x01 = INSERT  (id + 64-byte signature)
        0x02 = DELETE  (id only)

A ``PerfectHashTable`` can be built from a frozen registry for guaranteed
O(1) lookup on static datasets.
"""

from __future__ import annotations

import os
import struct
import threading
from pathlib import Path
from typing import Callable, Dict, Iterator, List, Optional, Tuple, Union

from rbsp.core.encoding import SIGNATURE_SIZE, NaturalSignature

# -- Log operation codes ---------------------------------------------------
_OP_INSERT = 0x01
_OP_DELETE = 0x02


def _pack_insert(sig_id: str, sig: NaturalSignature) -> bytes:
    """Pack an INSERT record for the append-only log."""
    id_bytes = sig_id.encode("utf-8")
    sig_bytes = sig.to_bytes()
    return struct.pack(">BH", _OP_INSERT, len(id_bytes)) + id_bytes + sig_bytes


def _pack_delete(sig_id: str) -> bytes:
    """Pack a DELETE record for the append-only log."""
    id_bytes = sig_id.encode("utf-8")
    return struct.pack(">BH", _OP_DELETE, len(id_bytes)) + id_bytes


class SignatureRegistry:
    """Thread-safe in-memory signature store with optional file persistence.

    Provides O(1) insertion, deletion, and lookup by ``sig_id``.
    Supports secondary lookup by source path and iteration.

    Args:
        backing_path: Optional filesystem path for an append-only log.
                      If provided, the log is replayed on construction and
                      all mutations are durably appended.

    Usage::

        registry = SignatureRegistry()
        registry.add("abc123", sig, source_path="/foo/bar.py")
        retrieved = registry.get("abc123")
    """

    def __init__(self, backing_path: Optional[Union[str, Path]] = None) -> None:
        self._lock = threading.Lock()
        self._store: Dict[str, NaturalSignature] = {}
        self._path_index: Dict[str, str] = {}
        self._metadata: Dict[str, Dict[str, str]] = {}
        self._listeners: List[Callable[[str, str, Optional[NaturalSignature]], None]] = []

        self._log_path: Optional[Path] = None
        self._log_file = None

        if backing_path is not None:
            self._log_path = Path(backing_path)
            self._replay_log()
            self._log_file = open(self._log_path, "ab")

    # ------------------------------------------------------------------
    # Event listeners
    # ------------------------------------------------------------------

    def add_listener(
        self,
        callback: Callable[[str, str, Optional[NaturalSignature]], None],
    ) -> None:
        """Register a callback for registry mutations.

        The callback receives ``(event_type, sig_id, signature)`` where
        event_type is ``"add"`` or ``"remove"``.
        """
        self._listeners.append(callback)

    def _notify(self, event: str, sig_id: str, sig: Optional[NaturalSignature]) -> None:
        for cb in self._listeners:
            try:
                cb(event, sig_id, sig)
            except Exception:
                pass  # listeners must not crash the registry

    # ------------------------------------------------------------------
    # Mutation
    # ------------------------------------------------------------------

    def add(
        self,
        sig_id: str,
        sig: NaturalSignature,
        *,
        source_path: Optional[str] = None,
        metadata: Optional[Dict[str, str]] = None,
    ) -> None:
        """Insert or replace a signature.

        Args:
            sig_id:      Unique identifier for this signature.
            sig:         The NaturalSignature object.
            source_path: Optional originating file path for secondary index.
            metadata:    Optional key-value metadata.
        """
        with self._lock:
            self._store[sig_id] = sig
            if source_path:
                self._path_index[source_path] = sig_id
            if metadata:
                self._metadata[sig_id] = metadata
            if self._log_file is not None:
                self._log_file.write(_pack_insert(sig_id, sig))
        self._notify("add", sig_id, sig)

    def remove(self, sig_id: str) -> Optional[NaturalSignature]:
        """Remove and return a signature by *sig_id*, or ``None``."""
        with self._lock:
            sig = self._store.pop(sig_id, None)
            self._metadata.pop(sig_id, None)
            # Clean path index
            to_remove = [p for p, s in self._path_index.items() if s == sig_id]
            for p in to_remove:
                del self._path_index[p]
            if sig is not None and self._log_file is not None:
                self._log_file.write(_pack_delete(sig_id))
                self._log_file.flush()
        if sig is not None:
            self._notify("remove", sig_id, sig)
        return sig

    def remove_by_path(self, path: str) -> Optional[NaturalSignature]:
        """Remove and return the signature associated with *path*."""
        with self._lock:
            sid = self._path_index.pop(path, None)
            if sid is None:
                return None
            sig = self._store.pop(sid, None)
            self._metadata.pop(sid, None)
            if sig is not None and self._log_file is not None:
                self._log_file.write(_pack_delete(sid))
                self._log_file.flush()
        if sig is not None:
            self._notify("remove", sid, sig)
        return sig

    def bulk_add(self, entries: List[Tuple[str, NaturalSignature]]) -> int:
        """Insert multiple signatures atomically.

        Args:
            entries: List of (sig_id, signature) tuples.

        Returns:
            Number of signatures inserted.
        """
        with self._lock:
            buf = bytearray()
            for sig_id, sig in entries:
                self._store[sig_id] = sig
                buf.extend(_pack_insert(sig_id, sig))
            if self._log_file is not None and buf:
                self._log_file.write(buf)
                self._log_file.flush()
        return len(entries)

    # ------------------------------------------------------------------
    # Lookup
    # ------------------------------------------------------------------

    def get(self, sig_id: str) -> Optional[NaturalSignature]:
        """O(1) lookup by signature ID."""
        with self._lock:
            return self._store.get(sig_id)

    def get_by_path(self, path: str) -> Optional[NaturalSignature]:
        """O(1) lookup by source file path."""
        with self._lock:
            sid = self._path_index.get(path)
            if sid is None:
                return None
            return self._store.get(sid)

    def get_sig_id_by_path(self, path: str) -> Optional[str]:
        """O(1) lookup of sig_id by source file path."""
        with self._lock:
            return self._path_index.get(path)

    def get_metadata(self, sig_id: str) -> Optional[Dict[str, str]]:
        """Retrieve metadata for a signature."""
        with self._lock:
            return self._metadata.get(sig_id)

    # ------------------------------------------------------------------
    # Iteration / info
    # ------------------------------------------------------------------

    def __len__(self) -> int:
        with self._lock:
            return len(self._store)

    def __contains__(self, sig_id: str) -> bool:
        with self._lock:
            return sig_id in self._store

    def __iter__(self) -> Iterator[NaturalSignature]:
        with self._lock:
            return iter(list(self._store.values()))

    def all_ids(self) -> List[str]:
        """Return a snapshot list of all signature IDs."""
        with self._lock:
            return list(self._store.keys())

    def all_paths(self) -> List[str]:
        """Return a snapshot list of all indexed source paths."""
        with self._lock:
            return list(self._path_index.keys())

    def items(self) -> List[Tuple[str, NaturalSignature]]:
        """Return a snapshot list of (sig_id, signature) pairs."""
        with self._lock:
            return list(self._store.items())

    def clear(self) -> None:
        """Remove all entries."""
        with self._lock:
            self._store.clear()
            self._path_index.clear()
            self._metadata.clear()

    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    def _replay_log(self) -> None:
        """Replay the append-only log to reconstruct in-memory state."""
        if self._log_path is None or not self._log_path.exists():
            return

        data = self._log_path.read_bytes()
        pos = 0
        while pos < len(data):
            if pos + 3 > len(data):
                break  # truncated header
            op, id_len = struct.unpack_from(">BH", data, pos)
            pos += 3

            if pos + id_len > len(data):
                break  # truncated id
            sig_id = data[pos : pos + id_len].decode("utf-8", errors="replace")
            pos += id_len

            if op == _OP_INSERT:
                if pos + SIGNATURE_SIZE > len(data):
                    break  # truncated signature
                sig_bytes = data[pos : pos + SIGNATURE_SIZE]
                pos += SIGNATURE_SIZE
                try:
                    sig = NaturalSignature.from_bytes(sig_bytes)
                    self._store[sig_id] = sig
                except ValueError:
                    pass  # skip corrupted record
            elif op == _OP_DELETE:
                self._store.pop(sig_id, None)
            else:
                break  # unknown op, stop replay

    def compact(self) -> None:
        """Rewrite the backing log to contain only current state.

        Removes delete tombstones and superseded inserts, reducing disk usage.
        """
        if self._log_path is None:
            return

        with self._lock:
            if self._log_file is not None:
                self._log_file.close()

            with open(self._log_path, "wb") as f:
                for sig_id, sig in self._store.items():
                    f.write(_pack_insert(sig_id, sig))

            self._log_file = open(self._log_path, "ab")

    def flush(self) -> None:
        """Ensure all pending writes are flushed to disk."""
        if self._log_file is not None:
            self._log_file.flush()
            os.fsync(self._log_file.fileno())

    def close(self) -> None:
        """Close the backing log file."""
        if self._log_file is not None:
            self._log_file.flush()
            self._log_file.close()
            self._log_file = None

    def __del__(self) -> None:
        try:
            self.close()
        except Exception:
            pass

    # ------------------------------------------------------------------
    # Perfect hash support (static datasets)
    # ------------------------------------------------------------------

    def build_perfect_hash(self) -> "PerfectHashTable":
        """Build a minimal perfect hash table from the current registry state.

        Useful for frozen datasets where no further inserts are expected.
        Provides guaranteed O(1) lookup with no collision handling overhead.
        """
        with self._lock:
            return PerfectHashTable(list(self._store.items()))


class PerfectHashTable:
    """Minimal perfect hash table for static RBSP signature collections.

    Uses a two-level hashing scheme (CHD-style) to achieve O(1) guaranteed
    lookup with zero collisions for the frozen keyset.

    This table is immutable after construction.
    """

    def __init__(self, entries: List[Tuple[str, NaturalSignature]]) -> None:
        n = len(entries)
        if n == 0:
            self._table: List[Optional[Tuple[str, NaturalSignature]]] = []
            self._displacements: List[Tuple[int, int]] = []
            self._size = 0
            self._bucket_count = 0
            return

        self._size = n
        table_size = max(n, 4)
        self._table = [None] * table_size
        self._bucket_count = max(1, n // 4)
        self._displacements = [(0, 0)] * self._bucket_count

        # Bucket entries by primary hash
        buckets: List[List[Tuple[str, NaturalSignature]]] = [
            [] for _ in range(self._bucket_count)
        ]
        for key, sig in entries:
            bucket_idx = hash(key) % self._bucket_count
            buckets[bucket_idx].append((key, sig))

        # Process largest buckets first for greedy placement
        bucket_order = sorted(
            range(self._bucket_count),
            key=lambda i: len(buckets[i]),
            reverse=True,
        )

        for bi in bucket_order:
            bucket = buckets[bi]
            if not bucket:
                continue

            placed = False
            for d1 in range(table_size):
                for d2 in range(table_size):
                    slots = [
                        ((hash(key) * (d1 + 1) + d2) % table_size)
                        for key, _ in bucket
                    ]
                    if len(set(slots)) != len(slots):
                        continue
                    if all(self._table[s] is None for s in slots):
                        for slot, entry in zip(slots, bucket):
                            self._table[slot] = entry
                        self._displacements[bi] = (d1, d2)
                        placed = True
                        break
                if placed:
                    break

            if not placed:
                # Fallback: linear probe (exceedingly rare)
                for key, sig in bucket:
                    for s in range(table_size):
                        if self._table[s] is None:
                            self._table[s] = (key, sig)
                            break
                self._displacements[bi] = (0, 0)

    def get(self, sig_id: str) -> Optional[NaturalSignature]:
        """Look up a signature by *sig_id*.  Guaranteed O(1).

        Returns ``None`` if the key was not in the original keyset.
        """
        if self._size == 0:
            return None
        key_hash = hash(sig_id)
        bucket_idx = key_hash % self._bucket_count
        d1, d2 = self._displacements[bucket_idx]
        table_size = len(self._table)
        slot = ((key_hash * (d1 + 1) + d2) % table_size)
        entry = self._table[slot]
        if entry is not None and entry[0] == sig_id:
            return entry[1]

        # Fallback scan (only if construction used fallback)
        for entry in self._table:
            if entry is not None and entry[0] == sig_id:
                return entry[1]
        return None

    def __contains__(self, sig_id: str) -> bool:
        return self.get(sig_id) is not None

    def __len__(self) -> int:
        return self._size
