"""Extract sub-checkpoints from safetensors shards without loading the full model.

Memory-efficient extraction of arbitrary subnets from HuggingFace-style
sharded checkpoints. Uses safetensors mmap (lazy tensor read) so peak RAM
is bounded by the single largest tensor, not by total model size.

This works for ANY model stored as safetensors shards with an index JSON
(HuggingFace convention: ``model.safetensors.index.json`` with a
``weight_map`` dict mapping ``key -> shard_filename``).

Typical workflow::

    from agiws_neural_quant import extract_subcheckpoint

    # Extract Qwen3-VL vision encoder weights (prefix ``model.visual.``)
    extract_subcheckpoint(
        index_path="models/qwen3_vl_8b/model.safetensors.index.json",
        key_prefix="model.visual.",
        output_path="models/qwen3_vl_vision.safetensors",
        strip_prefix="model.visual.",
    )

    # Then load into a vision-only model:
    from safetensors.torch import load_file
    sd = load_file("models/qwen3_vl_vision.safetensors")
    vision_model.load_state_dict(sd, strict=True)

For 1.5 TB models (e.g. GLM-5.2) this never touches more than one tensor
at a time in RAM — the shard files are mmap'd, only matching keys are read,
and the output file is written incrementally.
"""

import fnmatch
import json
import re
from pathlib import Path
from typing import Optional, Union

from safetensors import safe_open
from safetensors.torch import save_file

PathLike = Union[str, Path]


class ExtractReport:
    """Summary of an extraction operation."""

    def __init__(
        self,
        matched: int,
        skipped: int,
        total_keys: int,
        shards_read: list[str],
        output_path: str,
        total_bytes: int,
    ):
        self.matched = matched
        self.skipped = skipped
        self.total_keys = total_keys
        self.shards_read = sorted(shards_read)
        self.output_path = output_path
        self.total_bytes = total_bytes

    def __repr__(self) -> str:
        return (
            f"ExtractReport(matched={self.matched}, skipped={self.skipped}, "
            f"total_keys={self.total_keys}, shards_read={len(self.shards_read)}, "
            f"output={self.total_bytes / 1e6:.1f} MB)"
        )

    def __str__(self) -> str:
        lines = [
            f"Extraction complete:",
            f"  Matched keys:   {self.matched} / {self.total_keys}",
            f"  Skipped keys:    {self.skipped}",
            f"  Shards read:     {len(self.shards_read)}",
            f"  Output size:     {self.total_bytes / 1e6:.2f} MB",
            f"  Output path:     {self.output_path}",
        ]
        if self.shards_read:
            lines.append(f"  Shards:          {', '.join(self.shards_read)}")
        return "\n".join(lines)


def _load_weight_map(index_path: PathLike) -> dict[str, str]:
    """Read ``weight_map`` from a HuggingFace safetensors index JSON.

    Args:
        index_path: Path to ``model.safetensors.index.json``.

    Returns:
        Dict mapping ``key -> shard_filename``.
    """
    index_path = Path(index_path)
    if not index_path.is_file():
        raise FileNotFoundError(f"Index file not found: {index_path}")
    with open(index_path, "r", encoding="utf-8") as f:
        index = json.load(f)
    weight_map = index.get("weight_map")
    if not weight_map:
        raise ValueError(
            f"Index file {index_path} has no 'weight_map' key. "
            "Expected HuggingFace safetensors index format."
        )
    return weight_map


def _match_keys(
    weight_map: dict[str, str],
    key_prefix: Optional[str] = None,
    key_glob: Optional[str] = None,
    key_regex: Optional[str] = None,
) -> list[str]:
    """Select keys from weight_map by prefix, glob, or regex.

    Args:
        weight_map: Full weight_map dict.
        key_prefix: If set, keep only keys starting with this prefix.
        key_glob: If set, keep only keys matching this fnmatch glob.
        key_regex: If set, keep only keys matching this regex pattern.

    Returns:
        Sorted list of matching keys.

    Raises:
        ValueError: If none of prefix/glob/regex is provided.
    """
    if key_prefix is None and key_glob is None and key_regex is None:
        raise ValueError(
            "At least one of key_prefix, key_glob, key_regex must be provided."
        )
    keys = list(weight_map.keys())
    if key_prefix is not None:
        keys = [k for k in keys if k.startswith(key_prefix)]
    if key_glob is not None:
        keys = [k for k in keys if fnmatch.fnmatch(k, key_glob)]
    if key_regex is not None:
        pat = re.compile(key_regex)
        keys = [k for k in keys if pat.search(k)]
    return sorted(keys)


