#!/usr/bin/env python3
"""
helix_cdc/runtime/fibpi_runtime.py

FibPi-inspired CPU "Tesla valve" runtime for waves.

This runtime cycles jobs across CPU cores instead of letting PyTorch/OS
pile everything onto one core. Proven by receipts:
- FIBPI_RUNTIME_WAVES_V01: 200 jobs distributed 50/50/50/50 across 4 cores
- HELIX_LOBE_FIBPI_V01: Training steps distributed with preserved accuracy

Environment Variables:
    HELIX_FIBPI_ENABLED  - "1" (default) or "0" to disable core pinning
    HELIX_CORES          - Comma-separated core IDs (e.g. "0,1,2,3")
                           Default: all cores from sched_getaffinity(0)
    HELIX_FIBPI_MODE     - "round_robin" (default), later: "fibpi3d", "se_weighted"
    HELIX_FIBPI_VERBOSE  - "0" (default) or "1" to log chosen core per step

Usage:
    # From environment
    runtime = FibPiRuntime.from_env()

    # Or explicit config
    runtime = FibPiRuntime(FibPiRuntimeConfig(core_pool=[0,1,2,3]))

    # Run a job on next core
    with runtime.on_core() as core:
        do_heavy_work()

    # Or functional style
    result, core = runtime.run(heavy_function)
"""

from __future__ import annotations

import os
import sys
import itertools
import warnings
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Callable, Iterable, List, Optional, Set, Tuple, TypeVar, Generator, TYPE_CHECKING

if TYPE_CHECKING:
    from .se_router import SeRouter

T = TypeVar("T")


def _detect_core_pool() -> List[int]:
    """
    Detect available cores for this process.
    Uses Linux sched_getaffinity if available, falls back to range(os.cpu_count()).

    Returns a sorted list of core IDs.
    """
    try:
        affinity: Set[int] = os.sched_getaffinity(0)  # type: ignore[attr-defined]
        cores = sorted(int(c) for c in affinity)
    except (AttributeError, OSError):
        # Portable fallback: use all logical cores
        cpu_count = os.cpu_count() or 1
        cores = list(range(cpu_count))

    if not cores:
        cores = [0]  # extremely defensive
    return cores


def _parse_cores_env(env_val: str) -> List[int]:
    """Parse HELIX_CORES env var like '0,1,2,3' into sorted list of ints."""
    cores = []
    for part in env_val.split(","):
        part = part.strip()
        if not part:
            continue
        try:
            cores.append(int(part))
        except ValueError:
            warnings.warn(f"HELIX_CORES: ignoring invalid core ID '{part}'")
    return sorted(set(cores)) if cores else _detect_core_pool()


@dataclass
class FibPiRuntimeConfig:
    """
    Configuration for FibPiRuntime.

    Attributes:
        core_pool: List of CPU core IDs to cycle through.
        enabled: If False, on_core() becomes a no-op (no affinity changes).
        mode: Scheduling mode - "round_robin" for now, later "fibpi3d", "se_weighted".
        verbose: If True, log chosen core for each step.
    """

    core_pool: Optional[List[int]] = None
    enabled: bool = True
    mode: str = "round_robin"
    verbose: bool = False


