"""
BloomOS Kernel - Wave Scheduler Handlers
========================================

Syscall handlers for wave scheduler operations.

Syscalls:
- scheduler.acquire.v1: Acquire a core from the wave engine
- scheduler.release.v1: Release a core back to the wave engine
- scheduler.status.v1: Get wave engine status
- scheduler.pressure_map.v1: Get visual pressure map of cores
"""

from __future__ import annotations

import time
from pathlib import Path
from typing import Any, Dict, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from ..syscall_registry import SyscallRegistry


# Default state file for wave engine
DEFAULT_STATE_FILE = "/tmp/bloomos_wave_state.json"


def _get_wave_engine(state_file: Optional[str] = None):
    """Get or create wave engine instance."""
    from helix_cdc.scheduler.wave_engine import WaveEngine

    sf = state_file or DEFAULT_STATE_FILE
    return WaveEngine(state_file=sf, seed="bloomos_kernel")


def scheduler_acquire_v1(inputs: Dict[str, Any]) -> Dict[str, Any]:
    """
    Acquire a core from the wave engine.

    This uses the flywheel wave scheduler with pressure dynamics to
    select the optimal core for a job.

    Args (inputs dict):
        weight: Job weight (0.5=small, 1.0=normal, 2.0=large). Default: 1.0
        state_file: Path to wave engine state file (optional)

    Returns:
        {
            "status": "OK" | "FAIL",
            "core_id": int,           # Assigned core
            "weight": float,          # Weight used
            "flywheel_position": int, # Current flywheel position
            "flywheel_momentum": float, # Current momentum
        }
    """
    weight = inputs.get("weight", 1.0)
    state_file = inputs.get("state_file")

    try:
        engine = _get_wave_engine(state_file)
        core_id = engine.acquire(weight=weight)

        return {
            "status": "OK",
            "core_id": core_id,
            "weight": weight,
            "flywheel_position": engine.wave.position,
            "flywheel_momentum": round(engine.wave.momentum, 3),
        }

    except Exception as e:
        return {
            "status": "FAIL",
            "error": str(e),
            "core_id": -1,
            "weight": weight,
        }


def scheduler_release_v1(inputs: Dict[str, Any]) -> Dict[str, Any]:
    """
    Release a core back to the wave engine.

    This records job completion and updates pressure dynamics.

    Args (inputs dict):
        core_id: Core ID to release (required)
        weight: Job weight used (default: 1.0)
        state_file: Path to wave engine state file (optional)

    Returns:
        {
            "status": "OK" | "FAIL",
            "core_id": int,
            "core_pressure": float,    # Pressure after release
            "total_jobs": int,         # Total jobs processed by engine
        }
    """
    core_id = inputs.get("core_id")
    weight = inputs.get("weight", 1.0)
    state_file = inputs.get("state_file")

    if core_id is None:
        return {
            "status": "FAIL",
            "error": "Missing required field: core_id",
            "core_id": -1,
        }

    try:
        engine = _get_wave_engine(state_file)

        # Validate core_id exists in engine
        if core_id not in engine.cores:
            return {
                "status": "FAIL",
                "error": f"Core {core_id} not found in engine (valid cores: {list(engine.cores.keys())})",
                "core_id": core_id,
            }

        engine.release(core_id, weight=weight)

        core_pressure = engine.cores.get(core_id)
        pressure = core_pressure.pressure if core_pressure else 0.0

        return {
            "status": "OK",
            "core_id": core_id,
            "core_pressure": round(pressure, 3),
            "total_jobs": engine.wave.total_jobs,
        }

    except Exception as e:
        return {
            "status": "FAIL",
            "error": str(e),
            "core_id": core_id,
        }


def scheduler_status_v1(inputs: Dict[str, Any]) -> Dict[str, Any]:
    """
    Get wave engine status.

    Returns comprehensive status including flywheel state and per-core metrics.

    Args (inputs dict):
        state_file: Path to wave engine state file (optional)

    Returns:
        {
            "status": "OK" | "FAIL",
            "flywheel": {
                "position": int,
                "momentum": float,
                "total_jobs": int,
                "uptime_s": float,
            },
            "cores": {
                core_id: {
                    "pressure": float,
                    "jobs": int,
                    "active": bool,
                }
            },
            "firing_order": [int],
            "active_jobs": int,
        }
    """
    state_file = inputs.get("state_file")

    try:
        engine = _get_wave_engine(state_file)
        status = engine.status()

        return {
            "status": "OK",
            **status,
        }

    except Exception as e:
        return {
            "status": "FAIL",
            "error": str(e),
        }


