"""
helix_cdc/scheduler/wave_engine.py
==================================
Flywheel wave scheduler with pressure dynamics.

Instead of static core assignment, this models:
- Pressure buildup on cores (job load)
- Exhaust/cooldown after job completion
- Flywheel momentum (persistent scheduler state)
- Gear ratios (job size → scheduling granularity)

Usage:
    from helix_cdc.scheduler.wave_engine import WaveEngine

    engine = WaveEngine()
    core = engine.acquire()      # Get next core in wave
    engine.release(core, job_weight=1.0)  # Release with pressure

    # Or use context manager:
    with engine.job(weight=2.0) as core:
        run_on_core(core)
"""

import os
import time
import math
import json
import threading
import fcntl
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Any
from dataclasses import dataclass, field
from contextlib import contextmanager

# Import FibPi3D for base firing order
try:
    from helix_cdc.scheduler.wave_scheduler import get_runtime_schedule
    WAVE_AVAILABLE = True
except ImportError:
    WAVE_AVAILABLE = False


@dataclass
class CoreState:
    """State of a single core in the wave engine."""
    core_id: int
    pressure: float = 0.0          # Current "heat" (0 = cold, 1 = hot)
    last_fire_time: float = 0.0    # When it last ran a job
    jobs_completed: int = 0         # Total jobs run
    total_weight: float = 0.0       # Cumulative job weight


@dataclass
class WaveState:
    """Flywheel state - persists across jobs."""
    position: int = 0              # Current position in firing order
    momentum: float = 1.0          # Flywheel speed (higher = faster cycling)
    total_jobs: int = 0            # Jobs processed
    created_at: float = field(default_factory=time.time)


