#!/usr/bin/env python3
"""
manifest_regen.py - Plane C: Deterministic Regen from Manifest

Phase 3 Part 6: Verify & Replay regeneration from manifest.

Invariant R-REGEN-001:
    Regen MUST be deterministic and must not re-run routing heuristics.
    It only verifies and executes the manifest codec plan.

This module:
  1. Loads manifest (v2, v3, or v3.1 schema)
  2. For each shard:
     a. Verifies shard hash matches (if stored)
     b. Loads codec_id from compression_decision (v3+) or infers from codec field (v2)
     c. Executes THAT EXACT codec (no re-routing)
     d. Verifies output tensor matches expected hash (if stored)
  3. Assembles GGUF

Schema versions:
  - v2: codec field at shard level
  - v3: compression_decision with codec_id (routing intent) and codec_executed (actual)
  - v3.1: Phase 3.5.2 "Truth in Manifest" - codec_id is ALWAYS actual codec,
          routing_intent stores original routing decision if fallback was used

Supported codecs:
  - cdna_vq_hybrid_v1: k-means VQ with optional HXZO sidecar
  - svd_lowrank: SVD low-rank approximation (TODO)
  - svd_axis_aware: SVD with axis-aware residual (TODO)
  - dct_quant: DCT quantization (TODO)
  - binary_v1: Lossless binary format
  - fp8_block_affine_v1: FP8 block-wise quantization
"""

from __future__ import annotations

import hashlib
import json
import struct
import zlib
from dataclasses import dataclass, asdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np


# =============================================================================
# Phase 3.5.1: Codec Alias Table
# =============================================================================
# Maps Phase 3.5 axis-aware codec IDs to their canonical decoder.
# svd_row_aware and svd_col_aware both use the generic svd_axis_aware decoder.

CODEC_ALIASES = {
    "svd_row_aware": "svd_axis_aware",   # 2A -> generic SVD handler
    "svd_col_aware": "svd_axis_aware",   # 2B -> generic SVD handler
}


@dataclass
class RegenReceipt:
    """Receipt for a regeneration operation."""
    manifest_path: str
    manifest_schema: str
    output_path: str
    output_sha256: str
    tensors_total: int
    tensors_ok: int
    tensors_fail: int
    tensor_details: List[Dict[str, Any]]
    regen_time_s: float
    created_utc: str

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


def now_utc() -> str:
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")


def sha256_file(path: Path) -> str:
    """Compute SHA256 of a file."""
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()


def sha256_bytes(data: bytes) -> str:
    """Compute SHA256 of bytes."""
    return hashlib.sha256(data).hexdigest()


def load_manifest(manifest_path: Path) -> Dict[str, Any]:
    """
    Load manifest from JSON file.

    Supports both v2 and v3 schemas.
    """
    with manifest_path.open("r") as f:
        manifest = json.load(f)

    schema = manifest.get("schema", "unknown")
    # Phase 3.5.2: Added v3.1 schema (codec_id is always actual codec, not routing intent)
    if schema not in ("cdna_hybrid_manifest_v2", "cdna_hybrid_manifest_v3", "cdna_hybrid_manifest_v3.1"):
        raise ValueError(f"Unsupported manifest schema: {schema}")

    return manifest


def verify_shard_integrity(shard: Dict[str, Any], base_dir: Path) -> Tuple[bool, str]:
    """
    Verify shard file exists and hash matches (if stored).

    Args:
        shard: Shard entry from manifest
        base_dir: Base directory for resolving relative paths

    Returns:
        (ok, message) tuple
    """
    path_str = shard.get("path", "")
    if not path_str:
        return False, "No path in shard entry"

    shard_path = Path(path_str)
    if not shard_path.is_absolute():
        shard_path = base_dir / shard_path

    if not shard_path.exists():
        return False, f"Shard file not found: {shard_path}"

    # Check hash if stored
    expected_hash = shard.get("sha256") or shard.get("cdna_sha256")
    if expected_hash:
        actual_hash = sha256_file(shard_path)
        if actual_hash != expected_hash:
            return False, f"Hash mismatch: expected {expected_hash[:16]}..., got {actual_hash[:16]}..."

    return True, "OK"


