#!/usr/bin/env python3
"""
petal_seed.py - Regeneration shim for FlowTorch engine

Enables CDC mode by regenerating full model weights from .hxz seed
via Crystal Vault before inference. Caches by seed hash.
"""
from dataclasses import dataclass
from pathlib import Path
import hashlib
import json
import os
import subprocess
import tempfile
import time

# Env knobs
VAULT = Path(os.environ.get("HELIX_VAULT_PATH",
    "/home/voidstr3m33/echo_labs/helix/proof/helix_proof_v2/vault"))
REGEN = Path(os.environ.get("HELIX_REGEN_SCRIPT",
    "/home/voidstr3m33/echo_labs/helix/proof/packs/mistral-7b-q4km_2025-10-10/interpreter/regrow_from_braid.py"))
CACHE = Path(os.environ.get("HELIX_PETAL_CACHE_DIR", "/tmp/helix_petal_cache")).expanduser()
CACHE.mkdir(parents=True, exist_ok=True)

# Target size for Mistral-7B Q4_K_M
TARGET_SIZE = 4_368_439_584  # 4.1GB


@dataclass
class PetalSeed:
    """Container for .hxz seed bytes"""
    hxz_bytes: bytes
    tag: str = "mistral-7b.q4_k_m"


def _sha(b: bytes) -> str:
    """SHA256 hash of bytes"""
    return hashlib.sha256(b).hexdigest()


def regrow_petal(petal_seed: PetalSeed, petal_delta_h=None):
    """
    Regenerate full model weights from .hxz seed via Crystal Vault.

    Returns: (petal_seed_updated, petal_delta_h_or_meta)

    Where petal_delta_h_or_meta is dict with:
        - "gguf_path": str path to regenerated model
        - "receipt": dict with regen proof
        - "cached": bool whether cache was used
    """
    # Compute seed hash for cache key
    sid = _sha(petal_seed.hxz_bytes)[:16]
    out_path = CACHE / f"{sid}.gguf"
    rcp_path = CACHE / f"{sid}.regen.json"

    # Check cache
    if out_path.exists() and rcp_path.exists():
        receipt = json.loads(rcp_path.read_text())
        meta = {
            "gguf_path": str(out_path),
            "receipt": receipt,
            "cached": True
        }
        return petal_seed, meta

    # Write temp seed file
    with tempfile.NamedTemporaryFile(delete=False, suffix=".hxz", mode='wb') as f:
        f.write(petal_seed.hxz_bytes)
        tmp_seed = f.name

    try:
        # Call regrow script with positional args: hxz_path, output_path, target_size
        # (NOT --seed, --vault, --out flags - script uses positional args)
        start = time.time()

        print(f"🧬 Regenerating from seed {sid}...", flush=True)
        print(f"   Seed: {tmp_seed} ({len(petal_seed.hxz_bytes)} bytes)", flush=True)
        print(f"   Output: {out_path}", flush=True)
        print(f"   Target: {TARGET_SIZE:,} bytes", flush=True)

        # Add HELIX_RSA_LABS path for helix_regrow import
        env = os.environ.copy()
        helix_rsa = "/home/voidstr3m33/echo_labs/helix/labs/HELIX_RSA_LABS"
        env["PYTHONPATH"] = f"{helix_rsa}:{env.get('PYTHONPATH', '')}"

        result = subprocess.run(
            ["python3", str(REGEN), tmp_seed, str(out_path), str(TARGET_SIZE)],
            capture_output=True,
            text=True,
            timeout=1800,  # 30 min timeout (first regen can be slow)
            env=env
        )

        elapsed = time.time() - start

        if result.returncode != 0:
            raise RuntimeError(f"Regen failed: {result.stderr[:500]}")

        print(f"✅ Regeneration complete in {elapsed:.1f}s", flush=True)

        # Create receipt
        receipt = {
            "seed_sha256": _sha(petal_seed.hxz_bytes),
            "seed_bytes": len(petal_seed.hxz_bytes),
            "vault": str(VAULT),
            "out": str(out_path),
            "out_bytes": out_path.stat().st_size if out_path.exists() else 0,
            "target_bytes": TARGET_SIZE,
            "mode": "full_regen",
            "elapsed_sec": elapsed,
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            "regen_script": str(REGEN),
        }

        rcp_path.write_text(json.dumps(receipt, indent=2))

        meta = {
            "gguf_path": str(out_path),
            "receipt": receipt,
            "cached": False
        }

        return petal_seed, meta

    finally:
        # Clean up temp seed
        try:
            os.unlink(tmp_seed)
        except:
            pass


# Expose both for compatibility
__all__ = ["PetalSeed", "regrow_petal"]
