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

REAL substrate driver - handles both .hxz formats:
  1. Current: compressed seed after #====
  2. Future: explicit seed/size/dtype/shape in header

Orchestrates:
    .hxz capsule → parse header → seed_to_bytes → tensor → receipt

NO synthetic codecs, pure orchestration over proven primitives.
"""

from __future__ import annotations

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

import numpy as np


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

@dataclass
class CapsuleHeader:
    path: Path
    entropy: Optional[float]
    sha256_header: Optional[str]
    seed: str
    size_bytes: int
    dtype: str
    shape: Tuple[int, ...]
    format_version: str  # "v1_compressed" or "v2_explicit"
    raw_header_lines: List[str]


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


# ──────────────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────────────

def _now_iso() -> str:
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())


def _parse_shape(s: str) -> Tuple[int, ...]:
    """Parse shape like '1024,1024' or '4096x32000' into tuple."""
    parts = [p.strip() for p in s.replace("x", ",").split(",") if p.strip()]
    if not parts:
        raise ValueError(f"Invalid shape string: {s!r}")
    return tuple(int(p) for p in parts)


def _parse_dtype(s: str) -> str:
    """Normalize dtype string."""
    s = s.strip().lower()
    if s in ("f32", "float32"):
        return "float32"
    if s in ("f16", "float16"):
        return "float16"
    if s in ("bf16", "bfloat16"):
        return "float32"  # Adjust if you use bfloat16
    return s


# ──────────────────────────────────────────────────────────────────────────────
# .hxz parsing (handles both formats)
# ──────────────────────────────────────────────────────────────────────────────

def parse_hxz_capsule(path: Path) -> CapsuleHeader:
    """
    Parse .hxz capsule - handles two formats:

    Format v1 (current - compressed):
        #HXZ_START
        # entropy: 2.662047
        # sha256: 62d243a7...
        #====
        <zlib compressed seed>

    Format v2 (future - explicit):
        #HXZ_START
        # entropy: 2.0
        # sha256: 3e0039fd...
        # seed: 1f647d91f6...
        # size: 4194304
        # dtype: float32
        # shape: 1024,1024

    Returns CapsuleHeader with format_version indicating which was parsed.
    """
    path = Path(path)
    raw_bytes = path.read_bytes()

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

    # Find header end
    header_end = raw_bytes.find(b"#====\n")
    if header_end == -1:
        # No #==== marker - try to parse header until first non-# line
        header_end = len(raw_bytes)
        for i, byte in enumerate(raw_bytes):
            if byte == ord('\n'):
                # Check if next line starts with #
                next_line_start = i + 1
                if next_line_start < len(raw_bytes) and raw_bytes[next_line_start] != ord('#'):
                    header_end = next_line_start
                    break
        payload_bytes = b""
        has_payload = False
    else:
        payload_bytes = raw_bytes[header_end + 6:]  # Skip "#====\n"
        has_payload = True

    header_bytes = raw_bytes[:header_end]
    header_text = header_bytes.decode('utf-8')
    header_lines = [line.strip() for line in header_text.split('\n') if line.strip()]

    # Parse header fields
    entropy: Optional[float] = None
    sha256_val: Optional[str] = None
    seed: Optional[str] = None
    size_bytes: Optional[int] = None
    dtype: str = "float32"
    shape: Optional[Tuple[int, ...]] = None

    for line in header_lines:
        if not line.startswith('#'):
            continue
        body = line.lstrip('#').strip()
        if ':' not in body:
            continue
        key, val = [p.strip() for p in body.split(':', 1)]
        key_l = key.lower()

        if key_l == "entropy":
            try:
                entropy = float(val)
            except ValueError:
                pass
        elif key_l == "sha256":
            sha256_val = val
        elif key_l == "seed":
            seed = val
        elif key_l == "size":
            try:
                size_bytes = int(val)
            except ValueError:
                pass
        elif key_l == "dtype":
            dtype = _parse_dtype(val)
        elif key_l == "shape":
            try:
                shape = _parse_shape(val)
            except Exception:
                pass

    # Determine format and extract seed
    if seed is not None and size_bytes is not None:
        # Format v2: explicit metadata in header
        format_version = "v2_explicit"
        if shape is None:
            # Infer shape from size
            shape = (size_bytes // np.dtype(dtype).itemsize,)
    elif has_payload:
        # Format v1: compressed seed in payload
        format_version = "v1_compressed"
        try:
            decompressed_seed = zlib.decompress(payload_bytes)
            seed = decompressed_seed.hex()
        except zlib.error as e:
            raise ValueError(f"Failed to decompress seed payload: {e}") from e

        # Infer size from seed length (for format v1, size not in header)
        # This is heuristic - adjust based on your actual needs
        if size_bytes is None:
            # Assume we'll be told shape via CLI args
            size_bytes = len(decompressed_seed)
        if shape is None:
            shape = (size_bytes // np.dtype(dtype).itemsize,)
    else:
        raise ValueError(f"Cannot parse .hxz: no explicit seed and no payload: {path}")

    if seed is None:
        raise ValueError(f"No seed found in .hxz capsule: {path}")

    return CapsuleHeader(
        path=path,
        entropy=entropy,
        sha256_header=sha256_val,
        seed=seed,
        size_bytes=size_bytes,
        dtype=dtype,
        shape=shape,
        format_version=format_version,
        raw_header_lines=header_lines,
    )


# ──────────────────────────────────────────────────────────────────────────────
# seed_to_bytes wrapper (your cleaner import)
# ──────────────────────────────────────────────────────────────────────────────

def seed_to_bytes(seed: str, size_bytes: int) -> bytes:
    """
    Call REAL seed_to_bytes from CDR_Steven_Package.

    Tries multiple import paths for flexibility.
    """
    import sys
    from pathlib import Path

    # Try import paths in order (modify sys.path then import)
    search_paths = [
        Path.home() / "CDR_Steven_Package" / "code",
        Path.home() / "helix-cdc",
    ]

    for search_path in search_paths:
        if search_path.exists():
            sys.path.insert(0, str(search_path))
            try:
                from helix_core.helix_regrow import seed_to_bytes as seed_fn
                return seed_fn(seed, size_bytes)
            except ImportError:
                sys.path.pop(0)
                continue

    raise RuntimeError(
        f"Cannot import seed_to_bytes from any of: {search_paths}\n"
        "Check that CDR_Steven_Package/code/helix_core/helix_regrow.py exists"
    )


# ──────────────────────────────────────────────────────────────────────────────
# Main API
# ──────────────────────────────────────────────────────────────────────────────

def regenerate_tensor_from_capsule(
    capsule_path: Path | str,
    *,
    shape_override: Optional[Sequence[int]] = None,
    dtype_override: Optional[str] = None,
) -> RegenResult:
    """
    .hxz → tensor + receipt

    Handles both v1 (compressed) and v2 (explicit) formats.
    Shape/dtype overrides allow CLI to specify when header doesn't have them.
    """
    header = parse_hxz_capsule(Path(capsule_path))

    # Apply overrides
    if dtype_override is not None:
        header.dtype = _parse_dtype(dtype_override)
    if shape_override is not None:
        header.shape = tuple(int(x) for x in shape_override)

    # Calculate required bytes from shape
    np_dtype = np.dtype(header.dtype)
    expected_elements = int(np.prod(header.shape))
    required_bytes = expected_elements * np_dtype.itemsize

    # Use header size_bytes if reasonable, otherwise use calculated
    if header.size_bytes != required_bytes:
        print(f"Warning: header size_bytes={header.size_bytes} != calculated={required_bytes}")
        print(f"         Using calculated size for regeneration")
        header.size_bytes = required_bytes

    # Regenerate bytes
    print(f"Regenerating {header.size_bytes:,} bytes from seed (format={header.format_version})")
    start = time.time()
    regen_bytes = seed_to_bytes(header.seed, header.size_bytes)
    dt = time.time() - start

    if len(regen_bytes) != header.size_bytes:
        raise RuntimeError(
            f"seed_to_bytes returned {len(regen_bytes)} bytes, expected {header.size_bytes}"
        )

    # Hash
    sha256_regen = hashlib.sha256(regen_bytes).hexdigest()

    # Reshape
    flat = np.frombuffer(regen_bytes, dtype=np_dtype)
    if flat.size != expected_elements:
        raise RuntimeError(
            f"Size mismatch: shape={header.shape} → {expected_elements} elements, "
            f"but regen gave {flat.size}"
        )
    tensor = flat.reshape(header.shape)

    # Receipt
    receipt: Dict[str, Any] = {
        "version": "superglyph_cdr_driver_v2_real_substrate",
        "timestamp": _now_iso(),
        "capsule": {
            "path": str(header.path),
            "format": header.format_version,
            "entropy": header.entropy,
            "sha256_header": header.sha256_header,
        },
        "seed": {
            "value_hex": header.seed[:64] + "..." if len(header.seed) > 64 else header.seed,
            "size_bytes": header.size_bytes,
        },
        "tensor": {
            "dtype": header.dtype,
            "shape": list(header.shape),
            "num_elements": int(tensor.size),
            "nbytes": int(tensor.nbytes),
            "sha256": sha256_regen,
        },
        "regen": {
            "backend": "seed_to_bytes (CDR_Steven_Package/helix_regrow.py)",
            "implementation": "real",
            "duration_s": round(dt, 6),
            "throughput_bytes_per_s": round(header.size_bytes / dt, 2) if dt > 0 else None,
        },
    }

    return RegenResult(tensor=tensor, receipt=receipt)


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

def _cli() -> None:
    parser = argparse.ArgumentParser(
        description="Regenerate tensor from .hxz capsule (REAL substrate)"
    )
    parser.add_argument("--capsule", required=True, help="Path to .hxz capsule")
    parser.add_argument("--shape", help="Shape override (e.g., '1024,1024' or '4096x32000')")
    parser.add_argument("--dtype", help="Dtype override (e.g., 'float32')")
    parser.add_argument("--receipt-out", help="Path to write JSON receipt")
    parser.add_argument("--tensor-out", help="Path to write tensor (.npy)")

    args = parser.parse_args()

    shape_override = _parse_shape(args.shape) if args.shape else None
    dtype_override = args.dtype

    result = regenerate_tensor_from_capsule(
        capsule_path=args.capsule,
        shape_override=shape_override,
        dtype_override=dtype_override,
    )

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

    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}")

    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()