def _decode_cdna_vq_hybrid(
    cdna_path: Path,
    sidecar_path: Optional[Path] = None,
) -> np.ndarray:
    """
    Decode CDNA VQ hybrid format (k-means + optional sidecar).

    This is the ONLY decoder for codec_id="kmeans_vq" or codec="cdna_vq_hybrid_v1".

    Args:
        cdna_path: Path to .cdna.hxz file
        sidecar_path: Optional path to .hxzo sidecar file

    Returns:
        Decoded tensor as float32 numpy array
    """
    with cdna_path.open("rb") as f:
        magic = f.read(4)
        if magic != b"HXZC":
            raise RuntimeError(f"{cdna_path} invalid magic {magic!r}")
        version = struct.unpack(">H", f.read(2))[0]
        if version != 1:
            raise RuntimeError(f"{cdna_path} unsupported CDNA version {version}")
        header_len = struct.unpack(">H", f.read(2))[0]
        header = json.loads(f.read(header_len).decode("utf-8"))
        codebook = np.frombuffer(f.read(256 * 4), dtype=np.float32)
        compressed = f.read()

    shape = header.get("shape", [])
    if not shape or len(shape) != 2:
        raise RuntimeError(f"{cdna_path} missing/invalid shape in header: {shape}")

    indices = np.frombuffer(zlib.decompress(compressed), dtype=np.uint8).reshape(shape)
    tensor = codebook[indices]

    # Apply sidecar if provided and exists
    if sidecar_path is not None and sidecar_path.exists():
        try:
            from helix_cdc.regrow.hxzo_sidecar import read_outlier_sidecar
            positions, values, meta = read_outlier_sidecar(str(sidecar_path))
            if len(positions) > 0:
                flat = tensor.flatten()
                flat[positions] = values.astype(np.float32)
                tensor = flat.reshape(shape)
        except Exception as e:
            # Sidecar failed - continue without it
            pass

    return tensor.astype(np.float32, copy=False)


def _decode_svd_lowrank(
    shard_path: Path,
    params: Dict[str, Any],
) -> np.ndarray:
    """
    Decode SVD low-rank format.

    TODO: Implement when SVD codec is added.
    """
    raise NotImplementedError("SVD low-rank codec not yet implemented")


def _decode_svd_axis_aware(
    shard_path: Path,
    params: Dict[str, Any],
) -> np.ndarray:
    """
    Decode SVD axis-aware format.

    TODO: Implement when SVD codec is added.
    """
    raise NotImplementedError("SVD axis-aware codec not yet implemented")


def _decode_dct_quant(
    shard_path: Path,
    params: Dict[str, Any],
) -> np.ndarray:
    """
    Decode DCT quantization format.

    TODO: Implement when DCT codec is added.
    """
    raise NotImplementedError("DCT quantization codec not yet implemented")


def _decode_binary_v1(shard_path: Path) -> np.ndarray:
    """Decode lossless binary_v1 format."""
    from helix_cdc.regrow.hxz_model_loader import load_hxz_to_state_dict
    state_dict, _meta = load_hxz_to_state_dict(str(shard_path))
    if len(state_dict) != 1:
        raise RuntimeError(f"{shard_path} expected 1 tensor, found {len(state_dict)}")
    _name, tensor = next(iter(state_dict.items()))
    return np.asarray(tensor)


