#!/usr/bin/env python3
"""
helix_cdc/regrow/regen_from_helix64_manifest.py

WO-REGEN-HELIX64-MANIFEST-V1:
Regenerate a GGUF model from a Helix64 v2 seed manifest.

This proves the full regeneration loop:
  v2 manifest → decode seeds → assemble GGUF → functionally identical to vendor

The manifest contains:
  - Per-tensor: name, shape, dtype, sha256, helix64_seed
  - Model metadata (architecture, etc.)

NOTE: Tokenizer is NOT in the manifest yet (future WO).
      For now, we copy tokenizer from a reference GGUF (vendor or clone).

Usage:
  PYTHONPATH=. python3 helix_cdc/regrow/regen_from_helix64_manifest.py \
    --manifest seeds/mistral_vendor_clone_cdna_helix64_manifest_v2.json \
    --tokenizer-source artifacts/mistral_vendor_clone_cdna_v3.gguf \
    --out artifacts/mistral_from_seed_cdna_v1.gguf \
    --receipt receipts/kat_seed_regen_build.json
"""

import argparse
import hashlib
import json
import pathlib
import sys
import time
from typing import Any, Dict, List

import numpy as np

# Add repo root to path
ROOT = pathlib.Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT))

try:
    from gguf import GGUFReader, GGUFWriter, GGMLQuantizationType
except ImportError:
    print("[FATAL] gguf package not installed. Run: pip install gguf")
    sys.exit(1)

from tools.helix64 import decode as helix64_decode


# Map dtype strings to GGMLQuantizationType
DTYPE_MAP = {
    "F32": GGMLQuantizationType.F32,
    "F16": GGMLQuantizationType.F16,
    "Q4_0": GGMLQuantizationType.Q4_0,
    "Q4_1": GGMLQuantizationType.Q4_1,
    "Q5_0": GGMLQuantizationType.Q5_0,
    "Q5_1": GGMLQuantizationType.Q5_1,
    "Q8_0": GGMLQuantizationType.Q8_0,
    "Q8_1": GGMLQuantizationType.Q8_1,
    "Q2_K": GGMLQuantizationType.Q2_K,
    "Q3_K": GGMLQuantizationType.Q3_K,
    "Q4_K": GGMLQuantizationType.Q4_K,
    "Q5_K": GGMLQuantizationType.Q5_K,
    "Q6_K": GGMLQuantizationType.Q6_K,
    "Q8_K": GGMLQuantizationType.Q8_K,
    "IQ2_XXS": GGMLQuantizationType.IQ2_XXS,
    "IQ2_XS": GGMLQuantizationType.IQ2_XS,
    "IQ3_XXS": GGMLQuantizationType.IQ3_XXS,
    "IQ1_S": GGMLQuantizationType.IQ1_S,
    "IQ4_NL": GGMLQuantizationType.IQ4_NL,
    "IQ3_S": GGMLQuantizationType.IQ3_S,
    "IQ2_S": GGMLQuantizationType.IQ2_S,
    "IQ4_XS": GGMLQuantizationType.IQ4_XS,
}


def extract_string_array(field) -> List[str]:
    """Extract string array from GGUF field."""
    tokens = []
    idx = 5
    while idx + 1 < len(field.parts):
        str_len_part = field.parts[idx]
        idx += 1
        str_bytes_part = field.parts[idx]
        idx += 1
        if hasattr(str_bytes_part, 'tobytes'):
            tok = str_bytes_part.tobytes().decode('utf-8', errors='replace')
        else:
            tok = bytes(str_bytes_part).decode('utf-8', errors='replace')
        tokens.append(tok)
    return tokens


def extract_numeric_array(field, dtype=float) -> List:
    """Extract numeric array from GGUF field."""
    values = []
    for i in range(5, len(field.parts)):
        val = dtype(field.parts[i][0])
        values.append(val)
    return values


def extract_field_value(field, default=None):
    """Extract scalar value from GGUF field."""
    try:
        if len(field.parts) >= 1:
            val = field.parts[-1]
            if hasattr(val, 'shape'):
                if val.shape == () or (len(val.shape) == 1 and val.shape[0] == 1):
                    return val.item() if hasattr(val, 'item') else val
            if isinstance(val, (int, float, str, bool)):
                return val
            if hasattr(val, 'tobytes'):
                return val.tobytes().decode('utf-8', errors='replace')
    except Exception:
        pass
    return default


