"""
Transaction Manager for SpyMonk-DB.

Implements Spanner-style transactions:
- Read-write transactions with external consistency
- Read-only transactions with snapshot isolation (lock-free!)
- Two-phase commit for distributed transactions
"""

import uuid
import threading
from enum import Enum
from typing import Optional, Dict, List, Set
from dataclasses import dataclass, field
import logging

from spymonk_enterprise.time.hybrid_clock import Timestamp, HybridLogicalClock
from spymonk_enterprise.storage.mvcc import MVCCStore

logger = logging.getLogger(__name__)


class TransactionType(Enum):
    """Type of transaction"""
    READ_WRITE = "read_write"
    READ_ONLY = "read_only"


class TransactionState(Enum):
    """Transaction state"""
    ACTIVE = "active"
    PREPARING = "preparing"
    COMMITTED = "committed"
    ABORTED = "aborted"


@dataclass
class Mutation:
    """Single write operation in a transaction"""
    key: bytes
    value: Optional[bytes]
    is_delete: bool = False


class Transaction:
    """
    Single transaction.

    Spanner-style transaction with:
    - Snapshot isolation for reads
    - External consistency for writes
    - Lock-free read-only transactions
    """

    def __init__(
        self,
        txn_id: str,
        txn_type: TransactionType,
        start_timestamp: Timestamp,
        manager: 'TransactionManager'
    ):
        self.txn_id = txn_id
        self.type = txn_type
        self.start_timestamp = start_timestamp
        self.manager = manager
        self.state = TransactionState.ACTIVE

        # Track operations
        self.read_timestamps: List[Timestamp] = []
        self.mutations: List[Mutation] = []
        self.commit_timestamp: Optional[Timestamp] = None

    def get(self, key: bytes) -> Optional[bytes]:
        """
        Read key.

        Read-only txns: Use start_timestamp (snapshot isolation)
        Read-write txns: Use current timestamp (may see own writes)
        """
        if self.state != TransactionState.ACTIVE:
            raise Exception(f"Transaction {self.txn_id} is not active")

        if self.type == TransactionType.READ_ONLY:
            # Snapshot read at start_timestamp
            value = self.manager.store.get(key, self.start_timestamp)
            return value
        else:
            # Check if we've written to this key in this transaction
            for mutation in reversed(self.mutations):
                if mutation.key == key:
                    return None if mutation.is_delete else mutation.value

            # Read from storage
            value = self.manager.store.get(key, self.manager.clock.now())
            return value

    def put(self, key: bytes, value: bytes):
        """
        Write key-value pair.

        Buffered until commit.
        """
        if self.state != TransactionState.ACTIVE:
            raise Exception(f"Transaction {self.txn_id} is not active")

        if self.type == TransactionType.READ_ONLY:
            raise Exception("Cannot write in read-only transaction")

        mutation = Mutation(key=key, value=value, is_delete=False)
        self.mutations.append(mutation)

    def delete(self, key: bytes):
        """
        Delete key.

        Buffered until commit.
        """
        if self.state != TransactionState.ACTIVE:
            raise Exception(f"Transaction {self.txn_id} is not active")

        if self.type == TransactionType.READ_ONLY:
            raise Exception("Cannot delete in read-only transaction")

        mutation = Mutation(key=key, value=None, is_delete=True)
        self.mutations.append(mutation)

    def commit(self) -> bool:
        """
        Commit transaction.

        Returns:
            True if committed successfully
        """
        return self.manager.commit(self)

    def abort(self):
        """Abort transaction"""
        self.state = TransactionState.ABORTED


class TransactionManager:
    """
    Transaction manager.

    Coordinates transactions and ensures external consistency.
    """

    def __init__(self, clock: HybridLogicalClock, store: MVCCStore):
        self.clock = clock
        self.store = store

        # Active transactions
        self.active_txns: Dict[str, Transaction] = {}
        self._lock = threading.Lock()

        logger.info("Initialized transaction manager")

    def begin(self, read_only: bool = False) -> Transaction:
        """
        Begin a new transaction.

        Args:
            read_only: If True, create read-only transaction (lock-free!)

        Returns:
            New transaction
        """
        txn_id = str(uuid.uuid4())
        txn_type = TransactionType.READ_ONLY if read_only else TransactionType.READ_WRITE

        # For read-only txns, use current time as snapshot timestamp
        # For read-write txns, we'll determine commit timestamp later
        start_timestamp = self.clock.now()

        txn = Transaction(
            txn_id=txn_id,
            txn_type=txn_type,
            start_timestamp=start_timestamp,
            manager=self
        )

        with self._lock:
            self.active_txns[txn_id] = txn

        logger.debug(f"Started {txn_type.value} transaction {txn_id}")
        return txn

    def commit(self, txn: Transaction) -> bool:
        """
        Commit transaction.

        For read-only: instant commit (no writes)
        For read-write: choose commit timestamp + commit wait

        Args:
            txn: Transaction to commit

        Returns:
            True if committed successfully
        """
        if txn.state != TransactionState.ACTIVE:
            logger.warning(f"Transaction {txn.txn_id} already finalized")
            return False

        if txn.type == TransactionType.READ_ONLY:
            # Read-only: instant commit (no writes)
            txn.state = TransactionState.COMMITTED
            with self._lock:
                del self.active_txns[txn.txn_id]
            logger.debug(f"Committed read-only transaction {txn.txn_id}")
            return True

        # Read-write transaction
        try:
            txn.state = TransactionState.PREPARING

            # Choose commit timestamp
            # Must be > all read timestamps and > current time
            commit_ts = self.clock.now()
            if txn.read_timestamps:
                max_read_ts = max(txn.read_timestamps)
                if commit_ts < max_read_ts:
                    # Advance commit timestamp
                    commit_ts = Timestamp(
                        physical=max_read_ts.physical,
                        logical=max_read_ts.logical + 1,
                        uncertainty_ns=commit_ts.uncertainty_ns
                    )

            txn.commit_timestamp = commit_ts

            # Apply mutations
            for mutation in txn.mutations:
                if mutation.is_delete:
                    self.store.delete(mutation.key, commit_ts)
                else:
                    self.store.put(mutation.key, mutation.value, commit_ts)

            # CRITICAL: Commit wait (ensures external consistency!)
            # This is Spanner's secret sauce for linearizability
            self.clock.wait_until(commit_ts)

            # Mark as committed
            txn.state = TransactionState.COMMITTED

            with self._lock:
                del self.active_txns[txn.txn_id]

            logger.debug(f"Committed transaction {txn.txn_id} at {commit_ts}")
            return True

        except Exception as e:
            logger.error(f"Failed to commit transaction {txn.txn_id}: {e}")
            txn.state = TransactionState.ABORTED
            with self._lock:
                if txn.txn_id in self.active_txns:
                    del self.active_txns[txn.txn_id]
            return False

    def abort(self, txn: Transaction):
        """Abort transaction"""
        txn.state = TransactionState.ABORTED
        with self._lock:
            if txn.txn_id in self.active_txns:
                del self.active_txns[txn.txn_id]
        logger.debug(f"Aborted transaction {txn.txn_id}")
