#!/usr/bin/env python3
"""
Full GGUF Regeneration via Streaming

Regenerates a complete GGUF model from:
- DNA seed (5 bytes)
- Entropy seed (4 bytes)
- Centroid atlas (~1KB)

This produces PROCEDURAL weights - statistically similar to the original
model's distribution but NOT exact vendor weight recovery.

Architecture:
    DNA + Entropy + Atlas → Deterministic RNG → Q4_K_M blocks → Full GGUF

Usage as library:
    from helix_cdc.regrow import regenerate_full_gguf, RegenConfig

    config = RegenConfig(
        dna="1f647d91f6",
        entropy="bf690b54",
        atlas_path=Path("artifacts/centroids/mistral_v1/crystal_centroids_v1.npy"),
        header_source=Path("/path/to/reference.gguf"),
        output_path=Path("/path/to/output.gguf"),
    )
    receipt = regenerate_full_gguf(config)

What this DOES:
    - Deterministic regeneration (same seed → same bytes)
    - Q4_K_M block format compatible with llama.cpp
    - Streaming/resumable for 4GB+ models
    - Parallel synthesis (multi-core)

What this does NOT claim:
    - NOT vendor weight recovery
    - NOT lossless reconstruction
    - NOT Shannon-violating compression
"""

import json
import hashlib
import struct
import os
from dataclasses import dataclass, field
from pathlib import Path
from multiprocessing import Pool, cpu_count
from typing import Tuple, Optional, Dict, Any, List
from datetime import datetime, timezone

import numpy as np

# Wave scheduler integration
from helix_cdc.scheduler.wave_scheduler import (
    engine_firing_order,
    get_runtime_schedule,
    get_cpu_count as wave_get_cpu_count,
)
from helix_cdc.scheduler.wave_engine import WaveEngine, get_engine as get_wave_engine

# Q4_K_M block constants
BLOCK_BYTES = 24  # 2 (scale) + 2 (zero) + 16 (codes) + 4 (sign_mask)
VALUES_PER_BLOCK = 32
DEFAULT_SLICE_MB = 64
DEFAULT_HEADER_BYTES = 735520  # Standard Mistral-7B Q4_K_M header


@dataclass
class RegenConfig:
    """Configuration for full GGUF regeneration."""
    dna: str
    entropy: str
    atlas_path: Path
    header_source: Path  # Reference GGUF for header
    output_path: Path
    state_path: Optional[Path] = None
    label: str = "mistral_tensors"
    slice_mb: int = DEFAULT_SLICE_MB
    header_bytes: int = DEFAULT_HEADER_BYTES
    workers: int = field(default_factory=lambda: max(1, wave_get_cpu_count() - 1))
    receipts_dir: Optional[Path] = None
    scheduler_mode: str = "wave"  # "wave" (FibPi3D) or "pool" (legacy raw Pool)
    allow_procedural: bool = False  # Must explicitly enable procedural regen

    def __post_init__(self):
        self.atlas_path = Path(self.atlas_path)
        self.header_source = Path(self.header_source)
        self.output_path = Path(self.output_path)
        if self.state_path:
            self.state_path = Path(self.state_path)
        if self.receipts_dir:
            self.receipts_dir = Path(self.receipts_dir)


def _rng_from_context(dna_hex: str, entropy_hex: str, label: str, slice_index: int) -> np.random.Generator:
    """Create deterministic RNG from seed context."""
    key = f"{dna_hex}|{entropy_hex}|{label}|{slice_index}".encode()
    seed64 = int.from_bytes(hashlib.sha256(key).digest()[:8], "big")
    return np.random.default_rng(seed64)


def _sha256_file(p: Path) -> str:
    """Compute SHA256 of file."""
    return hashlib.sha256(p.read_bytes()).hexdigest()