class WaveEngine:
    """
    Flywheel wave scheduler with pressure dynamics.

    The engine maintains momentum between jobs. Each job:
    1. Acquires a core (based on pressure + firing order)
    2. Builds pressure on that core during execution
    3. Releases with exhaust that affects neighboring cores
    4. Advances the flywheel position

    Gear ratios:
    - weight=0.5: Small job, quick cycle, minimal pressure
    - weight=1.0: Normal job, standard cycle
    - weight=2.0: Large job, slower cycle, more pressure buildup
    """

    # Pressure dynamics constants
    PRESSURE_DECAY_RATE = 0.1      # Pressure decay per second
    EXHAUST_SPREAD = 0.3           # How much exhaust spreads to neighbors
    MOMENTUM_GAIN = 0.05           # Momentum increase per job
    MOMENTUM_DECAY = 0.01          # Momentum decay per second idle
    MIN_MOMENTUM = 0.5             # Minimum flywheel speed
    MAX_MOMENTUM = 3.0             # Maximum flywheel speed

    def __init__(
        self,
        num_cores: Optional[int] = None,
        state_file: Optional[str] = None,
        seed: str = "wave_engine",
    ):
        """
        Initialize wave engine.

        Args:
            num_cores: Number of cores (default: auto-detect)
            state_file: Path to persist flywheel state (optional)
            seed: Seed for FibPi3D firing order
        """
        self.seed = seed
        self.state_file = Path(state_file) if state_file else None
        self.lock = threading.Lock()

        # Get FibPi3D firing order
        if WAVE_AVAILABLE:
            schedule = get_runtime_schedule(seed=seed)
            self.firing_order = schedule["firing_order"]
            self.num_cores = len(self.firing_order)
        else:
            self.num_cores = num_cores or os.cpu_count() or 4
            self.firing_order = list(range(self.num_cores))

        # Initialize core states
        self.cores: Dict[int, CoreState] = {
            core_id: CoreState(core_id=core_id)
            for core_id in self.firing_order
        }

        # Initialize or load flywheel state
        self.wave = self._load_state() or WaveState()

        # Track active jobs
        self.active_jobs: Dict[int, float] = {}  # core_id -> start_time

    def _load_state(self) -> Optional[WaveState]:
        """Load flywheel state from file."""
        if not self.state_file or not self.state_file.exists():
            return None
        try:
            data = json.loads(self.state_file.read_text())
            state = WaveState(
                position=data.get("position", 0),
                momentum=data.get("momentum", 1.0),
                total_jobs=data.get("total_jobs", 0),
                created_at=data.get("created_at", time.time()),
            )
            # Decay momentum based on idle time
            idle_time = time.time() - data.get("last_save", time.time())
            state.momentum = max(
                self.MIN_MOMENTUM,
                state.momentum - idle_time * self.MOMENTUM_DECAY
            )

            # CRITICAL: Also load core pressure states for cross-process coordination
            core_states = data.get("core_states", {})
            for core_id_str, core_data in core_states.items():
                core_id = int(core_id_str)
                if core_id in self.cores:
                    # Apply time-decayed pressure from saved state
                    saved_pressure = core_data.get("pressure", 0.0)
                    saved_time = core_data.get("last_fire_time", 0.0)
                    if saved_time > 0:
                        elapsed = time.time() - saved_time
                        decayed_pressure = max(0.0, saved_pressure - elapsed * self.PRESSURE_DECAY_RATE)
                    else:
                        decayed_pressure = saved_pressure
                    self.cores[core_id].pressure = decayed_pressure
                    self.cores[core_id].last_fire_time = saved_time
                    self.cores[core_id].jobs_completed = core_data.get("jobs_completed", 0)
                    self.cores[core_id].total_weight = core_data.get("total_weight", 0.0)

            return state
        except Exception:
            return None

    def _save_state(self) -> None:
        """Persist flywheel state to file."""
        if not self.state_file:
            return
        try:
            self.state_file.parent.mkdir(parents=True, exist_ok=True)

            # CRITICAL: Also save core pressure states for cross-process coordination
            core_states = {
                str(core_id): {
                    "pressure": core.pressure,
                    "last_fire_time": core.last_fire_time,
                    "jobs_completed": core.jobs_completed,
                    "total_weight": core.total_weight,
                }
                for core_id, core in self.cores.items()
            }

            self.state_file.write_text(json.dumps({
                "position": self.wave.position,
                "momentum": self.wave.momentum,
                "total_jobs": self.wave.total_jobs,
                "created_at": self.wave.created_at,
                "last_save": time.time(),
                "core_states": core_states,  # NEW: persist pressure for cross-process
            }))
        except Exception:
            pass

    @contextmanager
    def _file_locked(self):
        """
        Context manager for atomic file-locked operations.

        This is CRITICAL for cross-process coordination:
        1. Acquires exclusive file lock (blocks other processes)
        2. Reloads state from file (gets latest from other processes)
        3. Yields for caller to make changes
        4. Saves state to file
        5. Releases lock
        """
        if not self.state_file:
            # No state file = in-process only, just use threading lock
            with self.lock:
                yield
            return

        lock_path = self.state_file.with_suffix('.lock')
        lock_path.parent.mkdir(parents=True, exist_ok=True)

        with open(lock_path, 'w') as lock_file:
            # Acquire exclusive lock (blocks until available)
            fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
            try:
                # Reload state from file (another process may have updated it)
                self._reload_state_from_file()

                yield

                # Save updated state
                self._save_state()
            finally:
                # Release lock
                fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)

    def _reload_state_from_file(self) -> None:
        """Reload state from file (called within file lock)."""
        if not self.state_file or not self.state_file.exists():
            return
        try:
            data = json.loads(self.state_file.read_text())

            # Update flywheel state
            self.wave.position = data.get("position", self.wave.position)
            self.wave.momentum = data.get("momentum", self.wave.momentum)
            self.wave.total_jobs = data.get("total_jobs", self.wave.total_jobs)

            # Update core states
            core_states = data.get("core_states", {})
            for core_id_str, core_data in core_states.items():
                core_id = int(core_id_str)
                if core_id in self.cores:
                    saved_pressure = core_data.get("pressure", 0.0)
                    saved_time = core_data.get("last_fire_time", 0.0)
                    # Apply time decay
                    if saved_time > 0:
                        elapsed = time.time() - saved_time
                        decayed = max(0.0, saved_pressure - elapsed * self.PRESSURE_DECAY_RATE)
                    else:
                        decayed = saved_pressure
                    self.cores[core_id].pressure = decayed
                    self.cores[core_id].last_fire_time = saved_time
                    self.cores[core_id].jobs_completed = core_data.get("jobs_completed", 0)
                    self.cores[core_id].total_weight = core_data.get("total_weight", 0.0)
        except Exception:
            pass

    def _decay_pressure(self) -> None:
        """Apply pressure decay to all cores based on time."""
        now = time.time()
        for core in self.cores.values():
            if core.pressure > 0:
                elapsed = now - core.last_fire_time
                decay = elapsed * self.PRESSURE_DECAY_RATE
                core.pressure = max(0.0, core.pressure - decay)

    def _find_best_core(self) -> int:
        """
        Find best core considering:
        1. FibPi3D firing order (wave position)
        2. Current pressure (prefer cooler cores)
        3. Momentum (how far ahead to look)
        """
        self._decay_pressure()

        # Look ahead based on momentum
        lookahead = max(1, int(self.wave.momentum * 2))

        best_core = None
        best_score = float('inf')

        for i in range(lookahead):
            pos = (self.wave.position + i) % self.num_cores
            core_id = self.firing_order[pos]
            core = self.cores[core_id]

            # Score = pressure + position penalty
            # Lower pressure = better, earlier in lookahead = better
            position_penalty = i * 0.1
            score = core.pressure + position_penalty

            # Skip if core is actively running a job
            if core_id in self.active_jobs:
                score += 10.0  # Heavy penalty

            if score < best_score:
                best_score = score
                best_core = core_id

        return best_core or self.firing_order[self.wave.position]

    def acquire(self, weight: float = 1.0) -> int:
        """
        Acquire a core for a job.

        Args:
            weight: Job weight (0.5=small, 1.0=normal, 2.0=large)

        Returns:
            Core ID to run on

        Uses file locking for cross-process coordination when state_file is set.
        """
        with self._file_locked():
            core_id = self._find_best_core()

            # Mark as active
            self.active_jobs[core_id] = time.time()

            # Apply initial pressure based on weight
            self.cores[core_id].pressure += weight * 0.5

            # Advance position to avoid next worker picking same core
            self.wave.position = (self.wave.position + 1) % self.num_cores

            return core_id

    def release(self, core_id: int, weight: float = 1.0) -> None:
        """
        Release a core after job completion.

        Args:
            core_id: The core that was used
            weight: Job weight (affects exhaust pressure)

        Uses file locking for cross-process coordination when state_file is set.
        """
        with self._file_locked():
            now = time.time()
            core = self.cores[core_id]

            # Calculate job duration
            start_time = self.active_jobs.pop(core_id, now)
            duration = now - start_time

            # Apply exhaust pressure
            exhaust = weight * (1.0 + duration * 0.1)
            core.pressure = min(1.0, core.pressure + exhaust * 0.3)
            core.last_fire_time = now
            core.jobs_completed += 1
            core.total_weight += weight

            # Spread exhaust to neighbors in firing order
            pos = self.firing_order.index(core_id)
            for offset in [-1, 1]:
                neighbor_pos = (pos + offset) % self.num_cores
                neighbor_id = self.firing_order[neighbor_pos]
                self.cores[neighbor_id].pressure += exhaust * self.EXHAUST_SPREAD * 0.5

            # Position already advanced in acquire(), just increment total_jobs
            self.wave.total_jobs += 1

            # Build momentum
            self.wave.momentum = min(
                self.MAX_MOMENTUM,
                self.wave.momentum + self.MOMENTUM_GAIN
            )
            # State saved by _file_locked() context manager

    @contextmanager
    def job(self, weight: float = 1.0):
        """
        Context manager for running a job.

        Usage:
            with engine.job(weight=2.0) as core:
                subprocess.run(["taskset", "-c", str(core), ...])
        """
        core = self.acquire(weight)
        try:
            yield core
        finally:
            self.release(core, weight)

    def status(self) -> Dict[str, Any]:
        """Get engine status."""
        self._decay_pressure()
        return {
            "flywheel": {
                "position": self.wave.position,
                "momentum": round(self.wave.momentum, 2),
                "total_jobs": self.wave.total_jobs,
                "uptime_s": round(time.time() - self.wave.created_at, 1),
            },
            "cores": {
                core_id: {
                    "pressure": round(core.pressure, 3),
                    "jobs": core.jobs_completed,
                    "active": core_id in self.active_jobs,
                }
                for core_id, core in self.cores.items()
            },
            "firing_order": self.firing_order,
            "active_jobs": len(self.active_jobs),
        }

    def pressure_map(self) -> str:
        """Visual pressure map of cores."""
        self._decay_pressure()

        lines = ["Wave Engine Pressure Map", "=" * 40]

        # Flywheel indicator
        lines.append(f"Flywheel: pos={self.wave.position} momentum={self.wave.momentum:.2f}")
        lines.append("")

        # Core pressure bars
        for i, core_id in enumerate(self.firing_order):
            core = self.cores[core_id]
            bar_len = int(core.pressure * 20)
            bar = "█" * bar_len + "░" * (20 - bar_len)

            # Mark current position
            marker = "→" if i == self.wave.position else " "
            active = "*" if core_id in self.active_jobs else " "

            lines.append(f"{marker}[{core_id:2d}]{active} [{bar}] {core.pressure:.2f}")

        return "\n".join(lines)


