#!/usr/bin/env python3
"""
cdr_state_manager.py - CDR-backed reversible state management for AI/QPU pipelines

Provides O(1) VM memory reversible checkpointing for:
- Token-level inference rollback
- Basin phase exploration
- QPU measurement replay
- DAG node re-execution

Uses CDR (Chained Deterministic Regeneration) bijective state transitions:
  Forward:  S' = P(S XOR M)
  Backward: S  = P^{-1}(S') XOR M

Integration points:
- FlowTorch decode loop (token checkpointing)
- Basin decode (phase state tracking)
- QPU nudge (measurement logging)
- BloomOS ledger (audit trail)

Usage:
    from helix_cdc.cdr_state_manager import CDRStateManager

    mgr = CDRStateManager(tape_size=65536)

    # Checkpoint inference state
    mgr.checkpoint("step_0", {"tokens": [1,2,3], "logits": logits_tensor})
    mgr.checkpoint("step_1", {"tokens": [1,2,3,4], "logits": logits_tensor})

    # Rollback to previous state
    state = mgr.rollback("step_0")

    # Or use CDR pure mode for O(1) memory
    mgr.absorb_state(state_bytes)  # Forward
    mgr.step_back()                 # Backward (requires LIFO order)
"""

from __future__ import annotations
import hashlib
import struct
import json
import time
import os
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Dict, List, Optional, Any, Tuple
from pathlib import Path
import sys

if TYPE_CHECKING:
    from helix_cdc.memory.cdr_echo_bridge import CDREchoBridge

# Persistence configuration (env toggles)
CDR_PERSIST = os.environ.get("CDR_PERSIST", "1") != "0"
CDR_PERSIST_EVERY = int(os.environ.get("CDR_PERSIST_EVERY", "10"))

# Add scripts to path for CDR primitives
sys.path.insert(0, str(Path(__file__).parent.parent / "scripts"))

try:
    from helix64_exec import (
        _cdr_expand_seed,
        _cdr_permute,
        _cdr_invert,
        _cdr_xor_rate,
        CDR_STATE_BYTES,
        CDR_RATE_BYTES,
    )
    CDR_AVAILABLE = True
except ImportError:
    CDR_AVAILABLE = False
    CDR_STATE_BYTES = 200
    CDR_RATE_BYTES = 64


@dataclass
class Checkpoint:
    """A single checkpoint in the state history."""
    checkpoint_id: str
    timestamp: float
    state_hash: str
    state_data: bytes
    step_index: int
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict:
        return {
            "checkpoint_id": self.checkpoint_id,
            "timestamp": self.timestamp,
            "state_hash": self.state_hash,
            "step_index": self.step_index,
            "metadata": self.metadata,
            "state_size": len(self.state_data),
        }


@dataclass
class CDRState:
    """CDR sponge state for pure O(1) mode."""
    state: bytearray
    step_count: int = 0
    initialized: bool = False