def _gen_slice(args: Tuple) -> Tuple[int, bytes, dict]:
    """
    Generate a single slice of tensor data.

    Args:
        args: Tuple of (dna_hex, entropy_hex, label, slice_index, size_bytes, atlas_path, scheduler_mode, wave_state_file)
              scheduler_mode: "wave" or "pool"
              wave_state_file: path to shared WaveEngine state file (for cross-process coordination)

    Returns: (slice_index, data_bytes, receipt_dict)
    """
    # Unpack args - now includes scheduler_mode and wave_state_file
    dna_hex, entropy_hex, label, slice_index, size_bytes, atlas_path, scheduler_mode, wave_state_file = args

    # Calculate job weight (normalized to 64MB = 1.0)
    job_weight = size_bytes / (64 * 1024 * 1024)

    # Use WaveEngine in THIS worker process via get_engine().job()
    # This is the key fix: core selection happens AT RUNTIME in the worker,
    # not pre-baked in the planning loop
    core_id = None
    if scheduler_mode == "wave" and wave_state_file:
        try:
            engine = get_wave_engine(state_file=wave_state_file, seed=dna_hex)
            # Use job() context manager - this does acquire() on enter, release() on exit
            with engine.job(weight=job_weight) as assigned_core:
                core_id = assigned_core
                # Set CPU affinity inside the job context
                try:
                    os.sched_setaffinity(0, {core_id})
                except (OSError, AttributeError):
                    pass  # Fall back gracefully

                # Generate slice data inside the wave job context
                data, n_blocks = _generate_slice_payload(
                    dna_hex, entropy_hex, label, slice_index, size_bytes, atlas_path
                )
        except Exception as e:
            # Fallback if WaveEngine fails - generate without affinity
            core_id = None
            data, n_blocks = _generate_slice_payload(
                dna_hex, entropy_hex, label, slice_index, size_bytes, atlas_path
            )
    else:
        # Legacy pool mode - OS decides core, no WaveEngine
        data, n_blocks = _generate_slice_payload(
            dna_hex, entropy_hex, label, slice_index, size_bytes, atlas_path
        )

    receipt = {
        "slice_index": slice_index,
        "size_bytes": len(data),
        "blocks": n_blocks,
        "rng_seed": int.from_bytes(
            hashlib.sha256(f"{dna_hex}|{entropy_hex}|{label}|{slice_index}".encode()).digest()[:8],
            "big"
        ),
        "core_affinity": core_id,  # None if pool mode, actual core_id if wave mode
        "scheduler_mode": scheduler_mode,
    }

    return (slice_index, data, receipt)