def decode_shard(
    shard: Dict[str, Any],
    manifest_dir: Path,
    cdna_dir: Optional[Path] = None,
    sidecar_dir: Optional[Path] = None,
) -> Tuple[np.ndarray, str]:
    """
    Decode a single shard using the codec specified in the manifest.

    Invariant R-REGEN-001: Uses ONLY the codec from the manifest.
    NO re-routing based on tensor content.

    Args:
        shard: Shard entry from manifest
        manifest_dir: Directory containing manifest (for relative paths)
        cdna_dir: Override CDNA directory
        sidecar_dir: Override sidecar directory

    Returns:
        (tensor, codec_used) tuple
    """
    tensor_name = shard.get("tensor_name", "unknown")

    # Determine codec from v3 compression_decision or v2 codec field
    compression_decision = shard.get("compression_decision", {})
    codec_id = compression_decision.get("codec_id")
    codec_params = compression_decision.get("params", {})

    # Fallback for v2 manifests
    if not codec_id:
        codec = shard.get("codec", "unknown")
        if codec == "cdna_vq_hybrid_v1":
            codec_id = "kmeans_vq"
        elif codec == "binary_v1":
            codec_id = "binary_v1"
        elif codec == "fp8_block_affine_v1":
            codec_id = "fp8_block_affine"
        else:
            # Default to k-means if unknown
            codec_id = "kmeans_vq"

    # Phase 3.5.2: Legacy support for v3.0 manifests
    # In v3.0, codec_id was routing intent and codec_executed was actual codec.
    # In v3.1+, codec_id is always the actual codec (routing_intent stores the original).
    # Check codec_executed for backward compatibility with v3.0 manifests.
    if compression_decision.get("codec_executed"):
        codec_id = compression_decision["codec_executed"]

    # Resolve paths
    shard_path_str = shard.get("path", "")
    shard_path = Path(shard_path_str)
    if not shard_path.is_absolute():
        if cdna_dir:
            shard_path = cdna_dir / shard_path.name
        else:
            shard_path = manifest_dir / shard_path_str

    sidecar_path = None
    sidecar_path_str = shard.get("outlier_sidecar_path")
    if sidecar_path_str:
        sidecar_path = Path(sidecar_path_str)
        if not sidecar_path.is_absolute():
            if sidecar_dir:
                sidecar_path = sidecar_dir / sidecar_path.name
            else:
                sidecar_path = manifest_dir / sidecar_path_str

    # Phase 3.5.1: Apply codec alias (svd_row_aware, svd_col_aware -> svd_axis_aware)
    canonical_codec = CODEC_ALIASES.get(codec_id, codec_id)

    # Decode using specified codec (NO RE-ROUTING)
    # Phase 3.5.1: Use canonical_codec for dispatch, but track original codec_id
    if canonical_codec == "kmeans_vq":
        tensor = _decode_cdna_vq_hybrid(shard_path, sidecar_path)
        codec_used = "kmeans_vq"
    elif canonical_codec == "svd_lowrank":
        tensor = _decode_svd_lowrank(shard_path, codec_params)
        codec_used = "svd_lowrank"
    elif canonical_codec == "svd_axis_aware":
        # Phase 3.5.1: This handles svd_axis_aware, svd_row_aware, and svd_col_aware
        # The axis info is in codec_params (from routing decision)
        tensor = _decode_svd_axis_aware(shard_path, codec_params)
        codec_used = codec_id  # Keep original codec_id (2A/2B) for traceability
    elif canonical_codec == "dct_quant":
        tensor = _decode_dct_quant(shard_path, codec_params)
        codec_used = "dct_quant"
    elif canonical_codec == "binary_v1":
        tensor = _decode_binary_v1(shard_path)
        codec_used = "binary_v1"
    else:
        # Unknown codec - try default k-means
        tensor = _decode_cdna_vq_hybrid(shard_path, sidecar_path)
        codec_used = f"kmeans_vq (fallback from {codec_id})"

    return tensor, codec_used


