#!/usr/bin/env python3
"""
scanner.py - GGUF File Scanner
==============================

Scans GGUF files to extract tensor metadata and dequantize tensors.
"""

import struct
from pathlib import Path
from typing import Dict, Any, List, Optional
import numpy as np

# Try to import gguf package
try:
    from gguf.quants import dequantize
    from gguf.constants import GGMLQuantizationType
    HAVE_GGUF = True
except ImportError:
    HAVE_GGUF = False
    GGMLQuantizationType = None

# GGML type mappings
GGML_TYPE_MAP = {}
if HAVE_GGUF:
    GGML_TYPE_MAP = {
        0: ('F32', GGMLQuantizationType.F32, (1, 4)),
        1: ('F16', GGMLQuantizationType.F16, (1, 2)),
        2: ('Q4_0', GGMLQuantizationType.Q4_0, (32, 18)),
        3: ('Q4_1', GGMLQuantizationType.Q4_1, (32, 20)),
        6: ('Q5_0', GGMLQuantizationType.Q5_0, (32, 22)),
        7: ('Q5_1', GGMLQuantizationType.Q5_1, (32, 24)),
        8: ('Q8_0', GGMLQuantizationType.Q8_0, (32, 34)),
        10: ('Q2_K', GGMLQuantizationType.Q2_K, (256, 84)),
        11: ('Q3_K', GGMLQuantizationType.Q3_K, (256, 110)),
        12: ('Q4_K', GGMLQuantizationType.Q4_K, (256, 144)),
        13: ('Q5_K', GGMLQuantizationType.Q5_K, (256, 176)),
        14: ('Q6_K', GGMLQuantizationType.Q6_K, (256, 210)),
    }
else:
    GGML_TYPE_MAP = {
        0: ('F32', None, (1, 4)),
        1: ('F16', None, (1, 2)),
        2: ('Q4_0', None, (32, 18)),
        3: ('Q4_1', None, (32, 20)),
        6: ('Q5_0', None, (32, 22)),
        7: ('Q5_1', None, (32, 24)),
        8: ('Q8_0', None, (32, 34)),
        10: ('Q2_K', None, (256, 84)),
        11: ('Q3_K', None, (256, 110)),
        12: ('Q4_K', None, (256, 144)),
        13: ('Q5_K', None, (256, 176)),
        14: ('Q6_K', None, (256, 210)),
    }


class GGUFScanner:
    """Scan GGUF file and extract tensor info."""

    def __init__(self, gguf_path: Path):
        self.path = Path(gguf_path)
        self.tensors: List[Dict[str, Any]] = []
        self.metadata: Dict[str, Any] = {}
        self.data_offset = 0

    def _read_string(self, f) -> str:
        length = struct.unpack('<Q', f.read(8))[0]
        return f.read(length).decode('utf-8')

    def _skip_kv_pairs(self, f, n_kv: int) -> Dict[str, Any]:
        """Skip key-value pairs but capture basic metadata."""
        metadata = {}
        GGUF_TYPES = {
            0: (1, '<B'), 1: (1, '<b'), 2: (2, '<H'), 3: (2, '<h'),
            4: (4, '<I'), 5: (4, '<i'), 6: (4, '<f'), 7: (1, '<B'),
            8: None, 9: None, 10: (8, '<Q'), 11: (8, '<q'), 12: (8, '<d'),
        }

        for _ in range(n_kv):
            key = self._read_string(f)
            vtype = struct.unpack('<I', f.read(4))[0]

            value = None
            if vtype == 8:  # String
                value = self._read_string(f)
            elif vtype == 9:  # Array
                arr_type = struct.unpack('<I', f.read(4))[0]
                arr_len = struct.unpack('<Q', f.read(8))[0]
                if arr_type == 8:
                    for _ in range(arr_len):
                        self._read_string(f)
                elif arr_type in GGUF_TYPES and GGUF_TYPES[arr_type]:
                    size, _ = GGUF_TYPES[arr_type]
                    f.read(size * arr_len)
            elif vtype in GGUF_TYPES and GGUF_TYPES[vtype]:
                size, fmt = GGUF_TYPES[vtype]
                raw = f.read(size)
                if size <= 8:
                    value = struct.unpack(fmt, raw)[0]

            if value is not None:
                metadata[key] = value

        return metadata

    def scan(self) -> 'GGUFScanner':
        """Scan GGUF file for all tensors."""
        print(f"Scanning: {self.path}")

        with open(self.path, 'rb') as f:
            # Header
            magic = f.read(4)
            if magic != b'GGUF':
                raise ValueError("Not a GGUF file")

            version = struct.unpack('<I', f.read(4))[0]
            n_tensors = struct.unpack('<Q', f.read(8))[0]
            n_kv = struct.unpack('<Q', f.read(8))[0]

            print(f"  GGUF v{version}: {n_tensors} tensors, {n_kv} KV pairs")

            self.metadata = self._skip_kv_pairs(f, n_kv)

            # Read tensor info
            for _ in range(n_tensors):
                name = self._read_string(f)
                n_dims = struct.unpack('<I', f.read(4))[0]
                dims = [struct.unpack('<Q', f.read(8))[0] for _ in range(n_dims)]
                dtype_int = struct.unpack('<I', f.read(4))[0]
                offset = struct.unpack('<Q', f.read(8))[0]

                dtype_info = GGML_TYPE_MAP.get(dtype_int, (f'UNK_{dtype_int}', None, (1, 1)))
                dtype_name, qtype, (block_size, type_size) = dtype_info

                n_elements = 1
                for d in dims:
                    n_elements *= d
                n_blocks = (n_elements + block_size - 1) // block_size
                byte_size = n_blocks * type_size

                self.tensors.append({
                    'name': name,
                    'dims': dims,
                    'dtype_int': dtype_int,
                    'dtype_name': dtype_name,
                    'qtype': qtype,
                    'offset': offset,
                    'n_elements': n_elements,
                    'byte_size': byte_size,
                    'block_size': block_size,
                    'type_size': type_size,
                })

            # Calculate data start offset
            pos = f.tell()
            self.data_offset = (pos + 31) // 32 * 32

        return self

    def extract_tensor(self, tensor_info: Dict[str, Any]) -> Optional[np.ndarray]:
        """Extract and dequantize a single tensor.

        Uses gguf library's GGUFReader for correct tensor loading.
        The gguf library handles Q8_0 and other quant types correctly.
        """
        if not HAVE_GGUF:
            print(f"  WARNING: gguf package not available, using synthetic")
            return np.random.randn(*tensor_info['dims']).astype(np.float32) * 0.01

        # Use gguf library's GGUFReader for correct dequantization
        # Our manual byte reading had wrong data layout for quantized formats
        from gguf import GGUFReader

        reader = GGUFReader(str(self.path))
        for t in reader.tensors:
            if t.name == tensor_info['name']:
                # The gguf library handles Q8_0 and other quant types correctly
                tensor = dequantize(t.data, t.tensor_type)
                return tensor.astype(np.float32)

        # Tensor not found
        print(f"  WARNING: {tensor_info['name']} not found in GGUF")
        return np.random.randn(*tensor_info['dims']).astype(np.float32) * 0.01

    def get_tensor_by_name(self, name: str) -> Optional[Dict[str, Any]]:
        """Get tensor info by name."""
        for t in self.tensors:
            if t['name'] == name:
                return t
        return None

    def filter_tensors(self, pattern: str) -> List[Dict[str, Any]]:
        """Filter tensors by name pattern."""
        return [t for t in self.tensors if pattern in t['name']]
