"""
Robust sandbox monitoring with atomic state transitions.

Designed to handle VM snapshot/resume cycles gracefully without race conditions.
"""

import json
import time
import threading
from pathlib import Path
from dataclasses import dataclass, asdict
from typing import Optional, Callable
from enum import Enum

from blocks_control_sdk.constants.core import is_micro_vm


class MonitorState(Enum):
    """
    Sandbox monitor states with clear transitions:

    STOPPED -> ACTIVE -> IDLE -> PAUSING -> (snapshot occurs)
                 ^         |
                 |_________|  (on heartbeat/activity)

    After resume: State is read from file, resume detected, transitions back to IDLE
    """
    STOPPED = "stopped"
    ACTIVE = "active"      # Agent is working
    IDLE = "idle"          # Agent completed, countdown running
    PAUSING = "pausing"    # Pause request sent, awaiting snapshot


@dataclass
class MonitorSnapshot:
    """Persistent state that survives across snapshot/resume"""
    state: str
    expiry_monotonic: Optional[float]  # monotonic time when idle expires
    expiry_wall: Optional[float]       # wall time for resume detection
    sequence_number: int               # for idempotency
    created_at: float


class SandboxMonitor:
    """
    Robust sandbox monitoring with VM snapshot-safe state transitions.

    Key design principles:
    1. State written to disk BEFORE any external calls
    2. Monotonic + wall clock for resume detection
    3. Single callback invocation guaranteed via sequence numbers
    4. No background threads - poll-based checking

    Singleton: Only one instance is created. Subsequent calls to SandboxMonitor()
    return the existing instance and skip __init__.
    """

    STATE_FILE = Path("/tmp/blocks_sandbox_monitor.json")
    RESUME_MARKER = Path("/tmp/blocks_pre_snapshot_marker")
    CLOCK_JUMP_THRESHOLD = 10.0  # seconds - threshold for resume detection

    _instance: Optional['SandboxMonitor'] = None

    def __new__(cls, *args, **kwargs):
        if cls._instance is not None:
            print("[SandboxMonitor] Returning existing instance")
            return cls._instance
        instance = super().__new__(cls)
        cls._instance = instance
        return instance

    def __init__(
        self,
        default_timeout_seconds: float = None,
        on_pause_requested: Optional[Callable[[], None]] = None,
        on_tick: Optional[Callable[[float], None]] = None,
    ):
        if hasattr(self, '_initialized') and self._initialized:
            return
        """
        Args:
            default_timeout_seconds: Default idle timeout
            on_pause_requested: Callback when pause is needed.
                                Called AT MOST ONCE per idle cycle (fire-and-forget).
            on_tick: Optional callback for progress updates (remaining seconds)
        """
        self._lock = threading.Lock()
        if default_timeout_seconds is not None:
            self._default_timeout = default_timeout_seconds
        else:
            _is_micro_vm = is_micro_vm()
            self._default_timeout = 2.5 * 60 if _is_micro_vm else 15 * 60
        self._on_pause_requested = on_pause_requested
        self._on_tick = on_tick

        # In-memory state
        self._state = MonitorState.STOPPED
        self._expiry_monotonic: Optional[float] = None
        self._expiry_wall: Optional[float] = None
        self._sequence_number = 0
        self._last_pause_sequence = -1  # Track which sequence we paused on

        # Resume detection
        self._last_monotonic = time.monotonic()
        self._last_wall = time.time()

        # Tick tracking
        self._last_tick_time = time.monotonic()

        self._initialized = True

    def configure(
        self,
        default_timeout_seconds: float = None,
        on_pause_requested: Optional[Callable[[], None]] = None,
        on_tick: Optional[Callable[[float], None]] = None,
    ):
        """Update callbacks and settings on an existing instance."""
        with self._lock:
            if default_timeout_seconds is not None:
                self._default_timeout = default_timeout_seconds
            if on_pause_requested is not None:
                self._on_pause_requested = on_pause_requested
            if on_tick is not None:
                self._on_tick = on_tick

    def _write_state(self) -> None:
        """Atomically write state to disk"""
        snapshot = MonitorSnapshot(
            state=self._state.value,
            expiry_monotonic=self._expiry_monotonic,
            expiry_wall=self._expiry_wall,
            sequence_number=self._sequence_number,
            created_at=time.time(),
        )
        temp_file = self.STATE_FILE.with_suffix('.tmp')
        temp_file.write_text(json.dumps(asdict(snapshot)))
        temp_file.rename(self.STATE_FILE)

    def _read_state(self) -> Optional[MonitorSnapshot]:
        """Read state from disk"""
        try:
            data = json.loads(self.STATE_FILE.read_text())
            return MonitorSnapshot(**data)
        except (FileNotFoundError, json.JSONDecodeError):
            return None

    def _detect_resume(self) -> bool:
        """
        Detect if we resumed from a VM snapshot.

        Uses two methods:
        1. Wall clock jumped significantly vs monotonic clock
        2. Pre-snapshot marker file exists
        """
        current_monotonic = time.monotonic()
        current_wall = time.time()

        monotonic_delta = current_monotonic - self._last_monotonic
        wall_delta = current_wall - self._last_wall

        self._last_monotonic = current_monotonic
        self._last_wall = current_wall

        # Method 1: Clock divergence
        if abs(wall_delta - monotonic_delta) > self.CLOCK_JUMP_THRESHOLD:
            print(f"[SandboxMonitor] Resume detected via clock jump: wall={wall_delta:.1f}s, mono={monotonic_delta:.1f}s")
            return True

        # Method 2: Marker file
        if self.RESUME_MARKER.exists():
            self.RESUME_MARKER.unlink()
            print("[SandboxMonitor] Resume detected via marker file")
            return True

        return False

    def _mark_pre_snapshot(self) -> None:
        """Write marker before triggering pause/snapshot"""
        self.RESUME_MARKER.write_text(str(time.time()))

    def start_active(self) -> None:
        """Signal that agent is actively working"""
        with self._lock:
            self._state = MonitorState.ACTIVE
            self._expiry_monotonic = None
            self._expiry_wall = None
            self._write_state()
            print("[SandboxMonitor] State -> ACTIVE")

    def start_idle(self, timeout_seconds: Optional[float] = None) -> None:
        """
        Signal that agent completed work and idle countdown should start.

        Args:
            timeout_seconds: Override default timeout. If None, uses default.
        """
        with self._lock:
            if self._state == MonitorState.PAUSING:
                # Don't restart idle if we're already pausing
                return

            timeout = timeout_seconds or self._default_timeout
            now_mono = time.monotonic()
            now_wall = time.time()

            self._state = MonitorState.IDLE
            self._expiry_monotonic = now_mono + timeout
            self._expiry_wall = now_wall + timeout
            self._sequence_number += 1
            self._write_state()
            print(f"[SandboxMonitor] State -> IDLE (timeout={timeout:.1f}s, seq={self._sequence_number})")

    def heartbeat(self, extend_seconds: float = 0) -> None:
        """
        Signal activity that should extend the idle timer.

        If in ACTIVE state, does nothing (already active).
        If in IDLE state, extends the timeout.
        If in PAUSING state, transitions back to IDLE (cancel pause).
        """
        with self._lock:
            if self._state == MonitorState.STOPPED:
                return

            if self._state == MonitorState.ACTIVE:
                # Already active, nothing to do
                return

            # Extend or restart idle timer
            now_mono = time.monotonic()
            now_wall = time.time()
            timeout = self._default_timeout + extend_seconds

            new_expiry_mono = now_mono + timeout
            new_expiry_wall = now_wall + timeout

            if self._state == MonitorState.PAUSING:
                # Cancel pause, go back to idle
                print("[SandboxMonitor] Heartbeat during PAUSING, cancelling pause")
                self._state = MonitorState.IDLE

            # Extend timeout (only if it would increase remaining time)
            if self._expiry_monotonic is None or new_expiry_mono > self._expiry_monotonic:
                self._expiry_monotonic = new_expiry_mono
                self._expiry_wall = new_expiry_wall
                self._sequence_number += 1

            self._write_state()

    def stop(self) -> None:
        """Stop monitoring completely"""
        with self._lock:
            self._state = MonitorState.STOPPED
            self._expiry_monotonic = None
            self._expiry_wall = None
            self._write_state()
            # Clean up marker if exists
            self.RESUME_MARKER.unlink(missing_ok=True)
            print("[SandboxMonitor] State -> STOPPED")

    def check(self) -> MonitorState:
        """
        Check current state and handle transitions.

        This is the main polling method. Should be called regularly from the main loop.

        Returns:
            Current state after any transitions
        """
        with self._lock:
            # First, check for resume from snapshot
            if self._detect_resume():
                self._handle_resume()

            # Handle tick callback
            if self._on_tick and self._state == MonitorState.IDLE:
                now = time.monotonic()
                if now - self._last_tick_time >= 60:
                    self._last_tick_time = now
                    remaining = self._get_remaining_seconds_unlocked()
                    # Call tick outside lock to prevent deadlocks
                    threading.Thread(
                        target=self._on_tick,
                        args=(remaining,),
                        daemon=True
                    ).start()

            # Check for idle expiry
            if self._state == MonitorState.IDLE and self._expiry_monotonic is not None:
                if time.monotonic() >= self._expiry_monotonic:
                    self._transition_to_pausing()

            return self._state

    def _handle_resume(self) -> None:
        """Handle resume from snapshot - reset to appropriate state"""
        print("[SandboxMonitor] Handling resume from snapshot")

        # Read persisted state
        persisted = self._read_state()
        if persisted:
            print(f"[SandboxMonitor] Persisted state: {persisted.state}, seq={persisted.sequence_number}")

        # DON'T override if we're already ACTIVE (agent already started working)
        if self._state == MonitorState.ACTIVE:
            print("[SandboxMonitor] Already ACTIVE, skipping resume transition to IDLE")
            self._last_pause_sequence = -1  # Allow new pause later
            return

        # Only transition to IDLE if we weren't already active
        # This prevents immediately re-triggering pause
        self._state = MonitorState.IDLE
        now_mono = time.monotonic()
        now_wall = time.time()
        self._expiry_monotonic = now_mono + self._default_timeout
        self._expiry_wall = now_wall + self._default_timeout
        self._sequence_number += 1
        self._last_pause_sequence = -1  # Allow new pause
        self._write_state()
        print(f"[SandboxMonitor] After resume: State -> IDLE (fresh timeout, seq={self._sequence_number})")

    def _transition_to_pausing(self) -> None:
        """
        Transition to PAUSING state and invoke callback.

        Key invariants:
        1. State written BEFORE callback
        2. Callback invoked AT MOST ONCE per sequence
        3. Marker written for resume detection
        """
        current_seq = self._sequence_number

        # Check if we already paused on this sequence
        if current_seq == self._last_pause_sequence:
            print(f"[SandboxMonitor] Already paused on sequence {current_seq}, skipping")
            return

        # Transition to PAUSING
        self._state = MonitorState.PAUSING
        self._last_pause_sequence = current_seq
        self._write_state()

        # Write marker BEFORE callback (for resume detection)
        self._mark_pre_snapshot()

        print(f"[SandboxMonitor] State -> PAUSING (seq={current_seq})")

        # Invoke callback (fire-and-forget, will likely trigger snapshot)
        if self._on_pause_requested:
            try:
                self._on_pause_requested()
            except Exception as e:
                # On failure, go back to idle with fresh timeout
                print(f"[SandboxMonitor] Pause callback error: {e}, back to IDLE")
                self._state = MonitorState.IDLE
                self._expiry_monotonic = time.monotonic() + self._default_timeout
                self._expiry_wall = time.time() + self._default_timeout
                self._sequence_number += 1
                self._last_pause_sequence = -1  # Allow retry
                self._write_state()

    def get_remaining_seconds(self) -> float:
        """Get remaining seconds until idle expires"""
        with self._lock:
            return self._get_remaining_seconds_unlocked()

    def _get_remaining_seconds_unlocked(self) -> float:
        """Get remaining seconds (must hold lock)"""
        if self._state != MonitorState.IDLE or self._expiry_monotonic is None:
            return 0.0
        return max(0.0, self._expiry_monotonic - time.monotonic())

    def get_state(self) -> MonitorState:
        """Get current state"""
        with self._lock:
            return self._state

    def is_idle(self) -> bool:
        """Check if in idle state"""
        with self._lock:
            return self._state == MonitorState.IDLE

    def is_active(self) -> bool:
        """Check if in active state"""
        with self._lock:
            return self._state == MonitorState.ACTIVE
