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

REAL substrate driver.

This is the thin orchestration layer that walks:

    .hxz capsule
      → parse header (entropy, sha256, SEED, size, optional dtype/shape)
      → seed_to_bytes()  [REAL Helix regrow from CDR_Steven_Package]
      → numpy tensor
      → JSON receipt (provenance, timings, hashes)

No synthetic codec, no fake Q4_K_M here. This just drives the **actual**
CDR regen primitive you've already proven (seed_to_bytes + receipts).
"""

from __future__ import annotations

import argparse
import hashlib
import json
import time
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, ...]
    raw_header_lines: List[str]


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


# ──────────────────────────────────────────────────────────────────────────────
# Low-level helpers
# ──────────────────────────────────────────────────────────────────────────────

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


def _parse_shape(s: str) -> Tuple[int, ...]:
    """
    Parse CLI/ header shape like "1024,1024" into (1024, 1024).
    """
    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 to numpy-compatible name.
    """
    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"  # you can adjust if you actually use bfloat16
    return s


# ──────────────────────────────────────────────────────────────────────────────
# 1) .hxz header parsing
# ──────────────────────────────────────────────────────────────────────────────

def parse_hxz_header(path: Path) -> CapsuleHeader:
    """
    Parse the textual header of a .hxz capsule.

    Expected format (what you showed from mistral-7b.hxz):

        #HXZ_START
        # entropy: 2.662047
        # sha256: 62d243a...
        # seed: 1f647d91f6bf690b54...
        # size: 4194304
        # dtype: float32
        # shape: 1024,1024
        <binary payload...>

    We stop reading when we hit the first non-`#` line or blank line.
    """
    path = Path(path)
    header_lines: List[str] = []
    with path.open("rb") as f:
        while True:
            line = f.readline()
            if not line:
                break
            try:
                s = line.decode("utf-8").rstrip("\n")
            except UnicodeDecodeError:
                # binary data -> header ended
                break
            if not s.startswith("#"):
                break
            header_lines.append(s)

    entropy: Optional[float] = None
    sha256_header: 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:
        # Strip leading "# " or "#"
        body = line.lstrip("#").strip()
        if not body:
            continue
        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:
                entropy = None
        elif key_l == "sha256":
            sha256_header = val
        elif key_l == "seed":
            seed = val
        elif key_l == "size":
            try:
                size_bytes = int(val)
            except ValueError:
                size_bytes = None
        elif key_l == "dtype":
            dtype = _parse_dtype(val)
        elif key_l == "shape":
            try:
                shape = _parse_shape(val)
            except Exception:
                shape = None

    # V1 format: derive seed from sha256 + entropy (no explicit seed field)
    if seed is None:
        if sha256_header is None:
            raise RuntimeError(f"No 'sha256' found in header of {path} (needed for v1 seed derivation)")

        # Combine sha256 (DNA seed) + entropy (entropy seed) to create full seed
        # This matches how superglyph_cdr_driver.py does it
        dna_seed = sha256_header[:32]  # First 32 chars of sha256
        if entropy is not None:
            entropy_seed = hashlib.sha256(str(entropy).encode()).hexdigest()[:16]
        else:
            entropy_seed = "0" * 16

        # Combined seed (matches v2 driver pattern)
        seed = hashlib.sha256((dna_seed + entropy_seed).encode()).hexdigest()

    if size_bytes is None:
        raise RuntimeError(f"No 'size' found in header of {path}")

    if shape is None:
        # Fallback: treat as flat vector
        shape = (size_bytes // np.dtype(dtype).itemsize,)

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


# ──────────────────────────────────────────────────────────────────────────────
# 2) Seed → bytes using REAL regen core
# ──────────────────────────────────────────────────────────────────────────────

def seed_to_bytes(seed: str, size_bytes: int) -> bytes:
    """
    Thin wrapper around your REAL CDR regen primitive.

    This calls regrow_helix_real.py which has the proven seed_to_bytes function.
    """
    # Use the proven regrow_helix_real.py implementation
    import sys
    from pathlib import Path

    # Add tools to path
    tools_path = Path.home() / "helix-cdc" / "tools"
    if str(tools_path) not in sys.path:
        sys.path.insert(0, str(tools_path))

    try:
        from regrow_helix_real import seed_to_bytes_streaming
    except ImportError as e:
        raise RuntimeError(
            "Cannot import seed_to_bytes_streaming from regrow_helix_real.py. "
            "Check that ~/helix-cdc/tools/regrow_helix_real.py exists."
        ) from e

    # Use streaming version (memory efficient)
    from io import BytesIO
    buffer = BytesIO()
    seed_to_bytes_streaming(seed, size_bytes, buffer)
    return buffer.getvalue()


# ──────────────────────────────────────────────────────────────────────────────
# 3) Capsule → tensor (np.ndarray) + receipt
# ──────────────────────────────────────────────────────────────────────────────

def regenerate_tensor_from_capsule(
    capsule_path: Path | str,
    *,
    shape_override: Optional[Sequence[int]] = None,
    dtype_override: Optional[str] = None,
) -> RegenResult:
    """
    High-level API:

      .hxz + optional shape/dtype override
        → header
        → seed_to_bytes(size_bytes)
        → numpy tensor
        → receipt

    If shape_override/dtype_override are provided, they override header values.
    """
    header = parse_hxz_header(Path(capsule_path))

    # Apply overrides if provided
    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)

    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, "
            f"expected {header.size_bytes}"
        )

    # Hash of regenerated bytes
    sha256_regen = hashlib.sha256(regen_bytes).hexdigest()

    # Interpret as numpy tensor
    np_dtype = np.dtype(header.dtype)
    flat = np.frombuffer(regen_bytes, dtype=np_dtype)

    expected_elems = int(np.prod(header.shape))
    if flat.size != expected_elems:
        raise RuntimeError(
            f"Shape mismatch: header.shape={header.shape} "
            f"→ {expected_elems} elements, but regen gave {flat.size}"
        )

    tensor = flat.reshape(header.shape)

    # Build receipt
    receipt: Dict[str, Any] = {
        "version": "superglyph_cdr_driver_v1_real_substrate",
        "timestamp": _now_iso(),
        "capsule": {
            "path": str(header.path),
            "entropy": header.entropy,
            "sha256_header": header.sha256_header,
            "header_lines": header.raw_header_lines,
        },
        "seed": {
            "value": header.seed,
            "size_bytes": header.size_bytes,
        },
        "tensor": {
            "dtype": header.dtype,
            "shape": list(header.shape),
            "num_elements": int(tensor.size),
            "nbytes": int(tensor.nbytes),
        },
        "regen": {
            "duration_s": round(dt, 6),
            "throughput_bytes_per_s": round(header.size_bytes / dt, 2)
            if dt > 0
            else None,
            "sha256_regen": sha256_regen,
        },
    }

    return RegenResult(tensor=tensor, receipt=receipt)