def _generate_slice_payload(dna_hex: str, entropy_hex: str, label: str,
                            slice_index: int, size_bytes: int, atlas_path: str) -> Tuple[bytes, int]:
    """
    Generate the actual slice payload (extracted for use inside/outside wave context).

    Returns: (data_bytes, n_blocks)
    """
    rng = _rng_from_context(dna_hex, entropy_hex, label, slice_index)
    atlas = np.load(atlas_path).astype(np.float32)

    n_blocks = (size_bytes + BLOCK_BYTES - 1) // BLOCK_BYTES

    # Sample codes (4-bit, 16 bins) and signs
    codes = rng.integers(0, 16, size=(n_blocks, VALUES_PER_BLOCK), dtype=np.uint8)
    signs = rng.integers(0, 2, size=(n_blocks, VALUES_PER_BLOCK), dtype=np.uint8)

    out = bytearray(n_blocks * BLOCK_BYTES)
    off = 0

    for b in range(n_blocks):
        # Get magnitudes from atlas
        mags = atlas[codes[b].astype(np.int64)]
        vals = mags * np.where(signs[b] == 1, -1.0, 1.0)

        vmin, vmax = float(vals.min()), float(vals.max())
        if abs(vmax - vmin) < 1e-6:
            scale, zero = 1.0, vmin
        else:
            scale, zero = (vmax - vmin) / 15.0, vmin

        # Pack codes as nibbles
        cb = bytearray(16)
        sign_mask = 0
        for i in range(VALUES_PER_BLOCK):
            if signs[b, i]:
                sign_mask |= (1 << i)
            if i % 2 == 0:
                cb[i // 2] = int(codes[b, i]) & 0x0F
            else:
                cb[i // 2] |= (int(codes[b, i]) & 0x0F) << 4

        # Pack block: scale(fp16) + zero(fp16) + codes(16B) + sign_mask(u32)
        s16 = np.float16(scale).tobytes()
        z16 = np.float16(zero).tobytes()
        out[off:off + BLOCK_BYTES] = s16 + z16 + bytes(cb) + struct.pack("<I", sign_mask)
        off += BLOCK_BYTES

    data = bytes(out[:size_bytes])
    return data, n_blocks


def _write_at_offset(path: Path, offset: int, data: bytes):
    """Write data at specific file offset."""
    with open(path, "r+b") as f:
        f.seek(offset)
        f.write(data)
        f.flush()


def _ensure_file_size(path: Path, total_size: int):
    """Pre-allocate file to exact size."""
    path.parent.mkdir(parents=True, exist_ok=True)
    if not path.exists() or path.stat().st_size != total_size:
        with open(path, "wb") as f:
            f.truncate(total_size)


def _copy_header(src: Path, dst: Path, header_bytes: int):
    """Copy GGUF header from source to destination."""
    with open(src, "rb") as fin, open(dst, "r+b") as fout:
        header = fin.read(header_bytes)
        fout.seek(0)
        fout.write(header)
        fout.flush()


def _load_state(state_path: Optional[Path]) -> Optional[dict]:
    """Load resume state if exists."""
    if not state_path or not state_path.exists():
        return None
    try:
        return json.loads(state_path.read_text())
    except Exception:
        return None


def _save_state(state_path: Path, state: dict):
    """Save resume state."""
    state_path.write_text(json.dumps(state, indent=2))


def regenerate_full_gguf(config: RegenConfig) -> Dict[str, Any]:
    """
    Regenerate a full GGUF model from DNA seed.

    Args:
        config: RegenConfig with all parameters

    Returns:
        Receipt dict with regeneration details

    WARNING: This produces PROCEDURAL weights (random noise in Q4_K_M format),
    NOT recovered vendor weights. The output will NOT produce coherent text.
    Use only for structural/format testing.
    """
    start_time = datetime.now(timezone.utc)

    # Safety gate: Require explicit consent for procedural regen
    env_allow = os.environ.get("HELIX_ALLOW_PROCEDURAL_REGEN", "").lower() in ("1", "true", "yes")
    if not (config.allow_procedural or env_allow):
        raise RuntimeError(
            "\n"
            "================================================================\n"
            "PROCEDURAL REGEN GATE: This operation is disabled by default.\n"
            "================================================================\n"
            "\n"
            "This regeneration produces PROCEDURAL NOISE, not recovered weights.\n"
            "The output GGUF will NOT produce coherent text - it's random Q4_K_M\n"
            "blocks that happen to be structurally valid.\n"
            "\n"
            "To enable (if you understand the limitations):\n"
            "  - Set allow_procedural=True in RegenConfig, OR\n"
            "  - Set HELIX_ALLOW_PROCEDURAL_REGEN=1 environment variable, OR\n"
            "  - Use --allow-procedural flag in CLI\n"
            "\n"
            "For actual inference, use vendor GGUF directly.\n"
            "================================================================\n"
        )

    # Validate inputs
    if not config.atlas_path.exists():
        raise FileNotFoundError(f"Atlas not found: {config.atlas_path}")
    if not config.header_source.exists():
        raise FileNotFoundError(f"Header source not found: {config.header_source}")

    # Calculate sizes
    total_size = config.header_source.stat().st_size
    tensor_bytes = total_size - config.header_bytes
    slice_bytes = config.slice_mb * 1024 * 1024
    slices_total = (tensor_bytes + slice_bytes - 1) // slice_bytes

    print(f"=== Full GGUF Regeneration ===")
    print(f"DNA seed: {config.dna}")
    print(f"Entropy: {config.entropy}")
    print(f"Atlas: {config.atlas_path}")
    print(f"Header: {config.header_bytes:,} bytes")
    print(f"Tensors: {tensor_bytes:,} bytes")
    print(f"Slices: {slices_total} x {config.slice_mb}MB")
    print(f"Workers: {config.workers}")
    print()

    # Prepare output file
    _ensure_file_size(config.output_path, total_size)
    _copy_header(config.header_source, config.output_path, config.header_bytes)

    # Load or init state
    state_path = config.state_path or config.output_path.with_suffix('.state.json')
    state = _load_state(state_path) or {
        "done": [],
        "total_slices": slices_total,
        "header_bytes": config.header_bytes,
        "dna": config.dna,
        "entropy": config.entropy,
        "label": config.label,
        "atlas": str(config.atlas_path),
    }
    done = set(state.get("done", []))

    # Wave scheduler setup - state file path for cross-process coordination
    # Workers will use get_engine() with this state file to coordinate
    wave_state_file = None
    firing_order = None

    if config.scheduler_mode == "wave":
        wave_state_file = str(config.output_path.with_suffix('.wave_state.json'))
        # Get firing order for display (workers will use their own engine instances)
        wave_schedule_info = get_runtime_schedule(N=4, seed=config.dna)
        firing_order = engine_firing_order(config.workers, N=4, seed=config.dna)
        print(f"Wave scheduler: firing_order={firing_order}")
        print(f"  Workers will coordinate via: {wave_state_file}")

    # Plan tasks - NO pre-assignment of cores!
    # Each worker will call get_engine().job() at runtime for dynamic core selection
    tasks = []
    for s in range(slices_total):
        if s in done:
            continue
        start = s * slice_bytes
        remaining = tensor_bytes - start
        this_size = min(slice_bytes, remaining)

        # Pass scheduler_mode and wave_state_file to worker
        # Worker will do the actual WaveEngine.job() acquire/release
        tasks.append((
            config.dna, config.entropy, config.label,
            s, this_size, str(config.atlas_path),
            config.scheduler_mode, wave_state_file
        ))

    if not tasks:
        print("All slices complete.")
    else:
        # Create receipts dir if needed
        receipts_dir = config.receipts_dir or config.output_path.parent / "slice_receipts"
        receipts_dir.mkdir(parents=True, exist_ok=True)

        # Generate slices in parallel
        # Workers coordinate via WaveEngine using shared state file (no release() needed here!)
        with Pool(processes=config.workers) as pool:
            for slice_index, data, rec in pool.imap_unordered(_gen_slice, tasks, chunksize=1):
                file_offset = config.header_bytes + slice_index * slice_bytes
                _write_at_offset(config.output_path, file_offset, data)

                # No wave_engine.release() here - workers handle their own acquire/release
                # via the job() context manager

                done.add(slice_index)
                state["done"] = sorted(list(done))
                _save_state(state_path, state)

                # Write slice receipt (includes core_affinity from worker)
                rec_path = receipts_dir / f"slice_{slice_index:05d}.json"
                rec.update({
                    "dna": config.dna,
                    "entropy": config.entropy,
                    "label": config.label,
                })
                rec_path.write_text(json.dumps(rec, indent=2))

                pct = len(done) / slices_total * 100
                core_info = f" core={rec.get('core_affinity')}" if rec.get('core_affinity') is not None else ""
                print(f"[{len(done)}/{slices_total}] Slice {slice_index} ({len(data):,} bytes){core_info} - {pct:.1f}%")

    # Generate wave_cores.json receipt by reading actual slice receipts
    # This reflects what cores were ACTUALLY used at runtime, not pre-assigned
    if config.scheduler_mode == "wave":
        receipts_dir = config.receipts_dir or config.output_path.parent / "slice_receipts"

        # Read back all slice receipts to get actual core assignments
        slice_to_core: Dict[int, int] = {}
        for rec_file in sorted(receipts_dir.glob("slice_*.json")):
            try:
                rec_data = json.loads(rec_file.read_text())
                slice_idx = rec_data.get("slice_index")
                core_aff = rec_data.get("core_affinity")
                if slice_idx is not None and core_aff is not None:
                    slice_to_core[slice_idx] = core_aff
            except Exception:
                pass

        # Load final WaveEngine state from shared state file
        wave_engine_status = None
        if wave_state_file:
            try:
                engine = get_wave_engine(state_file=wave_state_file, seed=config.dna)
                wave_engine_status = engine.status()
            except Exception:
                pass

        # Compute core distribution from actual slice receipts
        all_cores = set(slice_to_core.values())
        core_distribution = {
            str(core_id): sum(1 for c in slice_to_core.values() if c == core_id)
            for core_id in sorted(all_cores)
        }

        wave_cores_receipt = {
            "schema": "wave_cores_receipt_v2",  # v2 = derived from actual runtime
            "timestamp_utc": datetime.now(timezone.utc).isoformat(),
            "dna_seed": config.dna,
            "scheduler_mode": config.scheduler_mode,
            "firing_order": firing_order,
            "slice_assignments": {str(k): v for k, v in slice_to_core.items()},
            "core_distribution": core_distribution,
            "unique_cores_used": len(all_cores),
            "total_slices": len(slice_to_core),
        }
        if wave_engine_status:
            wave_cores_receipt["flywheel"] = wave_engine_status.get("flywheel")
            wave_cores_receipt["cores"] = wave_engine_status.get("cores")

        wave_cores_path = config.output_path.with_suffix('.wave_cores.json')
        wave_cores_path.write_text(json.dumps(wave_cores_receipt, indent=2))
        print(f"Wave cores receipt: {wave_cores_path}")
        print(f"  Unique cores used: {len(all_cores)}")
        print(f"  Distribution: {core_distribution}")

    # Verify GGUF structure
    with open(config.output_path, "rb") as f:
        magic = f.read(4)
        version = struct.unpack("<I", f.read(4))[0]

    end_time = datetime.now(timezone.utc)
    duration = (end_time - start_time).total_seconds()

    # Compute output hash
    output_sha256 = _sha256_file(config.output_path)

    print()
    print(f"=== Regeneration Complete ===")
    print(f"Output: {config.output_path}")
    print(f"Size: {config.output_path.stat().st_size:,} bytes")
    print(f"SHA256: {output_sha256[:16]}...")
    print(f"GGUF magic={magic!r} version={version}")
    print(f"Duration: {duration:.1f}s")

    # Build receipt
    receipt = {
        "schema": "gguf_full_regen_receipt_v1",
        "timestamp_utc": end_time.isoformat(),
        "duration_seconds": round(duration, 2),
        "inputs": {
            "dna_seed": config.dna,
            "entropy_seed": config.entropy,
            "atlas_path": str(config.atlas_path),
            "atlas_sha256": _sha256_file(config.atlas_path),
            "header_source": str(config.header_source),
            "header_bytes": config.header_bytes,
            "label": config.label,
        },
        "output": {
            "path": str(config.output_path),
            "size_bytes": config.output_path.stat().st_size,
            "sha256": output_sha256,
            "gguf_magic": magic.hex(),
            "gguf_version": version,
        },
        "config": {
            "slice_mb": config.slice_mb,
            "workers": config.workers,
            "slices_total": slices_total,
            "scheduler_mode": config.scheduler_mode,
            "firing_order": firing_order if config.scheduler_mode == "wave" else None,
            "wave_state_file": wave_state_file,  # Path to shared WaveEngine state
        },
        "claims": {
            "deterministic": True,
            "vendor_weight_recovery": False,
            "lossless": False,
            "method": "procedural_centroid_atlas",
        },
    }

    # Save receipt next to output
    receipt_path = config.output_path.with_suffix('.regen_receipt.json')
    receipt_path.write_text(json.dumps(receipt, indent=2))
    print(f"Receipt: {receipt_path}")

    return receipt


def get_default_atlas_path() -> Path:
    """Get path to default Mistral centroid atlas."""
    # Package-relative path
    pkg_root = Path(__file__).parent.parent.parent
    return pkg_root / "artifacts" / "centroids" / "mistral_v1" / "crystal_centroids_v1.npy"


# CLI entry point
if __name__ == "__main__":
    import argparse

    ap = argparse.ArgumentParser(
        description="Full GGUF Regeneration from DNA seed (procedural weights)"
    )
    ap.add_argument("--dna", required=True, help="DNA seed hex (e.g., 1f647d91f6)")
    ap.add_argument("--entropy", required=True, help="Entropy seed hex (e.g., bf690b54)")
    ap.add_argument("--atlas", help="Path to centroid atlas .npy (default: package atlas)")
    ap.add_argument("--header-source", required=True, help="Reference GGUF for header")
    ap.add_argument("--output", required=True, help="Output GGUF path")
    ap.add_argument("--slice-mb", type=int, default=DEFAULT_SLICE_MB, help="Slice size in MB")
    ap.add_argument("--workers", type=int, default=max(1, wave_get_cpu_count() - 1), help="Parallel workers")
    ap.add_argument("--label", default="mistral_tensors", help="Tensor label for RNG context")
    ap.add_argument("--scheduler-mode", choices=["wave", "pool"], default="wave",
                    help="Scheduler mode: 'wave' (FibPi3D affinity) or 'pool' (legacy OS-managed)")

    args = ap.parse_args()

    atlas_path = Path(args.atlas) if args.atlas else get_default_atlas_path()

    config = RegenConfig(
        dna=args.dna,
        entropy=args.entropy,
        atlas_path=atlas_path,
        header_source=Path(args.header_source),
        output_path=Path(args.output),
        slice_mb=args.slice_mb,
        workers=args.workers,
        label=args.label,
        scheduler_mode=args.scheduler_mode,
    )

    regenerate_full_gguf(config)