def regen_from_manifest(
    manifest_path: Path,
    output_gguf: Path,
    reference_gguf: Path,
    cdna_dir: Optional[Path] = None,
    sidecar_dir: Optional[Path] = None,
    verify_integrity: bool = True,
) -> RegenReceipt:
    """
    Regenerate GGUF from manifest using stored codec decisions.

    Invariant R-REGEN-001:
        Regen MUST be deterministic and must not re-run routing heuristics;
        it only verifies and executes the manifest codec plan.

    Args:
        manifest_path: Path to v2 or v3 manifest JSON
        output_gguf: Path to output GGUF file
        reference_gguf: Reference GGUF for shapes and metadata
        cdna_dir: Override CDNA shard directory (uses manifest cdna_dir if None)
        sidecar_dir: Override sidecar directory (uses manifest sidecar_dir if None)
        verify_integrity: Verify shard hashes before decoding

    Returns:
        RegenReceipt with details of regeneration
    """
    import time

    try:
        from gguf import GGUFReader, GGUFWriter
    except ImportError:
        raise ImportError("gguf package required. Install with: pip install gguf")

    t0 = time.time()

    # Load manifest
    manifest = load_manifest(manifest_path)
    schema = manifest.get("schema", "unknown")
    manifest_dir = manifest_path.parent

    # Use manifest dirs if not overridden
    if cdna_dir is None:
        cdna_dir_str = manifest.get("cdna_dir")
        if cdna_dir_str:
            cdna_dir = Path(cdna_dir_str)
            if not cdna_dir.is_absolute():
                cdna_dir = manifest_dir / cdna_dir_str

    if sidecar_dir is None:
        sidecar_dir_str = manifest.get("sidecar_dir")
        if sidecar_dir_str:
            sidecar_dir = Path(sidecar_dir_str)
            if not sidecar_dir.is_absolute():
                sidecar_dir = manifest_dir / sidecar_dir_str

    shards = manifest.get("shards", [])

    # Load reference GGUF for shapes and 1D tensors (norms)
    ref_reader = GGUFReader(str(reference_gguf))
    ref_tensors = {}
    ref_shapes = {}
    for t in ref_reader.tensors:
        arr = np.asarray(getattr(t, "data", t.tensor) if hasattr(t, "data") else t.tensor)
        ref_tensors[t.name] = arr
        ref_shapes[t.name] = arr.shape

    # Build tensor dict from shards
    tensors_ok = 0
    tensors_fail = 0
    tensor_details = []
    decoded_tensors = {}

    for shard in shards:
        tensor_name = shard.get("tensor_name", "unknown")

        detail = {"tensor_name": tensor_name, "status": "pending"}

        # Verify integrity if requested
        if verify_integrity:
            ok, msg = verify_shard_integrity(shard, cdna_dir or manifest_dir)
            if not ok:
                detail["status"] = "FAIL"
                detail["error"] = msg
                tensor_details.append(detail)
                tensors_fail += 1
                continue

        # Decode shard
        try:
            tensor, codec_used = decode_shard(
                shard, manifest_dir, cdna_dir, sidecar_dir
            )
            decoded_tensors[tensor_name] = tensor
            detail["status"] = "OK"
            detail["codec_used"] = codec_used
            detail["shape"] = list(tensor.shape)
            tensors_ok += 1
        except Exception as e:
            detail["status"] = "FAIL"
            detail["error"] = str(e)
            tensors_fail += 1

        tensor_details.append(detail)

    # Add 1D tensors (norms) from reference if not in shards
    shard_names = {s.get("tensor_name") for s in shards}
    for name, arr in ref_tensors.items():
        if name not in shard_names:
            # Copy from reference (1D tensors, biases)
            decoded_tensors[name] = arr.astype(np.float32, copy=False)
            tensor_details.append({
                "tensor_name": name,
                "status": "OK",
                "codec_used": "reference_copy",
                "shape": list(arr.shape),
            })
            tensors_ok += 1

    # Write output GGUF
    output_gguf.parent.mkdir(parents=True, exist_ok=True)

    # Copy metadata from reference
    from helix_cdc.regrow.regen_from_helix64_manifest import copy_tokenizer_and_metadata

    with open(output_gguf, "wb") as f:
        writer = GGUFWriter(f, ref_reader.architecture)

        # Copy metadata
        copy_tokenizer_and_metadata(ref_reader, writer)

        # Write tensors
        for name in ref_tensors.keys():
            if name in decoded_tensors:
                tensor = decoded_tensors[name]
                # Ensure shape matches reference
                ref_shape = ref_shapes[name]
                if tensor.shape != ref_shape:
                    # Try reshape
                    if tensor.size == np.prod(ref_shape):
                        tensor = tensor.reshape(ref_shape)
                    else:
                        raise ValueError(
                            f"Shape mismatch for {name}: got {tensor.shape}, expected {ref_shape}"
                        )
                writer.add_tensor(name, tensor)

        writer.write_header_to_file()
        writer.write_kv_data_to_file()
        writer.write_tensors_to_file()

    regen_time = time.time() - t0

    receipt = RegenReceipt(
        manifest_path=str(manifest_path),
        manifest_schema=schema,
        output_path=str(output_gguf),
        output_sha256=sha256_file(output_gguf),
        tensors_total=len(ref_tensors),
        tensors_ok=tensors_ok,
        tensors_fail=tensors_fail,
        tensor_details=tensor_details,
        regen_time_s=round(regen_time, 2),
        created_utc=now_utc(),
    )

    return receipt