def extract_subcheckpoint(
    index_path: PathLike,
    output_path: PathLike,
    key_prefix: Optional[str] = None,
    key_glob: Optional[str] = None,
    key_regex: Optional[str] = None,
    strip_prefix: Optional[str] = None,
    dtype: Optional[str] = None,
    verbose: bool = True,
) -> ExtractReport:
    """Extract matching tensors from sharded safetensors into a new single-file checkpoint.

    Reads only the shards that contain matching keys (mmap, lazy). Peak RAM
    is bounded by the largest single tensor, not by total model size.

    Args:
        index_path: Path to ``model.safetensors.index.json``.
        output_path: Where to write the extracted ``.safetensors`` file.
        key_prefix: Keep only keys starting with this prefix.
        key_glob: Keep only keys matching this fnmatch glob (e.g. ``"*.weight"``).
        key_regex: Keep only keys matching this regex.
        strip_prefix: If set, remove this prefix from each key in the output.
            Useful when the subnet checkpoint should load directly into a
            standalone model (e.g. strip ``"model.visual."`` so keys match
            ``Qwen3VLVisionModel`` state_dict).
        dtype: Optional cast — ``"fp32"``, ``"fp16"``, ``"bf16"``. Default:
            keep original dtype.
        verbose: Print progress and summary.

    Returns:
        ExtractReport with counts and file info.
    """
    index_path = Path(index_path)
    output_path = Path(output_path)
    shard_dir = index_path.parent

    weight_map = _load_weight_map(index_path)
    matched_keys = _match_keys(weight_map, key_prefix, key_glob, key_regex)

    if not matched_keys:
        raise ValueError(
            "No keys matched the given criteria. "
            f"Total keys in index: {len(weight_map)}."
        )

    # Group matched keys by shard file
    shard_to_keys: dict[str, list[str]] = {}
    for k in matched_keys:
        shard_name = weight_map[k]
        shard_to_keys.setdefault(shard_name, []).append(k)

    # Read tensors one shard at a time (mmap), collect into output dict
    output_sd: dict[str, "torch.Tensor"] = {}  # noqa: F821
    shards_read: list[str] = []
    total_bytes = 0

    import torch  # local import — extract module should be importable without torch

    for shard_name in sorted(shard_to_keys.keys()):
        shard_path = shard_dir / shard_name
        if not shard_path.is_file():
            raise FileNotFoundError(
                f"Shard file not found: {shard_path} "
                f"(referenced by index {index_path.name})"
            )
        if verbose:
            print(f"  Reading shard: {shard_name} "
                  f"({len(shard_to_keys[shard_name])} keys)")

        with safe_open(str(shard_path), framework="pt", device="cpu") as f:
            for key in shard_to_keys[shard_name]:
                tensor = f.get_tensor(key)
                if dtype is not None:
                    dt_map = {
                        "fp32": torch.float32,
                        "fp16": torch.float16,
                        "bf16": torch.bfloat16,
                    }
                    if dtype not in dt_map:
                        raise ValueError(
                            f"dtype={dtype!r} not supported. "
                            "Use 'fp32', 'fp16', or 'bf16'."
                        )
                    tensor = tensor.to(dt_map[dtype])
                # Strip prefix if requested
                out_key = key
                if strip_prefix and out_key.startswith(strip_prefix):
                    out_key = out_key[len(strip_prefix):]
                output_sd[out_key] = tensor
                total_bytes += tensor.numel() * tensor.element_size()
        shards_read.append(shard_name)

    # Write output checkpoint
    output_path.parent.mkdir(parents=True, exist_ok=True)
    save_file(output_sd, str(output_path), metadata={
        "format": "pt",
        "extracted_keys": str(len(output_sd)),
        "source_index": index_path.name,
    })

    report = ExtractReport(
        matched=len(matched_keys),
        skipped=len(weight_map) - len(matched_keys),
        total_keys=len(weight_map),
        shards_read=shards_read,
        output_path=str(output_path),
        total_bytes=total_bytes,
    )
    if verbose:
        print(report)
    return report


def list_shards(index_path: PathLike) -> dict[str, list[str]]:
    """List all shard files and their keys from an index JSON.

    Useful for exploring a checkpoint before extracting.

    Args:
        index_path: Path to ``model.safetensors.index.json``.

    Returns:
        Dict mapping ``shard_filename -> [sorted keys]``.
    """
    weight_map = _load_weight_map(index_path)
    shard_to_keys: dict[str, list[str]] = {}
    for k, shard in weight_map.items():
        shard_to_keys.setdefault(shard, []).append(k)
    for shard in shard_to_keys:
        shard_to_keys[shard].sort()
    return shard_to_keys


