#!/usr/bin/env python3
"""
superglyph_cdr_driver.py
========================

REAL substrate driver - orchestrates .hxz → tensor pipeline using existing primitives.

Pipeline:
    .hxz capsule → parse header (entropy, sha256, seed)
                 → seed_to_bytes() (REAL Helix regrow from CDR_Steven_Package)
                 → reshape to tensor
                 → receipt

NO synthetic codecs. Pure orchestration over proven regen functions.
"""

from __future__ import annotations

import json
import time
import zlib
import hashlib
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

import numpy as np


# ──────────────────────────────────────────────────────────────────────────────
# Dataclasses
# ──────────────────────────────────────────────────────────────────────────────

@dataclass
class CapsuleRef:
    """Parsed .hxz capsule reference"""
    path: Path
    entropy: float
    sha256: str
    seed_hex: str  # The compressed/encoded seed data
    metadata: Dict[str, Any]


@dataclass
class TensorRegenResult:
    tensor: np.ndarray
    receipt: Dict[str, Any]


# ──────────────────────────────────────────────────────────────────────────────
# .hxz Capsule Parser
# ──────────────────────────────────────────────────────────────────────────────

def parse_hxz_capsule(capsule_path: Path | str) -> CapsuleRef:
    """
    Parse a REAL .hxz capsule using the format you showed me:

    #HXZ_START
    # entropy: 2.662047
    # sha256: 62d243a7c667690d15b94c1bdaebc4a4c9636c00cfd58d5a2fa9b77ffe9a7d0a
    #====
    <zlib compressed seed data>

    This is the ACTUAL format from your seedpack.
    """
    capsule_path = Path(capsule_path)

    if not capsule_path.exists():
        raise FileNotFoundError(f"Capsule not found: {capsule_path}")

    raw_bytes = capsule_path.read_bytes()

    # Parse header
    if not raw_bytes.startswith(b"#HXZ_START\n"):
        raise ValueError(f"Invalid .hxz format (missing #HXZ_START): {capsule_path}")

    # Find header end marker
    header_end = raw_bytes.find(b"#====\n")
    if header_end == -1:
        raise ValueError(f"Invalid .hxz format (missing #====): {capsule_path}")

    header_bytes = raw_bytes[:header_end]
    seed_bytes = raw_bytes[header_end + 6:]  # Skip "#====\n"

    # Parse header fields
    header_text = header_bytes.decode('utf-8')
    entropy = None
    sha256_val = None

    for line in header_text.split('\n'):
        line = line.strip()
        if line.startswith('# entropy:'):
            entropy = float(line.split(':', 1)[1].strip())
        elif line.startswith('# sha256:'):
            sha256_val = line.split(':', 1)[1].strip()

    if entropy is None or sha256_val is None:
        raise ValueError(f"Missing entropy or sha256 in header: {capsule_path}")

    # Decompress seed (zlib compressed)
    try:
        decompressed_seed = zlib.decompress(seed_bytes)
        seed_hex = decompressed_seed.hex()
    except zlib.error as e:
        raise ValueError(f"Failed to decompress seed: {e}") from e

    return CapsuleRef(
        path=capsule_path,
        entropy=entropy,
        sha256=sha256_val,
        seed_hex=seed_hex,
        metadata={
            "capsule_size": len(raw_bytes),
            "header_size": header_end + 6,
            "compressed_seed_size": len(seed_bytes),
            "decompressed_seed_size": len(decompressed_seed)
        }
    )


# ──────────────────────────────────────────────────────────────────────────────
# Tensor Regeneration (using REAL seed_to_bytes)
# ──────────────────────────────────────────────────────────────────────────────

def _import_seed_to_bytes():
    """Import the REAL seed_to_bytes from CDR_Steven_Package"""
    try:
        import sys
        sys.path.insert(0, str(Path.home() / "CDR_Steven_Package" / "code"))
        from helix_core.helix_regrow import seed_to_bytes
        return seed_to_bytes
    except ImportError as e:
        raise RuntimeError(
            f"Failed to import seed_to_bytes from CDR_Steven_Package: {e}\n"
            "This driver requires the REAL helix_regrow.py module."
        ) from e


