import argparse
import json
import os
import re

from collections import defaultdict
from pathlib import Path


LAYER_KEY_PATTERN = re.compile(r"^model\.layers\.(\d+)\.(.+)$")
CANDIDATE_INDEX_FILES = [
    "model.safetensors.index.json",
    "consolidated.safetensors.index.json",
]


def normalize_numeric_segments(name: str) -> str:
    """
    Replace pure numeric path segments with '*'.
    Example: model.layers.12.mlp.experts.3.up_proj.weight
          -> model.layers.*.mlp.experts.*.up_proj.weight
    """
    return ".".join("*" if part.isdigit() else part for part in name.split("."))


def load_weight_map(input_json_path: str) -> dict[str, str]:
    with open(input_json_path, "r") as f:
        data = json.load(f)

    if isinstance(data, dict) and isinstance(data.get("weight_map"), dict):
        return data["weight_map"]

    if isinstance(data, dict) and all(
        isinstance(k, str) and isinstance(v, str) for k, v in data.items()
    ):
        return data

    raise ValueError(
        "Input json must be either {'weight_map': {...}} or a plain key->file mapping."
    )


def find_index_json(input_dir: str) -> str:
    for name in CANDIDATE_INDEX_FILES:
        path = os.path.join(input_dir, name)
        if os.path.exists(path):
            return path
    raise ValueError(f"Can not find index file in {input_dir}")


def load_weight_map_from_input_dir(input_dir: str) -> tuple[dict[str, str], str]:
    index_json_path = find_index_json(input_dir)
    return load_weight_map(index_json_path), index_json_path


def build_structure_report(
    weight_map: dict[str, str],
    max_examples: int,
    max_pattern_preview: int,
    include_all_layer_patterns: bool,
) -> dict:
    pattern_to_keys: dict[str, list[str]] = defaultdict(list)
    layer_to_suffix_patterns: dict[int, set[str]] = defaultdict(set)
    non_layer_patterns: set[str] = set()

    for key in weight_map.keys():
        normalized_full = normalize_numeric_segments(key)
        pattern_to_keys[normalized_full].append(key)

        layer_match = LAYER_KEY_PATTERN.match(key)
        if layer_match is None:
            non_layer_patterns.add(normalized_full)
            continue

        layer_id = int(layer_match.group(1))
        layer_suffix = layer_match.group(2)
        normalized_suffix = normalize_numeric_segments(layer_suffix)
        layer_to_suffix_patterns[layer_id].add(normalized_suffix)

    normalized_key_patterns = []
    for pattern, keys in sorted(
        pattern_to_keys.items(),
        key=lambda item: (-len(item[1]), item[0]),
    ):
        normalized_key_patterns.append(
            {
                "pattern": pattern,
                "count": len(keys),
                "example_keys": sorted(keys)[:max_examples],
            }
        )

    all_patterns = [item["pattern"] for item in normalized_key_patterns]

    signature_to_layers: dict[tuple[str, ...], list[int]] = defaultdict(list)
    for layer_id, suffix_patterns in layer_to_suffix_patterns.items():
        signature = tuple(sorted(suffix_patterns))
        signature_to_layers[signature].append(layer_id)

    layer_structure_groups = []
    sorted_signature_items = sorted(
        signature_to_layers.items(),
        key=lambda item: (len(item[1]), item[1][0]),
        reverse=True,
    )
    for idx, (signature, layer_ids) in enumerate(sorted_signature_items, start=1):
        layer_ids = sorted(layer_ids)
        entry = {
            "group_id": f"layer_struct_{idx:03d}",
            "layer_ids": layer_ids,
            "layer_count": len(layer_ids),
            "pattern_count": len(signature),
            "pattern_preview": list(signature)[:max_pattern_preview],
        }
        if include_all_layer_patterns:
            entry["patterns"] = list(signature)
        layer_structure_groups.append(entry)

    return {
        "all_patterns": all_patterns,
        "total_weight_keys": len(weight_map),
        "normalized_key_pattern_count": len(normalized_key_patterns),
        "normalized_key_patterns": normalized_key_patterns,
        "layer_ids_seen": sorted(layer_to_suffix_patterns.keys()),
        "layer_structure_group_count": len(layer_structure_groups),
        "layer_structure_groups": layer_structure_groups,
        "non_layer_pattern_count": len(non_layer_patterns),
        "non_layer_patterns": sorted(non_layer_patterns),
    }


def get_default_output_path(input_json_path: str) -> str:
    input_path = Path(input_json_path)
    return str(input_path.with_name(f"{input_path.stem}.layer_structures.json"))


def resolve_io_paths(
    input_dir: str | None, input_json: str | None, output_json: str | None
) -> tuple[dict[str, str], str, str]:
    if not input_dir and not input_json:
        raise ValueError("Either --input-dir or --input-json is required.")
    if input_dir and input_json:
        raise ValueError("Please provide only one of --input-dir or --input-json.")

    if input_dir:
        weight_map, input_json_path = load_weight_map_from_input_dir(input_dir)
    else:
        input_json_path = input_json
        weight_map = load_weight_map(input_json_path)

    output_json_path = output_json or get_default_output_path(input_json_path)
    return weight_map, input_json_path, output_json_path


def main():
    parser = argparse.ArgumentParser(
        description="Group weight_map keys by normalized pattern and layer structure."
    )
    parser.add_argument("--input-dir", default=None, type=str)
    parser.add_argument("--input-json", default=None, type=str)
    parser.add_argument("--output-json", default=None, type=str)
    parser.add_argument("--max-examples", default=3, type=int)
    parser.add_argument("--max-pattern-preview", default=20, type=int)
    parser.add_argument("--include-all-layer-patterns", action="store_true")

    args = parser.parse_args()
    weight_map, input_json_path, output_json = resolve_io_paths(
        input_dir=args.input_dir,
        input_json=args.input_json,
        output_json=args.output_json,
    )
    report = build_structure_report(
        weight_map=weight_map,
        max_examples=args.max_examples,
        max_pattern_preview=args.max_pattern_preview,
        include_all_layer_patterns=args.include_all_layer_patterns,
    )

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

    print(f"Saved grouped structure report to: {output_json}")
    print(f"Loaded weight_map from: {input_json_path}")
    print(f"Total keys: {report['total_weight_keys']}")
    print(f"Layer structure groups: {report['layer_structure_group_count']}")


if __name__ == "__main__":
    main()