# Singleton engine for module-level use
_ENGINE: Optional[WaveEngine] = None
_ENGINE_LOCK = threading.Lock()


def get_engine(
    state_file: Optional[str] = None,
    seed: str = "wave_engine",
) -> WaveEngine:
    """Get or create the singleton wave engine."""
    global _ENGINE
    with _ENGINE_LOCK:
        if _ENGINE is None:
            default_state = os.getenv("HELIX_WAVE_STATE", "/tmp/helix_wave_state.json")
            _ENGINE = WaveEngine(
                state_file=state_file or default_state,
                seed=seed,
            )
        return _ENGINE


def acquire(weight: float = 1.0) -> int:
    """Acquire a core from the global engine."""
    return get_engine().acquire(weight)


def release(core: int, weight: float = 1.0) -> None:
    """Release a core to the global engine."""
    get_engine().release(core, weight)


@contextmanager
def job(weight: float = 1.0):
    """Context manager for global engine job."""
    with get_engine().job(weight) as core:
        yield core


if __name__ == "__main__":
    # Demo
    import subprocess

    engine = WaveEngine()
    print(engine.pressure_map())
    print()

    # Simulate jobs
    for i in range(8):
        weight = 0.5 + (i % 3) * 0.5  # Vary weight
        with engine.job(weight=weight) as core:
            print(f"Job {i}: core={core}, weight={weight}")
            time.sleep(0.1)  # Simulate work

    print()
    print(engine.pressure_map())
    print()
    print(json.dumps(engine.status(), indent=2))
