#!/usr/bin/env python3
"""
helix_cdc/scheduler/regen_fibpi_scheduler.py
============================================

FibPi3D-based scheduler for tensor regeneration jobs.

Wraps wave_scheduler + UDR to parallelize regeneration work across CPU/GPU/QPU
using FibPi geometry for coherent spatial scheduling.

CRITICAL: This does NO regen math itself. It only:
  - Uses FibPi wave schedule for chunking/ordering
  - Uses UDR for backend selection (CPU/GPU/QPU)
  - Executes via user-provided regen function (callback)

Usage:
    from helix_cdc.scheduler.regen_fibpi_scheduler import schedule_regen_jobs

    jobs = [
        {
            "name": "output.weight",
            "offset": 0,
            "size": 65536,
            "fn": my_regen_function,  # callable(seed_ctx, job) -> bytes
            "seed_ctx": {"dna_seed": "...", "entropy_seed": "...", "atlas": ...}
        },
        ...
    ]

    results = schedule_regen_jobs(jobs, topology="fibpi", max_workers=None)
    # Results: list of dicts with "bytes", "sha256", "backend" keys
"""

from __future__ import annotations

import os
import time
import hashlib
import subprocess
import tempfile
import multiprocessing as mp
from typing import List, Dict, Any, Callable, Optional, TYPE_CHECKING
from dataclasses import dataclass

from .wave_scheduler import engine_firing_order, get_cpu_count
from helix_cdc.drivers.udr import UniversalDriver

if TYPE_CHECKING:
    from .codebook_manager import CodebookManager


@dataclass
class RegenJobResult:
    """Result from regenerating a single tensor chunk."""
    name: str
    offset: int
    size: int
    backend: str
    bytes_data: bytes
    sha256: str
    duration_s: float
    core_id: int


