"""
helix_cdc/runtime.py
====================
Unified Helix runtime with FibPi3D wave scheduling → C path routing.

Usage:
    from helix_cdc.runtime import helix_run, helix_regen, helix_infer

    # Auto-routes to C binary with wave scheduling
    result = helix_regen(seed="my_seed", size="10MB")
    result = helix_infer(prompt="What is entropy?")

    # Or use the generic dispatcher
    result = helix_run("regen", seed="my_seed", size="1MB")
    result = helix_run("infer", prompt="test", layer=0)

Environment variables:
    HELIX_WAVE_N=4          # FibPi dimensions (default 4)
    HELIX_WAVE_SEED=helix   # Wave schedule seed
    HELIX_USE_C=1           # Force C path (default: auto)
    HELIX_VERBOSE=1         # Print wave schedule info
"""

import os
import sys
import json
import subprocess
import time
from pathlib import Path
from typing import Dict, Any, Optional, List, Callable
from functools import wraps

# Find repo root
_REPO = Path(__file__).resolve().parent.parent
_C_BIN = _REPO / "c" / "bin"

# Import wave scheduler
try:
    from helix_cdc.scheduler.wave_scheduler import (
        get_runtime_schedule,
        engine_firing_order,
        get_cpu_count,
    )
    WAVE_AVAILABLE = True
except ImportError:
    WAVE_AVAILABLE = False
    get_runtime_schedule = None


# Global wave schedule (initialized once)
_WAVE_SCHEDULE: Optional[Dict[str, Any]] = None


def _init_wave_schedule() -> Dict[str, Any]:
    """Initialize FibPi3D wave schedule for this process."""
    global _WAVE_SCHEDULE

    if _WAVE_SCHEDULE is not None:
        return _WAVE_SCHEDULE

    if not WAVE_AVAILABLE:
        _WAVE_SCHEDULE = {"error": "wave_scheduler not available"}
        return _WAVE_SCHEDULE

    N = int(os.getenv("HELIX_WAVE_N", "4"))
    seed = os.getenv("HELIX_WAVE_SEED", "helix_runtime")

    _WAVE_SCHEDULE = get_runtime_schedule(N=N, seed=seed)

    if os.getenv("HELIX_VERBOSE"):
        print(f"[helix_runtime] FibPi3D wave schedule initialized", file=sys.stderr)
        print(f"  cores: {_WAVE_SCHEDULE['num_cores']}", file=sys.stderr)
        print(f"  firing_order: {_WAVE_SCHEDULE['firing_order']}", file=sys.stderr)
        print(f"  N={N}, seed={seed}", file=sys.stderr)

    return _WAVE_SCHEDULE


def get_wave_schedule() -> Dict[str, Any]:
    """Get the current FibPi3D wave schedule."""
    return _init_wave_schedule()


def set_cpu_affinity_from_wave(core_index: int = 0) -> bool:
    """
    Set CPU affinity based on wave schedule firing order.

    Args:
        core_index: Which position in firing order (0 = first to fire)

    Returns:
        True if affinity was set, False otherwise
    """
    try:
        import psutil
    except ImportError:
        return False

    schedule = get_wave_schedule()
    if "error" in schedule:
        return False

    firing_order = schedule.get("firing_order", [])
    if not firing_order or core_index >= len(firing_order):
        return False

    target_core = firing_order[core_index]
    try:
        psutil.Process().cpu_affinity([target_core])
        return True
    except Exception:
        return False


def _run_c_binary(
    binary_name: str,
    args: List[str],
    timeout: int = 60,
    env_extra: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
    """Run a C binary from helix-cdc/c/bin/"""
    binary_path = _C_BIN / binary_name

    if not binary_path.exists():
        return {
            "status": "error",
            "error": f"C binary not found: {binary_path}",
            "backend": "c_missing",
        }

    env = os.environ.copy()
    if env_extra:
        env.update(env_extra)

    # Pass wave schedule via env
    schedule = get_wave_schedule()
    if "error" not in schedule:
        env["HELIX_WAVE_SCHEDULE_JSON"] = json.dumps(schedule)

    cmd = [str(binary_path)] + args

    try:
        t0 = time.perf_counter()
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=timeout,
            env=env,
        )
        elapsed = time.perf_counter() - t0

        if result.returncode == 0:
            return {
                "status": "success",
                "backend": f"c_{binary_name}",
                "output": result.stdout.strip(),
                "stderr": result.stderr.strip(),
                "elapsed_s": elapsed,
                "wave_schedule": schedule,
            }
        else:
            return {
                "status": "error",
                "backend": f"c_{binary_name}",
                "returncode": result.returncode,
                "stderr": result.stderr.strip(),
                "elapsed_s": elapsed,
            }

    except subprocess.TimeoutExpired:
        return {
            "status": "error",
            "backend": f"c_{binary_name}",
            "error": f"timeout ({timeout}s)",
        }
    except Exception as e:
        return {
            "status": "error",
            "backend": f"c_{binary_name}",
            "error": str(e),
        }