def copy_tokenizer_and_metadata(reader: GGUFReader, writer: GGUFWriter) -> Dict[str, Any]:
    """Copy tokenizer and model metadata from reference GGUF to writer."""
    metadata = {}

    # Architecture - already set in GGUFWriter constructor
    if 'general.architecture' in reader.fields:
        arch = extract_field_value(reader.fields['general.architecture'])
        if arch:
            # Architecture is set via constructor, not add_architecture
            metadata['general.architecture'] = arch

    # Name
    if 'general.name' in reader.fields:
        name = extract_field_value(reader.fields['general.name'])
        if name:
            writer.add_name(name)
            metadata['general.name'] = name

    # LLaMA-specific fields
    llama_fields = [
        ('llama.context_length', 'add_context_length'),
        ('llama.embedding_length', 'add_embedding_length'),
        ('llama.block_count', 'add_block_count'),
        ('llama.feed_forward_length', 'add_feed_forward_length'),
        ('llama.attention.head_count', 'add_head_count'),
        ('llama.attention.head_count_kv', 'add_head_count_kv'),
        ('llama.rope.freq_base', 'add_rope_freq_base'),
        ('llama.attention.layer_norm_rms_epsilon', 'add_layer_norm_rms_eps'),
    ]

    for field_name, writer_method in llama_fields:
        if field_name in reader.fields:
            value = extract_field_value(reader.fields[field_name])
            if value is not None:
                try:
                    getattr(writer, writer_method)(value)
                    metadata[field_name] = value
                except Exception:
                    pass

    # Tokenizer vocabulary
    if 'tokenizer.ggml.tokens' in reader.fields:
        tokens_field = reader.fields['tokenizer.ggml.tokens']
        tokens = extract_string_array(tokens_field)
        writer.add_token_list(tokens)
        metadata['tokenizer.ggml.tokens_count'] = len(tokens)

        # Scores
        if 'tokenizer.ggml.scores' in reader.fields:
            scores_field = reader.fields['tokenizer.ggml.scores']
            scores = extract_numeric_array(scores_field, dtype=float)
            writer.add_token_scores(scores)

        # Token types
        if 'tokenizer.ggml.token_type' in reader.fields:
            types_field = reader.fields['tokenizer.ggml.token_type']
            toktypes = extract_numeric_array(types_field, dtype=int)
            writer.add_token_types(toktypes)

    # Tokenizer model
    if 'tokenizer.ggml.model' in reader.fields:
        model = extract_field_value(reader.fields['tokenizer.ggml.model'])
        if model:
            writer.add_tokenizer_model(model)
            metadata['tokenizer.ggml.model'] = model

    # Tokenizer pre (string)
    if 'tokenizer.ggml.pre' in reader.fields:
        pre = extract_field_value(reader.fields['tokenizer.ggml.pre'])
        if pre:
            writer.add_string('tokenizer.ggml.pre', pre)
            metadata['tokenizer.ggml.pre'] = pre

    # Special token IDs
    for tok_name in ['bos', 'eos', 'pad', 'unk', 'sep', 'unknown']:
        field_name = f'tokenizer.ggml.{tok_name}_token_id'
        if field_name in reader.fields:
            value = extract_field_value(reader.fields[field_name])
            if value is not None:
                try:
                    add_method = getattr(writer, f'add_{tok_name}_token_id', None)
                    if add_method:
                        add_method(int(value))
                    else:
                        # fallback: write as uint32 meta if writer lacks helper
                        writer.add_uint32(field_name, int(value))
                    metadata[field_name] = value
                except Exception:
                    pass

    # Chat template
    if 'tokenizer.chat_template' in reader.fields:
        template = extract_field_value(reader.fields['tokenizer.chat_template'])
        if template:
            writer.add_chat_template(template)
            metadata['tokenizer.chat_template'] = 'present'

    # BOS/EOS flags
    for flag in ['tokenizer.ggml.add_bos_token', 'tokenizer.ggml.add_eos_token', 'tokenizer.ggml.add_sep_token', 'tokenizer.ggml.add_space_prefix']:
        if flag in reader.fields:
            value = extract_field_value(reader.fields[flag])
            if value is not None:
                if 'add_bos' in flag:
                    writer.add_add_bos_token(bool(value))
                elif 'add_eos' in flag:
                    writer.add_add_eos_token(bool(value))
                elif 'add_sep' in flag:
                    # gguf writer lacks explicit sep support; store as bool meta
                    writer.add_bool(flag, bool(value))
                elif 'add_space_prefix' in flag:
                    writer.add_bool(flag, bool(value))
                metadata[flag] = value

    # File type and quantization version
    if 'general.file_type' in reader.fields:
        value = extract_field_value(reader.fields['general.file_type'])
        if value is not None:
            writer.add_file_type(int(value))
            metadata['general.file_type'] = value

    if 'general.quantization_version' in reader.fields:
        value = extract_field_value(reader.fields['general.quantization_version'])
        if value is not None:
            writer.add_quantization_version(int(value))
            metadata['general.quantization_version'] = value

    # Best-effort copy of any remaining scalar/string metadata not handled above
    # Preserve key types where possible; enforce u32 for known keys
    force_u32 = {
        "llama.rope.dimension_count",
        "llama.vocab_size",
        "general.file_type",
        "general.quantization_version",
        "tokenizer.ggml.bos_token_id",
        "tokenizer.ggml.eos_token_id",
        "tokenizer.ggml.pad_token_id",
        "tokenizer.ggml.unk_token_id",
        "tokenizer.ggml.unknown_token_id",
    }

    handled_keys = set(metadata.keys()) | {
        'general.architecture', 'GGUF.version', 'GGUF.tensor_count', 'GGUF.kv_count'
    }
    # mark tokenizer keys already handled
    handled_keys.update(k for k in reader.fields.keys() if k.startswith("tokenizer.ggml.") or k == "tokenizer.chat_template")
    for key, field in reader.fields.items():
        if key in handled_keys:
            continue
        # Never let the fallback touch tokenizer keys; they include arrays and must be copied explicitly.
        if key.startswith("tokenizer.ggml.") or key == "tokenizer.chat_template":
            continue
        # Skip internal GGUF meta
        if key in ("GGUF.version", "GGUF.tensor_count", "GGUF.kv_count"):
            continue
        val = extract_field_value(field)
        if val is None:
            continue
        try:
            if isinstance(val, bool):
                writer.add_bool(key, bool(val))
            elif isinstance(val, int):
                if key in force_u32:
                    writer.add_uint32(key, int(val))
                else:
                    writer.add_uint64(key, int(val))
            elif isinstance(val, float):
                writer.add_float64(key, float(val))
            elif isinstance(val, str):
                writer.add_string(key, val)
            else:
                continue
            metadata[key] = val
        except Exception:
            # If the writer lacks a matching add_* helper, just skip silently.
            continue

    return metadata


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Regenerate GGUF from Helix64 v2 seed manifest."
    )
    parser.add_argument(
        "--manifest",
        default="seeds/mistral_vendor_clone_cdna_helix64_manifest_v2.json",
        help="Path to v2 seed manifest",
    )
    parser.add_argument(
        "--tokenizer-source",
        default="artifacts/mistral_vendor_clone_cdna_v3.gguf",
        help="GGUF to copy tokenizer/metadata from (NOT used for weights!)",
    )
    parser.add_argument(
        "--out",
        default="artifacts/mistral_from_seed_cdna_v1.gguf",
        help="Path to output GGUF",
    )
    parser.add_argument(
        "--receipt",
        default="receipts/kat_seed_regen_build.json",
        help="Path to output JSON receipt",
    )
    parser.add_argument(
        "--max-tensors",
        type=int,
        default=0,
        help="Max tensors to process (0 = all, for testing)",
    )
    args = parser.parse_args()

    manifest_path = ROOT / args.manifest
    tokenizer_path = pathlib.Path(args.tokenizer_source)
    out_path = ROOT / args.out
    receipt_path = ROOT / args.receipt

    if not manifest_path.exists():
        raise SystemExit(f"Missing manifest: {manifest_path}")
    if not tokenizer_path.exists():
        raise SystemExit(f"Missing tokenizer source: {tokenizer_path}")

    out_path.parent.mkdir(parents=True, exist_ok=True)
    receipt_path.parent.mkdir(parents=True, exist_ok=True)

    t0 = time.time()

    # Load manifest
    print("[SEED-REGEN] Loading v2 seed manifest...")
    print(f"  Path: {manifest_path}")
    print("  (This may take a moment for ~5.5GB file...)")

    with manifest_path.open("r") as f:
        manifest = json.load(f)

    tensor_index = manifest["tensor_index"]
    print(f"  Loaded {len(tensor_index)} tensor entries")

    # Load tokenizer source
    print(f"\n[SEED-REGEN] Loading tokenizer from: {tokenizer_path}")
    tokenizer_reader = GGUFReader(str(tokenizer_path))
    print(f"  Tokenizer fields: {len(tokenizer_reader.fields)}")

    # Create writer
    print(f"\n[SEED-REGEN] Creating output GGUF: {out_path}")
    writer = GGUFWriter(str(out_path), "llama")

    # Copy tokenizer and metadata
    print("\n[SEED-REGEN] Copying tokenizer and metadata...")
    metadata = copy_tokenizer_and_metadata(tokenizer_reader, writer)
    print(f"  Copied {len(metadata)} metadata fields")

    # Process tensors from manifest
    if args.max_tensors > 0:
        tensor_index = tensor_index[:args.max_tensors]

    print(f"\n[SEED-REGEN] Decoding {len(tensor_index)} tensors from seeds...")

    tensor_stats: List[Dict[str, Any]] = []
    total_bytes = 0
    total_verified = 0

    for i, entry in enumerate(tensor_index):
        name = entry["name"]
        shape = tuple(entry["shape"])
        dtype_str = entry["dtype"]
        expected_sha256 = entry["sha256"]
        helix64_seed = entry["helix64_seed"]

        # Decode seed to bytes
        decoded_bytes = helix64_decode(helix64_seed)
        actual_sha256 = hashlib.sha256(decoded_bytes).hexdigest()

        # Verify SHA256
        sha256_match = (actual_sha256 == expected_sha256)
        if not sha256_match:
            print(f"  ❌ SHA256 MISMATCH for {name}!")
            print(f"     Expected: {expected_sha256}")
            print(f"     Actual:   {actual_sha256}")

        # Reconstruct numpy array with correct physical shape
        # Get quantization type
        if dtype_str not in DTYPE_MAP:
            raise SystemExit(f"Unknown dtype: {dtype_str} for tensor {name}")
        raw_dtype = DTYPE_MAP[dtype_str]

        n_bytes = len(decoded_bytes)

        # Compute physical shape from logical shape
        # - F32/F16 tensors: physical shape matches logical shape
        # - Quantized 2D tensors: physical shape is (logical[1], bytes_per_row)
        # - Quantized 1D tensors: physical shape is (n_bytes,) as uint8
        if dtype_str in ("F32", "F16"):
            # Float tensors - interpret as native dtype
            if dtype_str == "F32":
                tensor_data = np.frombuffer(decoded_bytes, dtype=np.float32).reshape(shape)
            else:  # F16
                tensor_data = np.frombuffer(decoded_bytes, dtype=np.float16).reshape(shape)
            # For float tensors, don't pass raw_dtype - let writer infer from tensor dtype
            writer.add_tensor(name=name, tensor=tensor_data)
        else:
            # Quantized tensors - stored as uint8 with physical byte shape
            if len(shape) == 2:
                # 2D quantized: physical shape is (logical[1], bytes_per_row)
                rows = shape[1]  # Note: GGUF uses (rows, cols) = (shape[1], shape[0])
                bytes_per_row = n_bytes // rows
                physical_shape = (rows, bytes_per_row)
            else:
                # 1D quantized (rare): just use flat shape
                physical_shape = (n_bytes,)

            tensor_data = np.frombuffer(decoded_bytes, dtype=np.uint8).reshape(physical_shape)
            writer.add_tensor(name=name, tensor=tensor_data, raw_dtype=raw_dtype)

        total_bytes += len(decoded_bytes)
        if sha256_match:
            total_verified += 1

        tensor_stats.append({
            "name": name,
            "size_bytes": len(decoded_bytes),
            "sha256_match": sha256_match,
        })

        if (i + 1) % 25 == 0 or i == 0 or i == len(tensor_index) - 1:
            status = "✅" if sha256_match else "❌"
            print(f"  [{i + 1:3d}/{len(tensor_index)}] {status} {name[:45]:45s} "
                  f"{len(decoded_bytes):>12,} bytes")

    # Write GGUF
    print("\n[SEED-REGEN] Writing GGUF...")
    writer.write_header_to_file()
    writer.write_kv_data_to_file()
    writer.write_tensors_to_file()
    writer.close()

    t1 = time.time()

    # Summary
    print()
    print("=" * 70)
    print("SUMMARY")
    print("=" * 70)
    print(f"Tensors decoded: {len(tensor_stats)}")
    print(f"SHA256 verified: {total_verified}/{len(tensor_stats)}")
    print(f"Total bytes decoded: {total_bytes:,}")
    print(f"Output GGUF size: {out_path.stat().st_size:,} bytes")
    print(f"Runtime: {t1 - t0:.2f}s")

    all_verified = (total_verified == len(tensor_stats))

    if all_verified:
        print()
        print("✅ ALL TENSORS VERIFIED - GGUF regenerated from seed manifest")
        print()
        print("Next steps:")
        print("  1. Run weight correlation KAT vs known-good clone")
        print("  2. Run MMLU KAT for functional equivalence")
    else:
        print()
        print("❌ SOME TENSORS FAILED SHA256 VERIFICATION")

    # Build receipt
    receipt = {
        "receipt_type": "seed_regen_build",
        "work_order": "WO-REGEN-HELIX64-MANIFEST-V1",
        "created_at_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "manifest_path": str(manifest_path),
        "tokenizer_source": str(tokenizer_path),
        "output_path": str(out_path),
        "manifest_version": manifest.get("version"),
        "num_tensors": len(tensor_stats),
        "num_verified": total_verified,
        "all_verified": all_verified,
        "total_bytes_decoded": total_bytes,
        "output_size_bytes": out_path.stat().st_size,
        "runtime_sec": round(t1 - t0, 2),
        "status": "PASS" if all_verified else "FAIL",
    }

    with receipt_path.open("w") as f:
        json.dump(receipt, f, indent=2)

    print(f"\n[SEED-REGEN] Receipt written: {receipt_path}")


if __name__ == "__main__":
    main()