def find_keys(
    index_path: PathLike,
    pattern: str,
    mode: str = "prefix",
    limit: int = 0,
) -> list[str]:
    """Find keys in an index by pattern, without loading any tensors.

    Args:
        index_path: Path to ``model.safetensors.index.json``.
        pattern: Search pattern.
        mode: ``"prefix"``, ``"glob"``, or ``"regex"``.
        limit: Max results (0 = unlimited).

    Returns:
        Sorted list of matching keys.
    """
    weight_map = _load_weight_map(index_path)
    if mode == "prefix":
        keys = sorted(k for k in weight_map if k.startswith(pattern))
    elif mode == "glob":
        keys = sorted(k for k in weight_map if fnmatch.fnmatch(k, pattern))
    elif mode == "regex":
        pat = re.compile(pattern)
        keys = sorted(k for k in weight_map if pat.search(k))
    else:
        raise ValueError(f"mode={mode!r} not supported. Use 'prefix', 'glob', or 'regex'.")
    if limit > 0:
        keys = keys[:limit]
    return keys


def load_subcheckpoint(
    checkpoint_path: PathLike,
    model: "torch.nn.Module",  # noqa: F821
    strict: bool = True,
    device: str = "cpu",
) -> "torch.nn.Module":  # noqa: F821
    """Load an extracted sub-checkpoint into a model.

    Convenience wrapper around ``safetensors.torch.load_file`` +
    ``model.load_state_dict``.

    Args:
        checkpoint_path: Path to ``.safetensors`` file (from extract_subcheckpoint).
        model: PyTorch model to load weights into.
        strict: If True, state_dict keys must match exactly.
        device: Device to load tensors onto before loading into model.

    Returns:
        The model (loaded in-place).
    """
    from safetensors.torch import load_file

    sd = load_file(str(checkpoint_path), device=device)
    model.load_state_dict(sd, strict=strict)
    return model


# ---------------------------------------------------------------------------
# Convenience: extract vision encoder from Qwen3-VL (and similar architectures)
# ---------------------------------------------------------------------------

def extract_vision_encoder(
    model_dir: PathLike,
    output_path: Optional[PathLike] = None,
    vision_prefix: str = "model.visual.",
    dtype: Optional[str] = None,
    verbose: bool = True,
) -> ExtractReport:
    """Extract vision encoder weights from a multimodal model checkpoint.

    Convenience function for the common case: extract all keys with
    ``model.visual.`` prefix, strip the prefix, and save as a standalone
    safetensors checkpoint that can be loaded directly into
    ``Qwen3VLVisionModel``.

    Args:
        model_dir: Directory containing ``model.safetensors.index.json``
            and shard files.
        output_path: Where to write the extracted checkpoint. Default:
            ``{model_dir}/vision_encoder.safetensors``.
        vision_prefix: Key prefix for vision encoder weights.
        dtype: Optional cast (``"fp32"``, ``"fp16"``, ``"bf16"``).
        verbose: Print progress.

    Returns:
        ExtractReport.
    """
    model_dir = Path(model_dir)
    index_path = model_dir / "model.safetensors.index.json"
    if output_path is None:
        output_path = model_dir / "vision_encoder.safetensors"
    return extract_subcheckpoint(
        index_path=index_path,
        output_path=output_path,
        key_prefix=vision_prefix,
        strip_prefix=vision_prefix,
        dtype=dtype,
        verbose=verbose,
    )


def extract_module_group(
    model_dir: PathLike,
    module_prefix: str,
    output_path: Optional[PathLike] = None,
    strip_prefix: Optional[str] = None,
    dtype: Optional[str] = None,
    verbose: bool = True,
) -> ExtractReport:
    """Extract an arbitrary module group from a checkpoint by prefix.

    General-purpose convenience: extract any subtree of a model by its
    parameter name prefix. Works for attention heads, MLP blocks, layer
    ranges, or any other logical grouping.

    Examples::

        # Extract attention layers from transformer block 5
        extract_module_group(
            model_dir="models/glm5",
            module_prefix="model.layers.5.self_attn.",
            output_path="models/glm5_attn_layer5.safetensors",
            strip_prefix="model.layers.5.self_attn.",
        )

        # Extract all MLP weights across all layers
        extract_module_group(
            model_dir="models/glm5",
            module_prefix="model.layers.",
            key_glob="*mlp.*",
            output_path="models/glm5_all_mlp.safetensors",
        )

    Args:
        model_dir: Directory with ``model.safetensors.index.json``.
        module_prefix: Parameter name prefix to match.
        output_path: Output ``.safetensors`` path. Default:
            ``{model_dir}/{prefix_sanitized}.safetensors``.
        strip_prefix: Prefix to strip from keys in output.
        dtype: Optional dtype cast.
        verbose: Print progress.

    Returns:
        ExtractReport.
    """
    model_dir = Path(model_dir)
    index_path = model_dir / "model.safetensors.index.json"
    if output_path is None:
        safe = module_prefix.replace(".", "_").strip("_")
        output_path = model_dir / f"{safe}.safetensors"
    if strip_prefix is None:
        strip_prefix = module_prefix
    return extract_subcheckpoint(
        index_path=index_path,
        output_path=output_path,
        key_prefix=module_prefix,
        strip_prefix=strip_prefix,
        dtype=dtype,
        verbose=verbose,
    )