def schedule_regen_jobs(
    jobs: List[Dict[str, Any]],
    topology: str = "fibpi",
    max_workers: Optional[int] = None,
    N: int = 4,
    seed: str = "helix_regen",
    use_c_backend: bool = False,
    codebook_manager: Optional['CodebookManager'] = None,
    total_size: Optional[int] = None
) -> List[Dict[str, Any]]:
    """
    Schedule regen jobs across CPU/GPU/QPU using FibPi3D wave scheduler + UDR.

    Args:
        jobs: List of job dicts, each containing:
            - name: Tensor name
            - offset: Byte offset
            - size: Size in bytes
            - fn: Regen function (seed_ctx, job) -> bytes
            - seed_ctx: Context dict (dna_seed, entropy_seed, atlas, etc.)
        topology: Scheduling topology ("fibpi" or "linear")
        max_workers: Max parallel workers (None = auto-detect)
        N: FibPi dimensionality (default 4D)
        seed: Deterministic seed for FibPi schedule
        use_c_backend: If True, use C binary (helix_regen_cli) for regen (144 MB/s)
                       If False, use Python functions (default)
        codebook_manager: Optional CodebookManager for zone-based codebook morphing.
                          If provided with a zone_map, each job gets zone-specific codebook.
        total_size: Total size for zone calculations (auto-detected if not provided)

    Returns:
        List of results (same order as input), each with:
            - All original job fields
            - backend: "cpu" | "gpu" | "qpu"
            - bytes: Regenerated bytes
            - sha256: Hash of bytes
            - duration_s: Time taken
            - core_id: Core that executed
            - implementation: "c" | "python"
            - codebook_zone: Zone name if codebook_manager provided

    Raises:
        RuntimeError: If non-CPU backend requested but not available
    """
    if not jobs:
        return []

    if max_workers is None:
        max_workers = max(1, get_cpu_count() - 1)

    num_jobs = len(jobs)

    # Calculate total_size for zone mapping if not provided
    if total_size is None and codebook_manager is not None:
        total_size = sum(job.get("size", 0) for job in jobs)
        if jobs:
            # Use max offset + size as total
            max_end = max(job.get("offset", 0) + job.get("size", 0) for job in jobs)
            total_size = max(total_size, max_end)

    # Inject zone-specific codebooks if codebook_manager is provided
    if codebook_manager is not None:
        for job in jobs:
            offset = job.get("offset", 0)
            zone_name = codebook_manager.get_zone_name(offset, total_size or 1)
            zone_codebook = codebook_manager.get_for_zone(offset, total_size or 1)

            # Add codebook to seed_ctx for the regen function
            if "seed_ctx" not in job:
                job["seed_ctx"] = {}
            job["seed_ctx"]["codebook"] = zone_codebook
            job["seed_ctx"]["codebook_zone"] = zone_name
            job["codebook_zone"] = zone_name  # Also at job level for logging

    # 1) Get FibPi firing order for job assignment
    if topology == "fibpi":
        firing_order = engine_firing_order(
            num_cores=min(num_jobs, max_workers),
            N=N,
            seed=seed
        )
        # Map job indices to core firing order
        job_to_core = {}
        for i in range(num_jobs):
            core_idx = i % len(firing_order)
            job_to_core[i] = firing_order[core_idx]
    else:
        # Linear fallback
        job_to_core = {i: i % max_workers for i in range(num_jobs)}

    # 2) Group jobs by assigned core (wave-front batches)
    core_batches: Dict[int, List[tuple]] = {}
    for i, job in enumerate(jobs):
        core_id = job_to_core[i]
        if core_id not in core_batches:
            core_batches[core_id] = []
        core_batches[core_id].append((i, job))

    # 3) Initialize UDR for backend routing (CPU/GPU/QPU)
    udr = UniversalDriver()

    # 4) Execute batches in parallel (multiprocessing)
    #    For now: CPU-only execution, but route via UDR for future GPU/QPU
    all_results: Dict[int, Dict[str, Any]] = {}

    # Flatten batches into (job_idx, job, core_id, udr, use_c_backend) tuples
    work_items = []
    for core_id, batch in core_batches.items():
        for job_idx, job in batch:
            work_items.append((job_idx, job, core_id, udr, use_c_backend))

    # Execute in parallel
    with mp.Pool(processes=min(len(work_items), max_workers)) as pool:
        results = pool.map(_execute_single_regen, work_items)

    # 5) Collect results by original job index
    for job_idx, result in results:
        all_results[job_idx] = result

    # 6) Return in original order
    return [all_results[i] for i in range(num_jobs)]


def _regen_with_c_backend(job: Dict[str, Any]) -> bytes:
    """
    Regenerate tensor chunk using C backend (helix_regen_cli binary).

    R1 FIX (2025-12): Now uses context mode with dna/entropy/label/offset
    to match Python's _rng_from_context() semantics.

    Args:
        job: Job dict containing seed_ctx and size

    Returns:
        Regenerated bytes

    Raises:
        RuntimeError: If C binary fails or missing
    """
    seed_ctx = job.get("seed_ctx", {})
    dna_hex = seed_ctx.get("dna_hex", "")
    entropy_seed = seed_ctx.get("entropy_seed", "")
    label = seed_ctx.get("label", "")
    offset = seed_ctx.get("offset", 0)

    size = job["size"]

    # Find C binary
    helix_cdc_path = os.path.expanduser("~/helix-cdc")
    c_binary = os.path.join(helix_cdc_path, "tools", "helix_regen_cli")

    if not os.path.exists(c_binary):
        raise RuntimeError(
            f"C binary not found at {c_binary}. "
            "Build it first or disable use_c_backend."
        )

    # Use temp file for output
    with tempfile.NamedTemporaryFile(mode='wb', delete=False, suffix='.bin') as tmpf:
        tmpfile_path = tmpf.name

    try:
        # Call C binary with context mode (matches Python RNG semantics)
        # Context: sha256(dna_hex|entropy_hex|label|offset)
        cmd = [
            c_binary,
            "--dna", dna_hex,
            "--entropy", entropy_seed,
            "--label", label,
            "--offset", str(offset),
            "--size", str(size),
            "--out", tmpfile_path
        ]

        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=300  # 5 min timeout
        )

        if result.returncode != 0:
            raise RuntimeError(
                f"C regen failed (exit {result.returncode}): {result.stderr}"
            )

        # Read generated bytes
        with open(tmpfile_path, 'rb') as f:
            regen_bytes = f.read()

        return regen_bytes

    finally:
        # Clean up temp file
        if os.path.exists(tmpfile_path):
            os.unlink(tmpfile_path)