def regenerate_tensor_from_capsule(
    capsule_path: Path | str,
    tensor_shape: Tuple[int, ...],
    dtype: str = "float32",
    *,
    use_real_regen: bool = True,
) -> TensorRegenResult:
    """
    REAL substrate driver entrypoint.

    Pipeline:
      1. parse_hxz_capsule() - Parse .hxz → seed_hex
      2. seed_to_bytes()     - REAL Helix regrow (from CDR_Steven_Package)
      3. Reshape to tensor   - numpy reshape
      4. Receipt             - Full provenance

    Args:
        capsule_path: Path to .hxz capsule
        tensor_shape: Target tensor shape (e.g., (4096, 32000) for output.weight)
        dtype: Numpy dtype (default: "float32")
        use_real_regen: If False, raises error (no fallback to synthetic)

    Returns:
        TensorRegenResult(tensor, receipt)
    """
    if not use_real_regen:
        raise RuntimeError(
            "use_real_regen=False not supported. This driver ONLY uses REAL substrate."
        )

    t0 = time.time()

    # Step 1: Parse capsule
    capsule = parse_hxz_capsule(capsule_path)

    # Step 2: Calculate required bytes
    np_dtype = np.dtype(dtype)
    total_elements = int(np.prod(tensor_shape))
    required_bytes = total_elements * np_dtype.itemsize

    # Step 3: Regenerate bytes using REAL seed_to_bytes
    seed_to_bytes = _import_seed_to_bytes()

    print(f"Regenerating {required_bytes:,} bytes from seed (shape={tensor_shape}, dtype={dtype})")
    regen_start = time.time()

    tensor_bytes = seed_to_bytes(capsule.seed_hex, required_bytes)

    regen_duration = time.time() - regen_start

    # Step 4: Reshape to tensor
    flat_array = np.frombuffer(tensor_bytes, dtype=np_dtype)

    if flat_array.size != total_elements:
        raise RuntimeError(
            f"Size mismatch: regenerated {flat_array.size} elements, "
            f"expected {total_elements} for shape {tensor_shape}"
        )

    tensor = flat_array.reshape(tensor_shape)

    total_duration = time.time() - t0

    # Step 5: Build receipt
    tensor_sha256 = hashlib.sha256(tensor_bytes).hexdigest()

    receipt = {
        "version": "superglyph_cdr_driver_v1_real_substrate",
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "capsule": {
            "path": str(capsule.path),
            "entropy": capsule.entropy,
            "sha256": capsule.sha256,
            "metadata": capsule.metadata,
        },
        "tensor": {
            "shape": list(tensor_shape),
            "dtype": dtype,
            "elements": total_elements,
            "bytes": required_bytes,
            "sha256": tensor_sha256,
        },
        "regen": {
            "backend": "seed_to_bytes (CDR_Steven_Package/helix_regrow.py)",
            "seed_hex_len": len(capsule.seed_hex),
            "duration_s": round(regen_duration, 6),
            "implementation": "real",
        },
        "results": {
            "duration_total_s": round(total_duration, 6),
            "duration_regen_s": round(regen_duration, 6),
            "bytes_per_second": int(required_bytes / regen_duration) if regen_duration > 0 else 0,
        },
    }

    return TensorRegenResult(tensor=tensor, receipt=receipt)


# ──────────────────────────────────────────────────────────────────────────────
# CLI
# ──────────────────────────────────────────────────────────────────────────────

def _cli() -> None:
    import argparse

    parser = argparse.ArgumentParser(
        description="Regenerate tensor from .hxz capsule via REAL substrate"
    )
    parser.add_argument("--capsule", required=True, help="Path to .hxz capsule")
    parser.add_argument("--shape", required=True, help="Tensor shape (e.g., '4096,32000')")
    parser.add_argument("--dtype", default="float32", help="Numpy dtype (default: float32)")
    parser.add_argument("--receipt-out", help="Optional path to write JSON receipt")
    parser.add_argument("--tensor-out", help="Optional path to write tensor (.npy)")

    args = parser.parse_args()

    # Parse shape
    try:
        shape = tuple(int(x) for x in args.shape.split(','))
    except ValueError as e:
        raise SystemExit(f"Invalid shape '{args.shape}': {e}") from e

    # Regenerate
    result = regenerate_tensor_from_capsule(
        capsule_path=args.capsule,
        tensor_shape=shape,
        dtype=args.dtype,
    )

    print()
    print(f"✅ Tensor regenerated from {Path(args.capsule).name}")
    print(f"   Shape: {result.tensor.shape}")
    print(f"   Dtype: {result.tensor.dtype}")
    print(f"   Bytes: {result.tensor.nbytes:,}")
    print(f"   SHA256: {result.receipt['tensor']['sha256'][:16]}...")
    print(f"   Duration: {result.receipt['results']['duration_regen_s']:.3f}s")
    print(f"   Throughput: {result.receipt['results']['bytes_per_second']:,} bytes/sec")
    print()

    # Save receipt
    if args.receipt_out:
        receipt_path = Path(args.receipt_out)
        receipt_path.parent.mkdir(parents=True, exist_ok=True)
        receipt_path.write_text(json.dumps(result.receipt, indent=2))
        print(f"📋 Receipt: {receipt_path}")

    # Save tensor
    if args.tensor_out:
        tensor_path = Path(args.tensor_out)
        tensor_path.parent.mkdir(parents=True, exist_ok=True)
        np.save(tensor_path, result.tensor)
        print(f"💾 Tensor: {tensor_path}")


if __name__ == "__main__":
    _cli()
