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

WO-VENDOR-CLONE-CDNA-01:
Build a vendor-clone GGUF using the proven Helix64 codec.

This implements the vendor_clone_cdna regen mode that:
1. Reads vendor GGUF (Q4_K_M tensors)
2. Encodes each tensor to Helix64 (base64) - LOSSLESS
3. Decodes back from Helix64 - LOSSLESS
4. Writes new GGUF with vendor-identical weights

The roundtrip is bit-perfect because:
- Helix64 = base64 encoding (proven lossless in kat_cdna_roundtrip.py)
- Q4_K blocks survive roundtrip with identical bytes
- No FP32 materialization needed (stays in quantized domain)

Usage:
  PYTHONPATH=. python3 helix_cdc/regrow/vendor_clone_cdna_gguf.py \
    --vendor /path/to/vendor.gguf \
    --out artifacts/mistral_vendor_clone_cdna.gguf \
    --receipt receipts/kat_vendor_clone_cdna.json
"""

from __future__ import annotations

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

import numpy as np

# Add tools to path for helix64
sys.path.insert(0, str(pathlib.Path(__file__).parent.parent.parent / "tools"))

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

from helix64 import encode as helix64_encode, decode as helix64_decode


def extract_field_value(field, expected_type: Optional[GGUFValueType] = None) -> Any:
    """Extract value from GGUF field."""
    if len(field.types) != 1:
        return None

    vtype = field.types[0]
    if expected_type and vtype != expected_type:
        return None

    if vtype == GGUFValueType.STRING:
        return bytes(field.parts[-1]).decode('utf-8')
    elif vtype in (GGUFValueType.UINT32, GGUFValueType.INT32):
        return int(field.parts[-1][0])
    elif vtype in (GGUFValueType.UINT64, GGUFValueType.INT64):
        return int(field.parts[-1][0])
    elif vtype == GGUFValueType.FLOAT32:
        return float(field.parts[-1][0])
    elif vtype == GGUFValueType.FLOAT64:
        return float(field.parts[-1][0])
    elif vtype == GGUFValueType.BOOL:
        return bool(field.parts[-1][0])
    return None


def extract_string_array(field) -> List[str]:
    """
    Extract string array from GGUF field.

    GGUF string arrays have parts structured as:
      parts[0-4]: header metadata (field name len, field name bytes, array type, element type, count)
      parts[5:]: alternating [length, bytes] pairs for each string
    """
    tokens = []
    idx = 5  # Start after header parts
    while idx + 1 < len(field.parts):
        # Get string length (uint64)
        str_len_part = field.parts[idx]
        idx += 1

        # Get string bytes
        str_bytes_part = field.parts[idx]
        idx += 1

        # Decode string
        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.

    GGUF numeric arrays have parts structured as:
      parts[0-4]: header metadata
      parts[5:]: individual values as (1,) arrays
    """
    values = []
    for i in range(5, len(field.parts)):
        val = dtype(field.parts[i][0])
        values.append(val)
    return values


