#!/usr/bin/env python3
"""
Unified Regen Pipeline - WO#FIBPI_BETTER_TENSOR_AUDIT_V1 Phase 1

This module provides a clean, unified interface to the Helix regeneration pipeline.
It wraps existing implementations without reinventing the wheel.

Core functions:
  - inspect_hxz_header(path): Parse HXZ header/manifest ONLY
  - hxz_to_state_dict(path): Full regen pipeline (seed → state_dict)
  - hxz_to_gguf(path, out_path): Regen straight to GGUF (optional stub)

Architecture:
  - Wraps helix_cdc.regrow.hxz_model_loader functions
  - Compatible with existing regrow() and Crystal Vault functions
  - Receipt-driven: All operations return metadata for KAT verification
"""

from pathlib import Path
from typing import Dict, Any, Tuple, Optional
import json
import time

# Import existing implementations (wrapping, not reinventing)
from helix_cdc.regrow.hxz_model_loader import (
    inspect_hxz_header as _inspect_hxz_header_impl,
    load_hxz_to_state_dict as _load_hxz_to_state_dict_impl,
    HXZHeader,
)


def inspect_hxz_header(path: str) -> dict:
    """
    Read HXZ header/manifest ONLY (fast, no full regen).

    Args:
        path: Path to .hxz file

    Returns:
        dict with schema, model_id, tensor registry, created_utc, sha256, etc.

    Example:
        header = inspect_hxz_header("model.hxz")
        print(f"Format: {header['format']}")
        print(f"Schema: {header['schema']}")
        print(f"Tensors: {len(header.get('registry', []))}")
    """
    # Call existing implementation
    hxz_header: HXZHeader = _inspect_hxz_header_impl(path)

    # Convert to dict for clean API
    header_dict = {
        "format": hxz_header.format,
        "schema": hxz_header.schema,
        "entropy": hxz_header.entropy,
        "sha256": hxz_header.sha256,
        "created_utc": hxz_header.created_utc,
        "encoding": hxz_header.encoding,
        "registry": hxz_header.registry,
        "header_bytes": hxz_header.header_bytes,
        "blob_offset": hxz_header.blob_offset,
        "registry_count": len(hxz_header.registry) if hxz_header.registry else 0,
    }

    return header_dict


def hxz_to_state_dict(path: str, dtype: str = "float32") -> Tuple[dict, dict]:
    """
    Full regen pipeline: HXZ seed → state_dict.

    Pipeline steps:
      1. Parse header (inspect_hxz_header)
      2. Walk seed/encoded payload
      3. Call existing regrow() / Crystal Vault functions
      4. Return (state_dict, mapping_info)

    Args:
        path: Path to .hxz file
        dtype: Target dtype ("float32" or "float16")

    Returns:
        Tuple of (state_dict, mapping_info):
          - state_dict: Dict[str, np.ndarray] - Tensor name → numpy array
          - mapping_info: Dict with {"format", "schema", "tensors_loaded", ...}

    Example:
        state_dict, info = hxz_to_state_dict("model.hxz")
        print(f"Loaded {info['tensors_loaded']} tensors")
        for name, arr in state_dict.items():
            print(f"  {name}: {arr.shape} ({arr.dtype})")
    """
    t0 = time.time()

    # Call existing implementation (wraps hxz_model_loader)
    state_dict, metadata = _load_hxz_to_state_dict_impl(path, dtype=dtype)

    regen_time = (time.time() - t0) * 1000  # ms

    # Add timing to metadata
    metadata["regen_time_ms"] = regen_time
    metadata["dtype_requested"] = dtype

    return state_dict, metadata


def hxz_to_gguf(path: str, out_path: str, dtype: str = "float32") -> dict:
    """
    Optional: Regen straight to GGUF format.

    This is a STUB for now - full implementation would require:
      - GGUF writer (llama.cpp format)
      - Tensor layout matching
      - Metadata conversion

    Args:
        path: Path to .hxz input
        out_path: Path to .gguf output
        dtype: Target dtype

    Returns:
        dict with metadata about tensors/bytes written

    Status:
        STUB - Returns error dict. Not needed for Phase 1 KATs.
    """
    return {
        "status": "stub",
        "error": "hxz_to_gguf not implemented yet - use hxz_to_state_dict instead",
        "note": "For GGUF conversion, use: state_dict = hxz_to_state_dict(path), then convert to GGUF manually",
        "hxz_path": path,
        "requested_gguf_path": out_path,
    }


def get_pipeline_status() -> dict:
    """
    Get status of regen pipeline components.

    Returns:
        dict with component availability and versions
    """
    status = {
        "pipeline_version": "v1.0_phase1",
        "components": {
            "inspect_hxz_header": "available",
            "hxz_to_state_dict": "available",
            "hxz_to_gguf": "stub_only",
        },
        "backends": {},
    }

    # Check for Crystal Vault
    try:
        from helix_cdc.seed_vault_codec import shake256_regrow
        status["backends"]["shake256_regrow"] = "available"
    except ImportError:
        status["backends"]["shake256_regrow"] = "unavailable"

    # Check for block_api regrow
    try:
        from helix_cdc.block_api import regrow
        status["backends"]["block_api_regrow"] = "available"
    except ImportError:
        status["backends"]["block_api_regrow"] = "unavailable"

    # Check for DNA seed loader
    try:
        from helix_cdc.regrow.dna_seed_loader import load_dna_seed
        status["backends"]["dna_seed_loader"] = "available"
    except ImportError:
        status["backends"]["dna_seed_loader"] = "unavailable"

    return status


# CLI for testing
if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage:")
        print("  python -m helix_cdc.regen_pipeline <model.hxz>          # Inspect + regen")
        print("  python -m helix_cdc.regen_pipeline --status             # Show pipeline status")
        sys.exit(1)

    if sys.argv[1] == "--status":
        status = get_pipeline_status()
        print(json.dumps(status, indent=2))
        sys.exit(0)

    hxz_path = sys.argv[1]

    print("=" * 70)
    print("  HELIX REGEN PIPELINE - WO#FIBPI Phase 1")
    print("=" * 70)
    print()

    # 1. Inspect header
    print(f"[1/2] Inspecting HXZ header: {hxz_path}")
    header = inspect_hxz_header(hxz_path)
    print(f"  Format: {header['format']}")
    print(f"  Schema: {header['schema']}")
    if header.get('entropy'):
        print(f"  Entropy: {header['entropy']:.6f}")
    if header.get('sha256'):
        print(f"  SHA256: {header['sha256'][:32]}...")
    if header.get('registry_count'):
        print(f"  Registry entries: {header['registry_count']}")
    print()

    # 2. Full regen
    print(f"[2/2] Running full regen pipeline...")
    state_dict, metadata = hxz_to_state_dict(hxz_path)
    print(f"  Loaded {metadata.get('tensors_loaded', len(state_dict))} tensors")
    print(f"  Regen time: {metadata.get('regen_time_ms', 0):.2f}ms")
    print()

    # Show first 5 tensors
    print("Sample tensors:")
    for i, (name, arr) in enumerate(list(state_dict.items())[:5]):
        print(f"  [{i+1}] {name}: {arr.shape} ({arr.dtype})")
    if len(state_dict) > 5:
        print(f"  ... and {len(state_dict) - 5} more")
    print()

    print("=" * 70)
    print("  REGEN COMPLETE")
    print("=" * 70)