def _execute_single_regen(args: tuple) -> tuple:
    """
    Execute single regen job (multiprocessing worker).

    Args:
        args: (job_idx, job, core_id, udr, use_c_backend)

    Returns:
        (job_idx, result_dict)
    """
    job_idx, job, core_id, udr, use_c_backend = args

    t0 = time.time()

    # Route via UDR (backend selection based on Se)
    # For regen, we treat as "inference" type job
    udr_job = {
        "version": "udr-v0.1",
        "problem": {
            "type": "inference",
            "encoding": "helix64",
            "payload": f"{job['name']}:{job['offset']}:{job['size']}",
            "shape": {"tokens": job['size'] // 4}  # Rough estimate for Se calc
        },
        "policy": {"prefs": ["cpu"]}  # For now, CPU only (GPU/QPU future work)
    }

    routing_preview = udr.preview(udr_job)
    backend = routing_preview["backend"]

    # If non-CPU requested, fail loud
    if backend != "cpu":
        raise RuntimeError(
            f"Non-CPU backend requested in regen scheduler: {backend}\n"
            f"GPU/QPU regen not yet implemented. Set UDR_ALLOW_HW=0 or fix routing."
        )

    # Execute regen (C backend or Python backend)
    if use_c_backend:
        regen_bytes = _regen_with_c_backend(job)
        implementation = "c"
    else:
        fn = job["fn"]
        seed_ctx = job.get("seed_ctx", {})

        try:
            regen_bytes = fn(seed_ctx, job)
        except Exception as e:
            raise RuntimeError(
                f"Regen function failed for {job['name']}:{job['offset']}:{job['size']}: {e}"
            ) from e

        implementation = "python"

    # Verify size
    if len(regen_bytes) != job["size"]:
        raise RuntimeError(
            f"Regen size mismatch: expected {job['size']}, got {len(regen_bytes)}"
        )

    dt = time.time() - t0
    sha = hashlib.sha256(regen_bytes).hexdigest()

    result = {
        **job,
        "backend": backend,
        "bytes": regen_bytes,
        "sha256": sha,
        "duration_s": round(dt, 6),
        "core_id": core_id,
        "Se": routing_preview["Se"],
        "implementation": implementation
    }

    # Remove 'fn' from result (not serializable for receipts)
    result.pop("fn", None)
    result.pop("seed_ctx", None)

    return job_idx, result


def create_chunked_jobs(
    name: str,
    total_offset: int,
    total_size: int,
    chunk_size: int,
    regen_fn: Callable,
    seed_ctx: Dict[str, Any]
) -> List[Dict[str, Any]]:
    """
    Split a large regen request into FibPi-sized chunks.

    Args:
        name: Tensor name
        total_offset: Starting offset
        total_size: Total bytes to regenerate
        chunk_size: Chunk size (e.g., 64KB)
        regen_fn: Regen function (seed_ctx, job) -> bytes
        seed_ctx: Context for regen function

    Returns:
        List of job dicts ready for schedule_regen_jobs()
    """
    num_chunks = (total_size + chunk_size - 1) // chunk_size

    jobs = []
    for i in range(num_chunks):
        chunk_offset = total_offset + i * chunk_size
        chunk_end = min(chunk_offset + chunk_size, total_offset + total_size)
        chunk_len = chunk_end - chunk_offset

        jobs.append({
            "name": name,
            "offset": chunk_offset,
            "size": chunk_len,
            "fn": regen_fn,
            "seed_ctx": seed_ctx
        })

    return jobs