class FibPiRuntime:
    """
    CPU "Tesla valve" runtime that cycles jobs across cores.

    Instead of letting PyTorch/OS pile work onto one core, this runtime
    deliberately distributes work across a core manifold using a firing order.

    Usage:
        runtime = FibPiRuntime.from_env()

        with runtime.on_core() as core:
            # Process pinned to `core` via sched_setaffinity
            do_work()
        # Affinity restored automatically

        # Or functional style:
        result, core = runtime.run(heavy_function, arg1, arg2)
    """

    def __init__(
        self,
        config: Optional[FibPiRuntimeConfig] = None,
        router: Optional["SeRouter"] = None,
    ):
        if config is None:
            config = FibPiRuntimeConfig()

        self._config = config
        self._enabled = config.enabled
        self._verbose = config.verbose
        self._mode = config.mode

        cores = config.core_pool or _detect_core_pool()
        if not cores:
            cores = [0]

        # If only 1 core, disable affinity changes (no point)
        if len(cores) == 1:
            self._enabled = False

        self._cores: List[int] = cores
        self._rr_iter: Iterable[int] = itertools.cycle(self._cores)
        self._last_core: Optional[int] = None

        # Se-based router (optional)
        self._router: Optional["SeRouter"] = router

    @classmethod
    def from_env(cls) -> "FibPiRuntime":
        """
        Create FibPiRuntime from environment variables.

        Environment Variables:
            HELIX_FIBPI_ENABLED  - "1" (default) or "0"
            HELIX_CORES          - Comma-separated core IDs (default: auto-detect)
            HELIX_FIBPI_MODE     - "round_robin" (default)
            HELIX_FIBPI_VERBOSE  - "0" (default) or "1"

        Returns:
            FibPiRuntime configured from environment.
        """
        enabled_str = os.environ.get("HELIX_FIBPI_ENABLED", "1").strip()
        enabled = enabled_str not in ("0", "false", "False", "FALSE", "no", "No", "NO")

        cores_str = os.environ.get("HELIX_CORES", "").strip()
        if cores_str:
            core_pool = _parse_cores_env(cores_str)
        else:
            core_pool = _detect_core_pool()

        mode = os.environ.get("HELIX_FIBPI_MODE", "round_robin").strip()

        verbose_str = os.environ.get("HELIX_FIBPI_VERBOSE", "0").strip()
        verbose = verbose_str in ("1", "true", "True", "TRUE", "yes", "Yes", "YES")

        config = FibPiRuntimeConfig(
            core_pool=core_pool,
            enabled=enabled,
            mode=mode,
            verbose=verbose,
        )

        return cls(config)

    @property
    def cores(self) -> List[int]:
        """Return copy of core pool."""
        return list(self._cores)

    @property
    def enabled(self) -> bool:
        """Return whether FibPi core pinning is active."""
        return self._enabled

    @property
    def last_core(self) -> Optional[int]:
        """Return the last core used (for tracking)."""
        return self._last_core

    @property
    def mode(self) -> str:
        """Return scheduling mode."""
        return self._mode

    @property
    def router(self) -> Optional["SeRouter"]:
        """Return the current Se-router (if any)."""
        return self._router

    def set_router(self, router: Optional["SeRouter"]) -> None:
        """
        Set the Se-based router for intelligent core selection.

        Args:
            router: SeRouter instance, or None to disable Se-routing.
        """
        self._router = router

    def next_core_for_se(self, se: float) -> int:
        """
        Get the next core based on an Se (Symbolic Entropy) score.

        If a SeRouter is configured, uses Se bucketing to select the core.
        Otherwise falls back to round-robin.

        Args:
            se: Symbolic entropy value for the current job.

        Returns:
            Core ID to use.
        """
        if self._router is not None:
            return self._router.observe(se)
        else:
            # Fallback to round-robin
            return self.next_core()

    def next_core(self) -> int:
        """Return the next core in the firing order (currently round-robin)."""
        return next(self._rr_iter)  # type: ignore[return-value]

    @contextmanager
    def on_core(self, core: Optional[int] = None) -> Generator[int, None, None]:
        """
        Context manager that temporarily pins this process to a single core.

        If FibPi is disabled or only 1 core exists, this becomes a no-op
        that still yields a core ID for tracking purposes.

        Args:
            core: Specific core to pin to. If None, uses next_core().

        Yields:
            The core ID being used (even if affinity wasn't changed).
        """
        if core is None:
            core = self.next_core()

        self._last_core = core

        if self._verbose:
            print(f"[FIBPI] step on core {core}", file=sys.stderr)

        # If disabled, just yield without changing affinity
        if not self._enabled:
            yield core
            return

        # Try to change affinity
        try:
            old_affinity = os.sched_getaffinity(0)  # type: ignore[attr-defined]
            restore_affinity = True
        except (AttributeError, OSError):
            old_affinity = None
            restore_affinity = False

        if restore_affinity:
            try:
                os.sched_setaffinity(0, {core})  # type: ignore[attr-defined]
            except (AttributeError, OSError, PermissionError):
                # Can't set affinity, just run anyway
                restore_affinity = False

        try:
            yield core
        finally:
            if restore_affinity and old_affinity is not None:
                try:
                    os.sched_setaffinity(0, old_affinity)  # type: ignore[attr-defined]
                except (AttributeError, OSError, PermissionError):
                    pass  # Best effort restore

    def run(self, fn: Callable[..., T], *args, **kwargs) -> Tuple[T, int]:
        """
        Run fn(*args, **kwargs) on the next core in the firing order.

        Returns:
            (result, core_id)
        """
        with self.on_core() as core:
            result = fn(*args, **kwargs)
        return result, core

    def __repr__(self) -> str:
        return (
            f"FibPiRuntime(cores={self._cores}, enabled={self._enabled}, "
            f"mode='{self._mode}')"
        )


if __name__ == "__main__":
    # Self-test: instantiate from env and print status
    print("FibPiRuntime self-test")
    print("=" * 40)

    runtime = FibPiRuntime.from_env()

    print(f"Detected cores: {runtime.cores}")
    print(f"FibPi enabled:  {runtime.enabled}")
    print(f"Mode:           {runtime.mode}")
    print()

    # Test a few cycles
    print("Testing 8 core cycles:")
    for i in range(8):
        with runtime.on_core() as core:
            print(f"  Step {i+1}: core {core}")

    print()
    print("Self-test complete.")