def scheduler_pressure_map_v1(inputs: Dict[str, Any]) -> Dict[str, Any]:
    """
    Get visual pressure map of cores.

    Returns a text-based visualization of core pressures.

    Args (inputs dict):
        state_file: Path to wave engine state file (optional)

    Returns:
        {
            "status": "OK" | "FAIL",
            "pressure_map": str,       # ASCII visualization
            "core_pressures": {int: float},  # Raw pressure values
        }
    """
    state_file = inputs.get("state_file")

    try:
        engine = _get_wave_engine(state_file)
        pressure_map = engine.pressure_map()

        core_pressures = {
            core_id: round(core.pressure, 3)
            for core_id, core in engine.cores.items()
        }

        return {
            "status": "OK",
            "pressure_map": pressure_map,
            "core_pressures": core_pressures,
        }

    except Exception as e:
        return {
            "status": "FAIL",
            "error": str(e),
        }


# Cost estimators - Wave scheduler syscalls are cheap (state read/write only)

def scheduler_acquire_cost_estimator(inputs: Dict[str, Any]) -> float:
    """
    Estimate IT cost for scheduler.acquire.v1.

    Cost model:
    - Base cost: 0.01 IT (minimal state read/write)
    - Weight factor: weight * 0.001 (higher weight jobs cost slightly more)
    """
    base_cost = 0.01
    weight = inputs.get("weight", 1.0)
    return base_cost + (weight * 0.001)


def scheduler_release_cost_estimator(inputs: Dict[str, Any]) -> float:
    """
    Estimate IT cost for scheduler.release.v1.

    Cost model:
    - Base cost: 0.01 IT (state write)
    - Weight factor: weight * 0.001
    """
    base_cost = 0.01
    weight = inputs.get("weight", 1.0)
    return base_cost + (weight * 0.001)


def scheduler_status_cost_estimator(inputs: Dict[str, Any]) -> float:
    """
    Estimate IT cost for scheduler.status.v1.

    Cost model:
    - Fixed: 0.005 IT (read-only)
    """
    return 0.005


def scheduler_pressure_map_cost_estimator(inputs: Dict[str, Any]) -> float:
    """
    Estimate IT cost for scheduler.pressure_map.v1.

    Cost model:
    - Fixed: 0.005 IT (read-only)
    """
    return 0.005


# Schema definitions
SCHEDULER_ACQUIRE_V1_SCHEMA = {
    "required": [],
    "types": {
        "weight": "float",
        "state_file": "str",
    },
}

SCHEDULER_RELEASE_V1_SCHEMA = {
    "required": ["core_id"],
    "types": {
        "core_id": "int",
        "weight": "float",
        "state_file": "str",
    },
}

SCHEDULER_STATUS_V1_SCHEMA = {
    "required": [],
    "types": {
        "state_file": "str",
    },
}

SCHEDULER_PRESSURE_MAP_V1_SCHEMA = {
    "required": [],
    "types": {
        "state_file": "str",
    },
}


def register_wave_scheduler_handlers(registry: "SyscallRegistry") -> None:
    """Register wave scheduler handlers with the registry."""
    registry.register(
        canonical_name="scheduler.acquire.v1",
        handler=scheduler_acquire_v1,
        schema=SCHEDULER_ACQUIRE_V1_SCHEMA,
        cost_estimator=scheduler_acquire_cost_estimator,
        description="Acquire a core from the wave engine (flywheel scheduler)",
    )

    registry.register(
        canonical_name="scheduler.release.v1",
        handler=scheduler_release_v1,
        schema=SCHEDULER_RELEASE_V1_SCHEMA,
        cost_estimator=scheduler_release_cost_estimator,
        description="Release a core back to the wave engine",
    )

    registry.register(
        canonical_name="scheduler.status.v1",
        handler=scheduler_status_v1,
        schema=SCHEDULER_STATUS_V1_SCHEMA,
        cost_estimator=scheduler_status_cost_estimator,
        description="Get wave engine status (flywheel, cores, pressure)",
    )

    registry.register(
        canonical_name="scheduler.pressure_map.v1",
        handler=scheduler_pressure_map_v1,
        schema=SCHEDULER_PRESSURE_MAP_V1_SCHEMA,
        cost_estimator=scheduler_pressure_map_cost_estimator,
        description="Get visual pressure map of cores",
    )
