"""
Multi-Version Concurrency Control (MVCC) Store.

MVCC allows multiple concurrent transactions to read and write
without blocking each other. Each transaction sees a consistent
snapshot of the database at a particular timestamp.

Key Features:
- Snapshot isolation for reads
- No read locks (readers never block writers)
- Garbage collection of old versions
- External consistency via commit timestamps
"""

import threading
from typing import Optional, Iterator, Tuple
from pathlib import Path
import logging

from spymonk_enterprise.time.hybrid_clock import Timestamp, HybridLogicalClock
from spymonk_enterprise.storage.engine.memtable import MemTable
from spymonk_enterprise.storage.engine.wal import WriteAheadLog

logger = logging.getLogger(__name__)


class MVCCStore:
    """
    MVCC storage engine with WAL and MemTable.

    This is the main storage interface that provides:
    - Durable writes (via WAL)
    - Fast reads (via MemTable)
    - MVCC snapshots
    - Crash recovery

    Future: Add SSTables for disk storage
    """

    def __init__(self, data_dir: Path, clock: HybridLogicalClock):
        """
        Initialize MVCC store.

        Args:
            data_dir: Directory for storage files
            clock: Hybrid logical clock for timestamps
        """
        self.data_dir = Path(data_dir)
        self.clock = clock

        # Create data directory
        self.data_dir.mkdir(parents=True, exist_ok=True)

        # Initialize components
        self.wal = WriteAheadLog(self.data_dir / 'wal')
        self.memtable = MemTable()

        # Thread safety
        self._lock = threading.Lock()

        # Recovery
        self._recover()

        logger.info(f"Initialized MVCC store in {data_dir}")

    def put(self, key: bytes, value: bytes, timestamp: Optional[Timestamp] = None) -> Timestamp:
        """
        Write key-value pair.

        Args:
            key: Key to write
            value: Value to write
            timestamp: Transaction timestamp (if None, use current time)

        Returns:
            Timestamp of the write
        """
        if timestamp is None:
            timestamp = self.clock.now()

        # Write to WAL first (durability)
        self.wal.append(timestamp, key, value, deleted=False)

        # Then write to MemTable
        should_flush = self.memtable.put(key, value, timestamp)

        # Check if we need to flush
        if should_flush:
            self._flush()

        return timestamp

    def get(self, key: bytes, timestamp: Optional[Timestamp] = None) -> Optional[bytes]:
        """
        Read value for key.

        Args:
            key: Key to read
            timestamp: Read timestamp (snapshot isolation)

        Returns:
            Value if found, None otherwise
        """
        if timestamp is None:
            timestamp = self.clock.now()

        # Read from MemTable
        value = self.memtable.get(key, timestamp)

        # TODO: If not found in MemTable, read from SSTables

        return value

    def delete(self, key: bytes, timestamp: Optional[Timestamp] = None) -> Timestamp:
        """
        Delete key (tombstone).

        Args:
            key: Key to delete
            timestamp: Transaction timestamp

        Returns:
            Timestamp of the delete
        """
        if timestamp is None:
            timestamp = self.clock.now()

        # Write tombstone to WAL
        self.wal.append(timestamp, key, None, deleted=True)

        # Mark as deleted in MemTable
        self.memtable.delete(key, timestamp)

        return timestamp

    def scan(self,
             start_key: Optional[bytes] = None,
             end_key: Optional[bytes] = None,
             timestamp: Optional[Timestamp] = None) -> Iterator[Tuple[bytes, bytes]]:
        """
        Range scan.

        Args:
            start_key: Start of range (inclusive)
            end_key: End of range (exclusive)
            timestamp: Read timestamp

        Yields:
            (key, value) tuples
        """
        if timestamp is None:
            timestamp = self.clock.now()

        # Scan MemTable
        yield from self.memtable.scan(start_key, end_key, timestamp)

        # TODO: Merge with SSTable scans

    def _flush(self):
        """Flush MemTable to disk (SSTable)"""
        logger.info("Flushing MemTable to disk")

        with self._lock:
            # TODO: Write MemTable to SSTable

            # Clear MemTable
            self.memtable.clear()

            # Checkpoint WAL
            self.wal.checkpoint()

    def _recover(self):
        """Recover from WAL on startup"""
        logger.info("Recovering from WAL")

        count = 0
        for entry in self.wal.replay():
            if entry.deleted:
                self.memtable.delete(entry.key, entry.timestamp)
            else:
                self.memtable.put(entry.key, entry.value, entry.timestamp)
            count += 1

        logger.info(f"Recovered {count} entries from WAL")

    def close(self):
        """Close the store"""
        self.wal.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