class CDRStateManager:
    """
    Reversible state manager using CDR for AI/QPU pipelines.

    Two modes:
    1. Checkpoint mode: Store full states, random access
    2. Pure CDR mode: O(1) memory, LIFO rollback only
    """

    def __init__(
        self,
        tape_size: int = 65536,
        max_checkpoints: int = 1024,
        enable_delta_compression: bool = True,
        ledger_path: Optional[Path] = None,
        echo_bridge: Optional["CDREchoBridge"] = None,
        component: str = "cdr",
    ):
        """
        Initialize CDR state manager.

        Args:
            tape_size: Size of MEMB tape for pure mode (default 64 KiB)
            max_checkpoints: Maximum checkpoints to store
            enable_delta_compression: Use XOR deltas between checkpoints
            ledger_path: Optional path for audit ledger
            echo_bridge: Optional CDREchoBridge for crash-resilient persistence
            component: Component name for persistence keys ("basin", "qpu", "cdna", etc.)
        """
        self.tape_size = tape_size
        self.max_checkpoints = max_checkpoints
        self.enable_delta_compression = enable_delta_compression
        self.ledger_path = ledger_path
        self.echo_bridge = echo_bridge
        self.component = component
        self._persist_counter = 0

        # Checkpoint storage (random access mode)
        self.checkpoints: Dict[str, Checkpoint] = {}
        self.checkpoint_order: List[str] = []

        # CDR pure mode state
        self.cdr_state = CDRState(
            state=bytearray(CDR_STATE_BYTES),
            step_count=0,
            initialized=False,
        )

        # MEMB tape for pure mode
        self.memb_tape = bytearray(tape_size)
        self.memb_ptr = 0

        # Message history for LIFO rollback (fast mode)
        self.message_history: List[bytes] = []

        # Delta compression cache
        self.prev_state_bytes: Optional[bytes] = None

        # Statistics
        self.stats = {
            "checkpoints_created": 0,
            "rollbacks_performed": 0,
            "cdr_absorbs": 0,
            "cdr_step_backs": 0,
            "bytes_stored": 0,
            "bytes_saved_delta": 0,
        }

    # ========================================================================
    # CHECKPOINT MODE (Random Access)
    # ========================================================================

    def checkpoint(
        self,
        checkpoint_id: str,
        state: Dict[str, Any],
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Checkpoint:
        """
        Create a checkpoint with random access.

        Args:
            checkpoint_id: Unique identifier for this checkpoint
            state: State dictionary to checkpoint
            metadata: Optional metadata (step number, etc.)

        Returns:
            Checkpoint object
        """
        # Serialize state
        state_bytes = self._serialize_state(state)
        state_hash = hashlib.sha256(state_bytes).hexdigest()

        # Delta compression
        if self.enable_delta_compression and self.prev_state_bytes is not None:
            delta_bytes = self._compute_delta(self.prev_state_bytes, state_bytes)
            if len(delta_bytes) < len(state_bytes) * 0.8:  # Only use if >20% savings
                self.stats["bytes_saved_delta"] += len(state_bytes) - len(delta_bytes)
                # Store delta with marker
                state_bytes = b"DELTA:" + delta_bytes

        self.prev_state_bytes = state_bytes if not state_bytes.startswith(b"DELTA:") else self.prev_state_bytes

        # Create checkpoint
        cp = Checkpoint(
            checkpoint_id=checkpoint_id,
            timestamp=time.time(),
            state_hash=state_hash,
            state_data=state_bytes,
            step_index=len(self.checkpoint_order),
            metadata=metadata or {},
        )

        # Store
        self.checkpoints[checkpoint_id] = cp
        self.checkpoint_order.append(checkpoint_id)

        # Enforce max checkpoints (LRU eviction)
        while len(self.checkpoints) > self.max_checkpoints:
            oldest_id = self.checkpoint_order.pop(0)
            del self.checkpoints[oldest_id]

        # Update stats
        self.stats["checkpoints_created"] += 1
        self.stats["bytes_stored"] += len(state_bytes)

        # Append to ledger if configured
        if self.ledger_path:
            self._append_ledger(cp)

        return cp

    def load_checkpoint(self, checkpoint_id: str) -> Optional[Dict[str, Any]]:
        """
        Load a checkpoint by ID.

        Args:
            checkpoint_id: Checkpoint to load

        Returns:
            State dictionary or None if not found
        """
        cp = self.checkpoints.get(checkpoint_id)
        if cp is None:
            return None

        state_bytes = cp.state_data

        # Handle delta-compressed checkpoints
        if state_bytes.startswith(b"DELTA:"):
            # Need to reconstruct from base + delta
            # For now, return None (delta reconstruction requires base state)
            return None

        return self._deserialize_state(state_bytes)

    def rollback(self, checkpoint_id: str) -> Optional[Dict[str, Any]]:
        """
        Rollback to a checkpoint and discard later checkpoints.

        Args:
            checkpoint_id: Checkpoint to rollback to

        Returns:
            State dictionary at checkpoint
        """
        if checkpoint_id not in self.checkpoints:
            return None

        cp = self.checkpoints[checkpoint_id]

        # Remove all checkpoints after this one
        idx = self.checkpoint_order.index(checkpoint_id)
        for later_id in self.checkpoint_order[idx + 1:]:
            # Only delete if it still exists (may have been overwritten)
            if later_id in self.checkpoints:
                del self.checkpoints[later_id]
        self.checkpoint_order = self.checkpoint_order[:idx + 1]

        self.stats["rollbacks_performed"] += 1

        return self.load_checkpoint(checkpoint_id)

    # ========================================================================
    # CDR PURE MODE (O(1) Memory, LIFO Rollback)
    # ========================================================================

    def cdr_init(self, seed: bytes = b"\x00" * 64) -> None:
        """
        Initialize CDR sponge state from seed.

        Args:
            seed: 64-byte seed for deterministic initialization
        """
        if not CDR_AVAILABLE:
            raise RuntimeError("CDR primitives not available (helix64_exec not found)")

        self.cdr_state.state = _cdr_expand_seed(seed[:64].ljust(64, b"\x00"))
        self.cdr_state.step_count = 0
        self.cdr_state.initialized = True
        self.message_history = []
        self.memb_ptr = 0

    def cdr_absorb(self, message: bytes, store_message: bool = True) -> bytes:
        """
        CDR forward step: S' = P(S XOR M)

        Args:
            message: Message bytes to absorb (max 64 bytes)
            store_message: If True, store for fast-mode rollback

        Returns:
            State hash after absorption
        """
        if not self.cdr_state.initialized:
            self.cdr_init()

        msg = message[:CDR_RATE_BYTES].ljust(CDR_RATE_BYTES, b"\x00")

        if store_message:
            # Fast mode: store message for automatic rollback
            self.message_history.append(msg)
        else:
            # Pure mode: write to MEMB tape
            if self.memb_ptr + len(msg) <= self.tape_size:
                self.memb_tape[self.memb_ptr:self.memb_ptr + len(msg)] = msg
                self.memb_ptr += len(msg)

        # CDR forward: S' = P(S XOR M)
        _cdr_xor_rate(self.cdr_state.state, msg)
        _cdr_permute(self.cdr_state.state)

        self.cdr_state.step_count += 1
        self.stats["cdr_absorbs"] += 1

        # Auto-persist to crash-resilient storage
        self._maybe_auto_persist()

        return hashlib.sha256(bytes(self.cdr_state.state)).hexdigest()

    def cdr_step_back(self, message: Optional[bytes] = None) -> bytes:
        """
        CDR backward step: S = P^{-1}(S') XOR M

        Args:
            message: Message to XOR back (None = pop from history)

        Returns:
            State hash after step back
        """
        if not self.cdr_state.initialized:
            raise RuntimeError("CDR not initialized")

        if message is None:
            # Fast mode: pop from message history (LIFO)
            if not self.message_history:
                raise RuntimeError("No messages to step back (history empty)")
            msg = self.message_history.pop()
        else:
            # Pure mode: caller provides message
            msg = message[:CDR_RATE_BYTES].ljust(CDR_RATE_BYTES, b"\x00")

        # CDR backward: S = P^{-1}(S') XOR M
        _cdr_invert(self.cdr_state.state)
        _cdr_xor_rate(self.cdr_state.state, msg)

        self.cdr_state.step_count -= 1
        self.stats["cdr_step_backs"] += 1

        return hashlib.sha256(bytes(self.cdr_state.state)).hexdigest()

    def cdr_squeeze(self, length: int = 64) -> bytes:
        """
        Extract bytes from CDR rate portion.

        Args:
            length: Number of bytes to extract (max 64)

        Returns:
            Extracted bytes
        """
        if not self.cdr_state.initialized:
            raise RuntimeError("CDR not initialized")

        return bytes(self.cdr_state.state[:min(length, CDR_RATE_BYTES)])

    def cdr_state_hash(self) -> str:
        """Get current CDR state hash."""
        if not self.cdr_state.initialized:
            return "uninitialized"
        return hashlib.sha256(bytes(self.cdr_state.state)).hexdigest()

    # ========================================================================
    # PERSISTENCE (Crash-Resilient Storage)
    # ========================================================================

    def persist_summary(self, key: Optional[str] = None) -> bool:
        """
        Persist CDR state summary to crash-resilient storage.

        Saves ONLY summary metadata (hash/step_count/checkpoint IDs).
        Does NOT persist prompts, outputs, or raw absorbed messages.

        Args:
            key: Optional custom key (default: "cdr:{component}:last")

        Returns:
            True on success, False if persistence disabled or unavailable
        """
        if not CDR_PERSIST or self.echo_bridge is None:
            return False

        try:
            from helix_cdc.memory.cdr_echo_bridge import build_record
        except ImportError:
            return False

        key = key or f"cdr:{self.component}:last"
        record = build_record(
            component=self.component,
            state_hash=self.cdr_state_hash(),
            step_count=int(self.cdr_state.step_count) if self.cdr_state.initialized else 0,
            checkpoints=list(self.checkpoints.keys())[-10:],
        )

        return self.echo_bridge.save(key, record)

    def load_summary(self, key: Optional[str] = None) -> Optional[Dict[str, Any]]:
        """
        Load CDR state summary from persistent storage.

        Use after crash/restart to see where we left off.

        Args:
            key: Optional custom key (default: "cdr:{component}:last")

        Returns:
            Dict with ts, component, state_hash, step_count, checkpoints, extra
            or None if not found
        """
        if self.echo_bridge is None:
            return None

        key = key or f"cdr:{self.component}:last"
        return self.echo_bridge.load(key)

    def _maybe_auto_persist(self) -> None:
        """Auto-persist if counter hits threshold."""
        if not CDR_PERSIST or self.echo_bridge is None:
            return

        self._persist_counter += 1
        if self._persist_counter % CDR_PERSIST_EVERY == 0:
            self.persist_summary()

    # ========================================================================
    # INFERENCE INTEGRATION HELPERS
    # ========================================================================

    def checkpoint_inference_step(
        self,
        step: int,
        tokens: List[int],
        logits: Optional[Any] = None,
        hidden_state: Optional[Any] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> str:
        """
        Checkpoint a single inference step.

        Args:
            step: Step number
            tokens: Generated tokens so far
            logits: Optional logits tensor (will be converted to list)
            hidden_state: Optional hidden state
            metadata: Additional metadata

        Returns:
            Checkpoint ID
        """
        state = {
            "step": step,
            "tokens": tokens,
            "token_count": len(tokens),
        }

        if logits is not None:
            # Handle torch/numpy tensors
            if hasattr(logits, "tolist"):
                state["logits_hash"] = hashlib.sha256(
                    str(logits.tolist()[:100]).encode()  # Hash first 100 for efficiency
                ).hexdigest()[:16]
            elif hasattr(logits, "numpy"):
                state["logits_hash"] = hashlib.sha256(
                    logits.numpy().tobytes()[:1000]
                ).hexdigest()[:16]

        if hidden_state is not None:
            if hasattr(hidden_state, "shape"):
                state["hidden_shape"] = list(hidden_state.shape)

        checkpoint_id = f"infer_step_{step}"
        meta = metadata or {}
        meta["inference_step"] = step

        self.checkpoint(checkpoint_id, state, meta)
        return checkpoint_id

    def checkpoint_basin_phase(
        self,
        phi0_deg: float,
        output_tokens: List[int],
        coherence: Optional[float] = None,
    ) -> str:
        """
        Checkpoint a basin phase exploration.

        Args:
            phi0_deg: Phase angle (0-360)
            output_tokens: Generated output for this phase
            coherence: Optional coherence measurement

        Returns:
            Checkpoint ID
        """
        state = {
            "phi0_deg": phi0_deg,
            "output_tokens": output_tokens,
            "output_length": len(output_tokens),
        }

        if coherence is not None:
            state["coherence"] = coherence

        checkpoint_id = f"basin_phi_{int(phi0_deg)}"
        self.checkpoint(checkpoint_id, state, {"basin_phase": phi0_deg})
        return checkpoint_id

    def checkpoint_qpu_measurement(
        self,
        measurement: int,
        p1: float,
        job_id: str,
        circuit_hash: Optional[str] = None,
    ) -> str:
        """
        Checkpoint a QPU measurement result.

        Args:
            measurement: Measurement result (0 or 1)
            p1: Probability of |1⟩
            job_id: IBM Quantum job ID
            circuit_hash: Optional circuit hash

        Returns:
            Checkpoint ID
        """
        state = {
            "measurement": measurement,
            "p1": p1,
            "job_id": job_id,
        }

        if circuit_hash:
            state["circuit_hash"] = circuit_hash

        checkpoint_id = f"qpu_{job_id[:8]}"
        self.checkpoint(checkpoint_id, state, {"qpu_measurement": True})
        return checkpoint_id

    # ========================================================================
    # INTERNAL HELPERS
    # ========================================================================

    def _serialize_state(self, state: Dict[str, Any]) -> bytes:
        """Serialize state dict to bytes."""
        return json.dumps(state, default=str, sort_keys=True).encode("utf-8")

    def _deserialize_state(self, state_bytes: bytes) -> Dict[str, Any]:
        """Deserialize bytes to state dict."""
        return json.loads(state_bytes.decode("utf-8"))

    def _compute_delta(self, base: bytes, target: bytes) -> bytes:
        """Compute XOR delta between two states."""
        # Pad to same length
        max_len = max(len(base), len(target))
        base = base.ljust(max_len, b"\x00")
        target = target.ljust(max_len, b"\x00")

        # XOR delta
        delta = bytes(a ^ b for a, b in zip(base, target))

        # Simple RLE compression of zeros
        compressed = bytearray()
        i = 0
        while i < len(delta):
            if delta[i] == 0:
                # Count consecutive zeros
                count = 0
                while i + count < len(delta) and delta[i + count] == 0 and count < 255:
                    count += 1
                compressed.extend([0, count])
                i += count
            else:
                compressed.append(delta[i])
                i += 1

        return bytes(compressed)

    def _append_ledger(self, cp: Checkpoint) -> None:
        """Append checkpoint to audit ledger."""
        if self.ledger_path is None:
            return

        entry = {
            "type": "cdr_checkpoint",
            "checkpoint": cp.to_dict(),
            "cdr_state_hash": self.cdr_state_hash(),
            "timestamp": time.time(),
        }

        with open(self.ledger_path, "a") as f:
            f.write(json.dumps(entry) + "\n")

    def get_stats(self) -> Dict[str, Any]:
        """Get manager statistics."""
        return {
            **self.stats,
            "active_checkpoints": len(self.checkpoints),
            "cdr_step_count": self.cdr_state.step_count,
            "cdr_initialized": self.cdr_state.initialized,
            "message_history_size": len(self.message_history),
            "memb_tape_used": self.memb_ptr,
            "memb_tape_capacity": self.tape_size,
        }


# ============================================================================
# CONVENIENCE FUNCTIONS
# ============================================================================

_default_manager: Optional[CDRStateManager] = None


def get_default_manager() -> CDRStateManager:
    """Get or create default CDR state manager."""
    global _default_manager
    if _default_manager is None:
        _default_manager = CDRStateManager()
    return _default_manager


def checkpoint(checkpoint_id: str, state: Dict[str, Any], **kwargs) -> Checkpoint:
    """Convenience: checkpoint using default manager."""
    return get_default_manager().checkpoint(checkpoint_id, state, **kwargs)


def rollback(checkpoint_id: str) -> Optional[Dict[str, Any]]:
    """Convenience: rollback using default manager."""
    return get_default_manager().rollback(checkpoint_id)


def cdr_absorb(message: bytes) -> str:
    """Convenience: CDR absorb using default manager."""
    return get_default_manager().cdr_absorb(message)


def cdr_step_back() -> str:
    """Convenience: CDR step back using default manager."""
    return get_default_manager().cdr_step_back()


# ============================================================================
# DEMO / TEST
# ============================================================================

if __name__ == "__main__":
    print("CDR State Manager Demo")
    print("=" * 60)

    mgr = CDRStateManager(tape_size=4096)

    # Demo 1: Checkpoint mode
    print("\n1. CHECKPOINT MODE (Random Access)")
    print("-" * 40)

    for i in range(5):
        tokens = list(range(i + 1))
        cp = mgr.checkpoint_inference_step(
            step=i,
            tokens=tokens,
            metadata={"demo": True}
        )
        print(f"   Step {i}: checkpoint '{cp}', tokens={tokens}")

    print(f"\n   Rolling back to step 2...")
    state = mgr.rollback("infer_step_2")
    print(f"   Restored state: {state}")
    print(f"   Active checkpoints: {list(mgr.checkpoints.keys())}")

    # Demo 2: CDR pure mode
    print("\n2. CDR PURE MODE (O(1) Memory)")
    print("-" * 40)

    if CDR_AVAILABLE:
        mgr.cdr_init(b"demo_seed_12345678901234567890123456789012345678901234567890")
        initial_hash = mgr.cdr_state_hash()
        print(f"   Initial state: {initial_hash[:32]}...")

        messages = [f"message_{i}".encode() for i in range(3)]
        for msg in messages:
            h = mgr.cdr_absorb(msg)
            print(f"   Absorbed '{msg.decode()}': {h[:32]}...")

        print(f"\n   Stepping back (LIFO)...")
        for _ in range(3):
            h = mgr.cdr_step_back()
            print(f"   After step_back: {h[:32]}...")

        final_hash = mgr.cdr_state_hash()
        print(f"\n   Final state:   {final_hash[:32]}...")
        print(f"   Match initial: {initial_hash == final_hash}")
    else:
        print("   CDR primitives not available (run from helix-cdc root)")

    # Stats
    print("\n3. STATISTICS")
    print("-" * 40)
    for k, v in mgr.get_stats().items():
        print(f"   {k}: {v}")