def main():
    """CLI entry point for manifest regen."""
    import argparse

    parser = argparse.ArgumentParser(
        description="Regenerate GGUF from manifest (Plane C: Verify & Replay)"
    )
    parser.add_argument("--manifest", required=True, type=Path, help="Path to v2/v3 manifest JSON")
    parser.add_argument("--reference", required=True, type=Path, help="Reference GGUF for shapes/metadata")
    parser.add_argument("--output", required=True, type=Path, help="Output GGUF path")
    parser.add_argument("--cdna-dir", type=Path, help="Override CDNA shard directory")
    parser.add_argument("--sidecar-dir", type=Path, help="Override sidecar directory")
    parser.add_argument("--no-verify", action="store_true", help="Skip shard integrity verification")
    parser.add_argument("--receipt", type=Path, help="Write receipt JSON to this path")
    args = parser.parse_args()

    if not args.manifest.exists():
        print(f"ERROR: Manifest not found: {args.manifest}")
        return 1

    if not args.reference.exists():
        print(f"ERROR: Reference GGUF not found: {args.reference}")
        return 1

    print("=" * 60)
    print("PLANE C: MANIFEST REGEN (Verify & Replay)")
    print("=" * 60)
    print(f"Manifest:   {args.manifest}")
    print(f"Reference:  {args.reference}")
    print(f"Output:     {args.output}")
    print(f"Verify:     {'YES' if not args.no_verify else 'NO'}")
    print()

    receipt = regen_from_manifest(
        manifest_path=args.manifest,
        output_gguf=args.output,
        reference_gguf=args.reference,
        cdna_dir=args.cdna_dir,
        sidecar_dir=args.sidecar_dir,
        verify_integrity=not args.no_verify,
    )

    print()
    print("=" * 60)
    print("REGEN COMPLETE")
    print("=" * 60)
    print(f"Output:       {receipt.output_path}")
    print(f"SHA256:       {receipt.output_sha256[:32]}...")
    print(f"Tensors OK:   {receipt.tensors_ok}/{receipt.tensors_total}")
    print(f"Tensors FAIL: {receipt.tensors_fail}")
    print(f"Time:         {receipt.regen_time_s}s")

    if args.receipt:
        args.receipt.parent.mkdir(parents=True, exist_ok=True)
        args.receipt.write_text(json.dumps(receipt.to_dict(), indent=2))
        print(f"\nReceipt:      {args.receipt}")

    return 0 if receipt.tensors_fail == 0 else 1


if __name__ == "__main__":
    import sys
    sys.exit(main())