def helix_regen(
    seed: str,
    size: str = "1MB",
    output_path: Optional[str] = None,
    use_fibpi: bool = True,
    timeout: int = 60,
) -> Dict[str, Any]:
    """
    Regenerate data from seed using C helix_regen_cli.

    Args:
        seed: Regen seed string
        size: Output size (e.g., "1MB", "10MB", "100MB")
        output_path: Output file path (default: /tmp/helix_regen_{seed[:8]}.bin)
        use_fibpi: Use FibPi per-block seeding
        timeout: Timeout in seconds

    Returns:
        Dict with status, output_path, sha256, throughput, wave_schedule
    """
    if output_path is None:
        safe_seed = seed[:8].replace("/", "_").replace(":", "_")
        output_path = f"/tmp/helix_regen_{safe_seed}.bin"

    args = ["--seed", seed, "--size", size, "--out", output_path]
    if not use_fibpi:
        args.append("--no-fibpi")

    result = _run_c_binary("helix_regen_cli", args, timeout=timeout)

    if result["status"] == "success":
        result["output_path"] = output_path
        result["seed"] = seed
        result["size"] = size

    return result


def helix_infer(
    prompt: str,
    cdna_path: Optional[str] = None,
    genome_path: Optional[str] = None,
    layer: int = 0,
    timeout: int = 60,
) -> Dict[str, Any]:
    """
    Run compressed-domain inference using C echo_infer.

    Args:
        prompt: Input prompt
        cdna_path: Path to CDNA file (default: artifacts/mistral_test.cdna)
        genome_path: Path to genome manifest JSON
        layer: Layer index
        timeout: Timeout in seconds

    Returns:
        Dict with status, output, float32_reduction, wave_schedule
    """
    if cdna_path is None:
        cdna_path = str(_REPO / "artifacts" / "mistral_test.cdna")
    if genome_path is None:
        genome_path = str(_REPO / "artifacts" / "mistral_genome.json")

    if not Path(cdna_path).exists():
        return {"status": "error", "error": f"CDNA not found: {cdna_path}"}
    if not Path(genome_path).exists():
        return {"status": "error", "error": f"Genome manifest not found: {genome_path}"}

    args = [
        "--cdna", cdna_path,
        "--genome", genome_path,
        "--layer", str(layer),
        "--prompt", prompt,
    ]

    return _run_c_binary("echo_infer", args, timeout=timeout)


def helix_run(op: str, **kwargs) -> Dict[str, Any]:
    """
    Generic Helix runtime dispatcher.

    Args:
        op: Operation type ("regen", "infer")
        **kwargs: Operation-specific arguments

    Returns:
        Operation result dict
    """
    if op == "regen":
        return helix_regen(
            seed=kwargs.get("seed", "default"),
            size=kwargs.get("size", "1MB"),
            output_path=kwargs.get("output_path"),
            use_fibpi=kwargs.get("use_fibpi", True),
            timeout=kwargs.get("timeout", 60),
        )
    elif op == "infer":
        return helix_infer(
            prompt=kwargs.get("prompt", ""),
            cdna_path=kwargs.get("cdna_path"),
            genome_path=kwargs.get("genome_path"),
            layer=kwargs.get("layer", 0),
            timeout=kwargs.get("timeout", 60),
        )
    else:
        return {"status": "error", "error": f"Unknown operation: {op}"}


def with_wave_schedule(func: Callable) -> Callable:
    """
    Decorator that ensures wave schedule is initialized before function runs.
    Also sets CPU affinity if psutil is available.
    """
    @wraps(func)
    def wrapper(*args, **kwargs):
        _init_wave_schedule()
        set_cpu_affinity_from_wave(0)
        return func(*args, **kwargs)
    return wrapper


# Auto-initialize on import if HELIX_WAVE_INIT=1
if os.getenv("HELIX_WAVE_INIT") == "1":
    _init_wave_schedule()