def copy_metadata(reader: GGUFReader, writer: GGUFWriter) -> Dict[str, Any]:
    """Copy key metadata from vendor to writer. Returns copied metadata dict."""
    metadata = {}

    # Essential model metadata to copy
    field_mapping = {
        'general.name': ('add_name', str),
        'llama.context_length': ('add_context_length', int),
        'llama.embedding_length': ('add_embedding_length', int),
        'llama.block_count': ('add_block_count', int),
        'llama.feed_forward_length': ('add_feed_forward_length', int),
        'llama.attention.head_count': ('add_head_count', int),
        'llama.attention.head_count_kv': ('add_head_count_kv', int),
        'llama.rope.dimension_count': ('add_rope_dimension_count', int),
        'llama.attention.layer_norm_rms_epsilon': ('add_layer_norm_rms_eps', float),
        'llama.rope.freq_base': ('add_rope_freq_base', float),
    }

    for field_name, (method_name, converter) in field_mapping.items():
        if field_name in reader.fields:
            value = extract_field_value(reader.fields[field_name])
            if value is not None:
                try:
                    converted = converter(value)
                    getattr(writer, method_name)(converted)
                    metadata[field_name] = converted
                except Exception:
                    pass

    # CRITICAL: Copy tokenizer vocabulary (required for inference!)
    if 'tokenizer.ggml.tokens' in reader.fields:
        tokens_field = reader.fields['tokenizer.ggml.tokens']
        # Use proper extraction for GGUF string arrays
        tokens = extract_string_array(tokens_field)
        writer.add_token_list(tokens)
        metadata['tokenizer.ggml.tokens'] = f"[{len(tokens)} tokens]"
        print(f"  Copied {len(tokens)} tokens to vocabulary")

        # Add token scores if available - use proper extraction for numeric arrays
        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)
            print(f"  Copied {len(scores)} token scores")

        # Add token types if available - use proper extraction for numeric arrays
        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)
            print(f"  Copied {len(toktypes)} token types")

    # Copy tokenizer model type (llama, gpt2, etc)
    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

    # Copy special token IDs
    for field_name in ['tokenizer.ggml.bos_token_id', 'tokenizer.ggml.eos_token_id',
                       'tokenizer.ggml.padding_token_id', 'tokenizer.ggml.unknown_token_id']:
        if field_name in reader.fields:
            value = extract_field_value(reader.fields[field_name])
            if value is not None:
                metadata[field_name] = value
                if 'bos_token_id' in field_name:
                    writer.add_bos_token_id(int(value))
                elif 'eos_token_id' in field_name:
                    writer.add_eos_token_id(int(value))
                elif 'padding_token_id' in field_name:
                    writer.add_pad_token_id(int(value))
                elif 'unknown_token_id' in field_name:
                    writer.add_unk_token_id(int(value))

    # Copy chat template if present
    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'] = f"[{len(template)} chars]"

    # Copy add_bos/add_eos flags
    for flag in ['tokenizer.ggml.add_bos_token', 'tokenizer.ggml.add_eos_token']:
        if flag in reader.fields:
            value = extract_field_value(reader.fields[flag])
            if value is not None:
                metadata[flag] = value
                if 'add_bos' in flag:
                    writer.add_add_bos_token(bool(value))
                elif 'add_eos' in flag:
                    writer.add_add_eos_token(bool(value))

    # Copy 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

    return metadata