# ──────────────────────────────────────────────────────────────────────────────
# 4) CLI entrypoint
# ──────────────────────────────────────────────────────────────────────────────

def _cli() -> None:
    parser = argparse.ArgumentParser(
        description="Regenerate tensor from a Helix .hxz capsule using REAL substrate."
    )
    parser.add_argument(
        "--capsule", required=True, help="Path to .hxz capsule (e.g. mistral-7b.hxz)"
    )
    parser.add_argument(
        "--shape",
        help='Optional shape override, e.g. "1024,1024" or "4096x4096". '
             "If omitted, use shape from .hxz header.",
    )
    parser.add_argument(
        "--dtype",
        help="Optional dtype override (default from header), e.g. float32, f16.",
    )
    parser.add_argument(
        "--receipt-out",
        help="Optional path to write JSON receipt (e.g. receipts/mistral_from_capsule.json)",
    )

    args = parser.parse_args()

    capsule_path = Path(args.capsule)

    shape_override: Optional[Tuple[int, ...]] = None
    if args.shape:
        shape_override = _parse_shape(args.shape)

    dtype_override = args.dtype

    result = regenerate_tensor_from_capsule(
        capsule_path=capsule_path,
        shape_override=shape_override,
        dtype_override=dtype_override,
    )

    t = result.tensor
    print(f"✅ Tensor regenerated from {capsule_path.name}")
    print(f"   shape: {t.shape}")
    print(f"   dtype: {t.dtype}")
    print(f"   nbytes: {t.nbytes}")
    print(f"   duration: {result.receipt['regen']['duration_s']} s")
    print(f"   throughput: {result.receipt['regen']['throughput_bytes_per_s']} bytes/s")

    if args.receipt_out:
        out_path = Path(args.receipt_out)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        out_path.write_text(json.dumps(result.receipt, indent=2))
        print(f"   receipt → {out_path}")


if __name__ == "__main__":
    _cli()