def vendor_clone_tensor_cdna(
    vendor_data: np.ndarray,
    tensor_name: str,
) -> tuple[np.ndarray, Dict[str, Any]]:
    """
    Clone vendor tensor through Helix64 CDNA roundtrip.

    Pipeline:
        vendor Q4_K bytes → helix64_encode (base64) → helix64_decode → clone bytes

    This is proven LOSSLESS by kat_cdna_roundtrip.py.

    Returns:
        (cloned_data, stats_dict)
    """
    t0 = time.time()

    # Get raw bytes
    raw_bytes = vendor_data.tobytes()
    original_sha256 = hashlib.sha256(raw_bytes).hexdigest()

    # Helix64 encode (base64)
    encoded = helix64_encode(raw_bytes)

    # Helix64 decode (base64)
    decoded_bytes = helix64_decode(encoded)

    # Verify lossless
    decoded_sha256 = hashlib.sha256(decoded_bytes).hexdigest()
    is_identical = (raw_bytes == decoded_bytes)

    # Reconstruct numpy array with same shape
    cloned_data = np.frombuffer(decoded_bytes, dtype=vendor_data.dtype)
    cloned_data = cloned_data.reshape(vendor_data.shape)

    dt = time.time() - t0

    stats = {
        "name": tensor_name,
        "original_size": len(raw_bytes),
        "encoded_size": len(encoded),
        "is_lossless": is_identical,
        "sha256_match": original_sha256 == decoded_sha256,
        "original_sha256": original_sha256,
        "decoded_sha256": decoded_sha256,
        "duration_s": round(dt, 4),
    }

    return cloned_data, stats


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Build vendor_clone_cdna GGUF using Helix64 codec."
    )
    parser.add_argument(
        "--vendor",
        default="/home/voidstr3m33/echo_labs/projects/Helix/HELIX_CONSOLIDATED/ECHO_AI_VAULT/models/mistral-7b-instruct-v0.2.Q4_K_M.gguf",
        help="Path to vendor GGUF model",
    )
    parser.add_argument(
        "--out",
        default="artifacts/mistral_vendor_clone_cdna.gguf",
        help="Path to output GGUF",
    )
    parser.add_argument(
        "--receipt",
        default="receipts/kat_vendor_clone_cdna.json",
        help="Path to output JSON receipt",
    )
    parser.add_argument(
        "--max-tensors",
        type=int,
        default=0,
        help="Max tensors to process (0 = all)",
    )
    args = parser.parse_args()

    vendor_path = pathlib.Path(args.vendor)
    out_path = pathlib.Path(args.out)
    receipt_path = pathlib.Path(args.receipt)

    if not vendor_path.exists():
        raise SystemExit(f"Vendor model not found: {vendor_path}")

    print("=" * 70)
    print("WO-VENDOR-CLONE-CDNA-01")
    print("Vendor Clone via Helix64 CDNA Roundtrip")
    print("=" * 70)
    print(f"Vendor: {vendor_path}")
    print(f"Output: {out_path}")
    print()

    t0 = time.time()

    # Read vendor model
    print("Reading vendor model...")
    reader = GGUFReader(vendor_path)
    print(f"  Tensors: {len(reader.tensors)}")
    print(f"  Fields: {len(reader.fields)}")

    # Get architecture
    arch = "llama"  # Default
    if 'general.architecture' in reader.fields:
        arch = extract_field_value(reader.fields['general.architecture'])
        print(f"  Architecture: {arch}")

    # Create output directory
    out_path.parent.mkdir(parents=True, exist_ok=True)

    # Create writer
    writer = GGUFWriter(str(out_path), arch)

    # Copy metadata
    print("\nCopying metadata...")
    metadata = copy_metadata(reader, writer)
    print(f"  Copied {len(metadata)} fields")

    # Process tensors
    tensors = list(reader.tensors)
    if args.max_tensors > 0:
        tensors = tensors[:args.max_tensors]

    print(f"\nProcessing {len(tensors)} tensors via CDNA roundtrip...")

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

    for i, t in enumerate(tensors):
        # Clone through CDNA roundtrip
        cloned_data, stats = vendor_clone_tensor_cdna(
            vendor_data=np.array(t.data),
            tensor_name=t.name,
        )
        tensor_stats.append(stats)

        total_bytes += stats["original_size"]
        if stats["is_lossless"]:
            total_lossless += 1

        # Add to writer - use vendor's physical shape
        writer.add_tensor(
            name=t.name,
            tensor=cloned_data,
            raw_dtype=GGMLQuantizationType(t.tensor_type),
        )

        status_icon = "✅" if stats["is_lossless"] else "❌"
        if (i + 1) % 50 == 0 or i == len(tensors) - 1:
            print(f"  [{i+1}/{len(tensors)}] {status_icon} {t.name[:40]:40s} "
                  f"({stats['original_size']:,} bytes)")

    # Write GGUF
    print("\nWriting 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 processed: {len(tensor_stats)}")
    print(f"Lossless roundtrips: {total_lossless}/{len(tensor_stats)}")
    print(f"Total bytes processed: {total_bytes:,}")
    print(f"Output size: {out_path.stat().st_size:,} bytes")
    print(f"Runtime: {t1 - t0:.2f}s")

    # Verify all lossless
    all_lossless = (total_lossless == len(tensor_stats))

    if all_lossless:
        print()
        print("✅ ALL TENSORS LOSSLESS - vendor_clone_cdna SUCCESSFUL")
        print()
        print("Next steps:")
        print("  1. Run kat_regen_weight_correlation.py to verify correlation")
        print("  2. Run llama-cli inference test on the cloned model")
    else:
        print()
        print("❌ SOME TENSORS HAD LOSS - investigate failures")

    # Build receipt
    receipt = {
        "receipt_type": "vendor_clone_cdna",
        "receipt_id": "kat_vendor_clone_cdna",
        "work_order": "WO-VENDOR-CLONE-CDNA-01",
        "created_at_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "vendor_model_path": str(vendor_path),
        "output_model_path": str(out_path),
        "vendor_model_size": vendor_path.stat().st_size,
        "output_model_size": out_path.stat().st_size,
        "num_tensors": len(tensor_stats),
        "num_lossless": total_lossless,
        "total_bytes_processed": total_bytes,
        "runtime_sec": round(t1 - t0, 2),
        "all_lossless": all_lossless,
        "status": "PASS" if all_lossless else "FAIL",
        "metadata_copied": metadata,
        "tensor_stats": tensor_stats[:50],  # First 50 for readability
    }

    # Write receipt
    receipt_path.parent.mkdir(parents=True, exist_ok=True)
    with receipt_path.open("w") as f:
        json.dump(receipt, f, indent=2)

    print(f"\nReceipt written to: {receipt_path}")


if __name__ == "__main__":
    main()
