# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import re
from collections import defaultdict
from dataclasses import dataclass, field
from typing import TYPE_CHECKING

import paddle
from paddle.distributed.fleet.utils.log_util import logger

if TYPE_CHECKING:
    from paddle.distributed.communication.group import Group

    from ..aoa.aoa_engine import AOAEngine
    from .metadata import Metadata

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

_MAX_TOTAL_LINES = 500
_MAX_KEYS_SHOWN = 50
_MAX_SHAPE_MISMATCHES = 20
_MAX_PATTERNS_SHOWN = 30
_SRC_FOLD_THRESHOLD = 5
_MAX_SLICE_DETAIL_KEYS = 5


def _get_rank() -> int:
    return paddle.distributed.get_rank()


# ---------------------------------------------------------------------------
# Color support (disabled by default)
# ---------------------------------------------------------------------------


class _C:
    """No-op color helpers. Colors are disabled by default."""

    @staticmethod
    def green(t):
        return t

    @staticmethod
    def yellow(t):
        return t

    @staticmethod
    def red(t):
        return t

    @staticmethod
    def cyan(t):
        return t


# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------


@dataclass
class ShapeMismatchInfo:
    key: str
    src_global_shape: tuple[int, ...]
    dst_global_shape: tuple[int, ...]
    src_dtype: str | None = None
    dst_dtype: str | None = None


@dataclass
class KeyValidationResult:
    missing_keys: set[str] = field(default_factory=set)
    unexpected_keys: set[str] = field(default_factory=set)
    shape_mismatches: list[ShapeMismatchInfo] = field(default_factory=list)
    randomly_initialized_keys: set[str] = field(default_factory=set)


@dataclass
class AOASliceMapping:
    src_key: str
    src_slice: tuple[slice, ...]
    dst_slice: tuple[slice, ...]
    postprocess: list[str] | None = None


@dataclass
class AOAMappingEntry:
    dst_key: str
    dst_global_shape: tuple[int, ...]
    slice_mappings: list[AOASliceMapping] = field(default_factory=list)
    is_identity: bool = False


# ---------------------------------------------------------------------------
# Public API: Standard (non-AOA) validation
# ---------------------------------------------------------------------------


def validate_and_report_keys_standard(
    metadata_list: list[Metadata],
    state_dict_param_names: set[str],
    process_group: Group | None,
    use_dist: bool,
    checkpoint_path: str,
    state_dict: dict,
) -> KeyValidationResult:
    """Validate keys for the standard (non-AOA) loading path.

    Gathers global dst keys across all ranks, compares with global src keys,
    checks shape mismatches. Prints report on rank 0 only.
    """
    # 1. Gather global dst keys
    if use_dist:
        global_dst_key_list = []
        paddle.distributed.all_gather_object(
            global_dst_key_list, list(state_dict_param_names), process_group
        )
        global_dst_keys = {
            k for sublist in global_dst_key_list for k in sublist
        }
    else:
        global_dst_keys = state_dict_param_names

    # 2. Collect global src keys from metadata
    global_src_keys = set()
    for metadata in metadata_list:
        for local_tensor_index in metadata.storage_metadata:
            if (
                local_tensor_index.replica_id is not None
                and local_tensor_index.replica_id != 0
            ):
                continue
            global_src_keys.add(local_tensor_index.tensor_key)

    # 3. Compute missing / unexpected
    missing_keys = global_dst_keys - global_src_keys
    unexpected_keys = global_src_keys - global_dst_keys

    # 4. Check shape mismatches for matching keys
    shape_mismatches = []
    assert state_dict is not None, "state_dict must not be None"
    # Gather dst global shapes: {key: global_shape}
    local_dst_shapes = {}
    for key, val in state_dict.items():
        k = key if isinstance(key, str) else key[0]
        if hasattr(val, "global_shape"):
            local_dst_shapes[k] = tuple(val.global_shape)
        else:
            local_dst_shapes[k] = tuple(val.shape)

    if use_dist:
        all_dst_shapes_list = []
        paddle.distributed.all_gather_object(
            all_dst_shapes_list, local_dst_shapes, process_group
        )
        global_dst_shapes = {}
        for d in all_dst_shapes_list:
            global_dst_shapes.update(d)
    else:
        global_dst_shapes = local_dst_shapes

    # Build src global shapes from metadata
    src_global_shapes: dict[str, tuple[int, ...]] = {}
    for metadata in metadata_list:
        if not metadata.state_dict_metadata:
            continue
        for key, src_metas in metadata.state_dict_metadata.items():
            if not src_metas or src_metas[0].global_shape is None:
                continue
            src_global_shapes[key] = tuple(src_metas[0].global_shape)

    matching_keys = global_dst_keys & global_src_keys
    for key in sorted(matching_keys):
        src_shape = src_global_shapes.get(key)
        dst_shape = global_dst_shapes.get(key)
        if src_shape is None or dst_shape is None:
            continue
        if src_shape != dst_shape:
            shape_mismatches.append(
                ShapeMismatchInfo(
                    key=key,
                    src_global_shape=src_shape,
                    dst_global_shape=dst_shape,
                )
            )

    result = KeyValidationResult(
        missing_keys=missing_keys,
        unexpected_keys=unexpected_keys,
        shape_mismatches=shape_mismatches,
        randomly_initialized_keys=set(),
    )

    # 5. Print on rank 0 (or always when not using dist)
    if not use_dist or _get_rank() == 0:
        _print_standard_report(result, checkpoint_path, len(global_dst_keys))

    return result


# ---------------------------------------------------------------------------
# Public API: AOA validation
# ---------------------------------------------------------------------------


def validate_and_report_keys_aoa(
    aoa_engine: AOAEngine,
    metadata: Metadata,
    checkpoint_path: str,
    use_dist: bool = True,
) -> KeyValidationResult:
    """Validate keys for the AOA loading path.

    Called AFTER AOAEngine is initialized. Uses output_vars/input_vars to
    compute truly missing/unexpected keys and builds the mapping table.
    """
    # 1. Covered dst keys
    aoa_covered_dst_keys = {
        k for k, v in aoa_engine.output_vars.items() if v is not None
    }
    randomly_initialized_keys = set(aoa_engine.need_add_output_vars)

    # 2. Consumed src keys
    consumed_src_keys = set()
    for tensor_desc in aoa_engine.output_vars.values():
        if tensor_desc is None:
            continue
        for src_key, _, _, _ in tensor_desc.slices:
            consumed_src_keys.add(src_key)

    # 3. Explicitly removed / all src keys
    explicitly_removed = set(aoa_engine.need_remove_input_vars)
    all_src_keys = set(aoa_engine.input_vars.keys())

    # 4. Compute truly missing / unexpected
    dst_state_keys = aoa_engine.context.get_all_dst_state_keys()
    truly_missing = (
        dst_state_keys - aoa_covered_dst_keys - randomly_initialized_keys
    )
    truly_unexpected = all_src_keys - consumed_src_keys - explicitly_removed

    # 5. Build AOA mapping entries
    aoa_mappings = _build_aoa_mappings(aoa_engine)

    result = KeyValidationResult(
        missing_keys=truly_missing,
        unexpected_keys=truly_unexpected,
        shape_mismatches=[],
        randomly_initialized_keys=randomly_initialized_keys,
    )

    # 6. Print on rank 0 (or always when not using dist)
    if not use_dist or _get_rank() == 0:
        _print_aoa_report(
            result, aoa_mappings, explicitly_removed, checkpoint_path
        )

    return result


def _build_aoa_mappings(aoa_engine: AOAEngine) -> list[AOAMappingEntry]:
    """Extract mapping entries from AOA engine's output_vars."""
    entries = []
    for dst_key, tensor_desc in sorted(aoa_engine.output_vars.items()):
        if tensor_desc is None:
            continue
        shape = tuple(tensor_desc.shape)
        slice_mappings = []
        for src_key, src_sl, dst_sl, pp_list in tensor_desc.slices:
            slice_mappings.append(
                AOASliceMapping(
                    src_key=src_key,
                    src_slice=src_sl,
                    dst_slice=dst_sl,
                    postprocess=pp_list,
                )
            )
        # Determine if identity
        is_identity = (
            len(slice_mappings) == 1
            and slice_mappings[0].src_key == dst_key
            and slice_mappings[0].postprocess is None
            and _slice_covers_full(slice_mappings[0].dst_slice, shape)
        )
        entries.append(
            AOAMappingEntry(
                dst_key=dst_key,
                dst_global_shape=shape,
                slice_mappings=slice_mappings,
                is_identity=is_identity,
            )
        )
    return entries


def _slice_covers_full(sl: tuple[slice, ...], shape: tuple[int, ...]) -> bool:
    """Check if a slice tuple covers the full tensor."""
    if len(sl) != len(shape):
        return False
    for s, dim in zip(sl, shape):
        if s.start != 0 or s.stop != dim:
            return False
    return True


# ---------------------------------------------------------------------------
# Printing: Standard report
# ---------------------------------------------------------------------------

_SEP = "=" * 70
_THIN_SEP = "-" * 70


def _print_standard_report(
    result: KeyValidationResult, path: str, total_keys: int
) -> None:
    lines = [_SEP, f"FlexCheckpoint Load Report (Checkpoint: {path})", _SEP]

    if (
        not result.missing_keys
        and not result.unexpected_keys
        and not result.shape_mismatches
    ):
        lines.append(
            _C.green(
                f"[OK] All {total_keys} keys matched successfully. "
                f"(missing: 0, unexpected: 0, shape_mismatch: 0)"
            )
        )
    else:
        matched = total_keys - len(result.missing_keys)
        lines.append(
            f"Matched: {matched}/{total_keys} keys | "
            f"Missing: {len(result.missing_keys)} | "
            f"Unexpected: {len(result.unexpected_keys)} | "
            f"Shape mismatch: {len(result.shape_mismatches)}"
        )
        if result.missing_keys:
            lines.append("")
            lines.append(
                _C.yellow(
                    f"[WARNING] Missing keys ({len(result.missing_keys)} total) "
                    f"- model expects but not in checkpoint:"
                )
            )
            lines.extend(_format_key_list(result.missing_keys))
        if result.unexpected_keys:
            lines.append("")
            lines.append(
                _C.yellow(
                    f"[WARNING] Unexpected keys ({len(result.unexpected_keys)} total) "
                    f"- in checkpoint but not used:"
                )
            )
            lines.extend(_format_key_list(result.unexpected_keys))
        if result.shape_mismatches:
            lines.append("")
            lines.append(
                _C.yellow(
                    f"[WARNING] Shape mismatches ({len(result.shape_mismatches)} total):"
                )
            )
            for m in result.shape_mismatches[:_MAX_SHAPE_MISMATCHES]:
                lines.append(
                    f"    {m.key}: ckpt={list(m.src_global_shape)} vs model={list(m.dst_global_shape)}"
                )
            remaining = len(result.shape_mismatches) - _MAX_SHAPE_MISMATCHES
            if remaining > 0:
                lines.append(f"    ... and {remaining} more")

    lines.append(_SEP)
    _emit(lines)


# ---------------------------------------------------------------------------
# Printing: AOA report
# ---------------------------------------------------------------------------


def _print_aoa_report(
    result: KeyValidationResult,
    aoa_mappings: list[AOAMappingEntry],
    explicitly_removed: set[str],
    path: str,
) -> None:
    lines = [
        _SEP,
        f"FlexCheckpoint Load Report (Checkpoint: {path}, AOA enabled)",
        _SEP,
    ]

    # Status
    total_dst = (
        len(aoa_mappings)
        + len(result.missing_keys)
        + len(result.randomly_initialized_keys)
    )
    if not result.missing_keys and not result.unexpected_keys:
        lines.append(
            _C.green(
                f"[OK] All {total_dst} keys resolved via AOA mapping. "
                f"(missing: 0, unexpected: 0)"
            )
        )
    else:
        matched = total_dst - len(result.missing_keys)
        lines.append(
            f"Matched: {matched}/{total_dst} keys | "
            f"Missing: {len(result.missing_keys)} | "
            f"Unexpected: {len(result.unexpected_keys)}"
        )
        if result.missing_keys:
            lines.append("")
            lines.append(
                _C.yellow(
                    f"[WARNING] Missing keys ({len(result.missing_keys)} total) "
                    f"- no AOA source mapping:"
                )
            )
            lines.extend(_format_key_list(result.missing_keys))
        if result.unexpected_keys:
            lines.append("")
            lines.append(
                _C.yellow(
                    f"[WARNING] Unexpected keys ({len(result.unexpected_keys)} total) "
                    f"- in checkpoint but not consumed by any AOA mapping:"
                )
            )
            lines.extend(_format_key_list(result.unexpected_keys))

    # AOA mapping table
    lines.append("")
    lines.append(_C.cyan(_THIN_SEP))

    # Classify mappings
    non_identity = [m for m in aoa_mappings if not m.is_identity]
    rename_only, with_transform, structural = _classify_mappings(non_identity)

    total_dst = len(aoa_mappings)
    total_src = len(
        {sm.src_key for m in aoa_mappings for sm in m.slice_mappings}
    )
    lines.append(
        _C.cyan(f"AOA Key Mapping ({total_dst} dst keys, {total_src} src keys)")
    )
    lines.append(_C.cyan(_THIN_SEP))

    # Summary
    lines.append("Summary:")
    lines.append(
        f"  1-to-1 rename (same shape, no transform):  {len(rename_only)} keys (not shown)"
    )
    lines.append(
        f"  1-to-1 with transform:                     {len(with_transform)} keys "
        f"({min(len(_group_by_signature(with_transform)), _MAX_PATTERNS_SHOWN)} pattern(s) below)"
    )
    lines.append(
        f"  Structural (N-to-1 / 1-to-N / reshape):    {len(structural)} keys "
        f"({min(len(_group_by_signature(structural)), _MAX_PATTERNS_SHOWN)} pattern(s) below)"
    )

    # Print transform patterns
    next_index = 1
    if with_transform:
        lines.append("")
        result_lines, next_index = _format_pattern_groups(
            _group_by_signature(with_transform), "1-to-1 transform", next_index
        )
        lines.extend(result_lines)

    # Print structural patterns
    if structural:
        lines.append("")
        result_lines, next_index = _format_pattern_groups(
            _group_by_signature(structural), "structural", next_index
        )
        lines.extend(result_lines)

    # Removed / Initialized
    lines.append("")
    removed_str = ", ".join(sorted(explicitly_removed)[:5])
    if len(explicitly_removed) > 5:
        removed_str += f" ... +{len(explicitly_removed) - 5} more"
    lines.append(f"Removed ({len(explicitly_removed)}): {removed_str or '-'}")
    init_keys = result.randomly_initialized_keys
    init_str = ", ".join(sorted(init_keys)[:5])
    if len(init_keys) > 5:
        init_str += f" ... +{len(init_keys) - 5} more"
    lines.append(f"Initialized ({len(init_keys)}): {init_str or '-'}")

    lines.append("")
    lines.append(_THIN_SEP)
    lines.append(_SEP)
    _emit(lines)


# ---------------------------------------------------------------------------
# Helpers: Classification & Pattern Merging
# ---------------------------------------------------------------------------


def _classify_mappings(
    non_identity: list[AOAMappingEntry],
) -> tuple[list[AOAMappingEntry], list[AOAMappingEntry], list[AOAMappingEntry]]:
    """Classify non-identity mappings into rename_only, with_transform, structural."""
    rename_only = []
    with_transform = []
    structural = []
    for entry in non_identity:
        if len(entry.slice_mappings) != 1:
            structural.append(entry)
            continue
        sm = entry.slice_mappings[0]
        src_norm = re.sub(r"\d+", "{N}", sm.src_key)
        dst_norm = re.sub(r"\d+", "{N}", entry.dst_key)
        if src_norm != dst_norm:
            structural.append(entry)
        elif sm.postprocess is None:
            rename_only.append(entry)
        else:
            with_transform.append(entry)
    return rename_only, with_transform, structural


def _get_signature(entry: AOAMappingEntry) -> str:
    """Compute a structure signature for pattern grouping."""
    dst_norm = re.sub(r"\d+", "{N}", entry.dst_key)
    parts = [dst_norm, str(len(entry.slice_mappings))]
    for sm in entry.slice_mappings:
        src_norm = re.sub(r"\d+", "{N}", sm.src_key)
        pp = "|".join(sm.postprocess) if sm.postprocess else ""
        parts.append(f"{src_norm}:{pp}")
    return "@@".join(parts)


def _group_by_signature(
    entries: list[AOAMappingEntry],
) -> dict[str, list[AOAMappingEntry]]:
    """Group entries by structure signature."""
    groups: dict[str, list[AOAMappingEntry]] = defaultdict(list)
    for entry in entries:
        groups[_get_signature(entry)].append(entry)
    return groups


def _format_pattern_groups(
    groups: dict[str, list[AOAMappingEntry]], label: str, start_index: int = 1
) -> tuple[list[str], int]:
    """Format grouped patterns with box-drawing style. Returns (lines, next_index)."""
    lines = []
    shown = 0
    idx = start_index
    for _sig, entries in sorted(groups.items(), key=lambda x: -len(x[1])):
        if shown >= _MAX_PATTERNS_SHOWN:
            remaining = len(groups) - shown
            lines.append(f"  ... and {remaining} more {label} pattern(s)")
            break
        shown += 1
        representative = entries[0]
        count = len(entries)

        # Build pattern title
        dst_pattern = re.sub(r"\d+", "*", representative.dst_key)
        lines.append(f"[Pattern #{idx}] {dst_pattern}  ({count} keys, {label})")
        lines.append("\u250c" + "\u2500" * 69)
        # DST line
        shape_str = list(representative.dst_global_shape)
        lines.append(f"\u2502 DST: {representative.dst_key}  {shape_str}")
        # SRC lines (with folding)
        _append_src_lines(lines, representative.slice_mappings)
        # OP line
        ops = _describe_ops(representative)
        if ops:
            lines.append(f"\u2502 OP:  {ops}")
        lines.append("\u2514" + "\u2500" * 69)
        lines.append("")
        idx += 1
    return lines, idx


def _append_src_lines(
    lines: list[str], slice_mappings: list[AOASliceMapping]
) -> None:
    """Append SRC lines, folding consecutive numeric patterns."""
    if len(slice_mappings) <= _SRC_FOLD_THRESHOLD:
        for i, sm in enumerate(slice_mappings):
            prefix = "\u2502 SRC:" if i == 0 else "\u2502    +"
            slice_info = _format_slice_range(sm.src_slice, sm.dst_slice)
            lines.append(f"{prefix} {sm.src_key}{slice_info}")
        return

    # Try to fold: find common pattern
    src_keys = [sm.src_key for sm in slice_mappings]
    folded = _try_fold_src_keys(src_keys)
    if folded:
        lines.append(f"\u2502 SRC: {folded}  (\u00d7{len(slice_mappings)})")
    else:
        # Show first 2 and last 1
        lines.append(f"\u2502 SRC: {src_keys[0]}")
        lines.append(f"\u2502    + {src_keys[1]}")
        lines.append(f"\u2502    + ... ({len(src_keys) - 3} more)")
        lines.append(f"\u2502    + {src_keys[-1]}")


def _format_slice_range(
    src_slice: tuple[slice, ...], dst_slice: tuple[slice, ...]
) -> str:
    """Format slice info when same src_key appears multiple times."""
    src_str = ",".join(f"{s.start}:{s.stop}" for s in src_slice)
    dst_str = ",".join(f"{s.start}:{s.stop}" for s in dst_slice)
    return f"  [{src_str}] -> dst[{dst_str}]"


def _try_fold_src_keys(keys: list[str]) -> str | None:
    """Try to fold src keys like experts.0, experts.1, ..., experts.255 into a pattern."""
    if len(keys) < 2:
        return None
    # Find varying digit segments
    pattern = re.sub(r"\d+", "{}", keys[0])
    for k in keys[1:]:
        if re.sub(r"\d+", "{}", k) != pattern:
            return None
    # Extract the varying numbers
    nums_per_key = [re.findall(r"\d+", k) for k in keys]
    num_positions = len(nums_per_key[0])
    # Find which position varies
    varying_pos = []
    for pos in range(num_positions):
        vals = [int(n[pos]) for n in nums_per_key]
        if len(set(vals)) > 1:
            varying_pos.append(pos)
    if len(varying_pos) != 1:
        return None
    vpos = varying_pos[0]
    vals = [int(n[vpos]) for n in nums_per_key]
    lo, hi = min(vals), max(vals)
    # Reconstruct pattern with {lo..hi}
    segments = re.split(r"\d+", keys[0])
    digits = re.findall(r"\d+", keys[0])
    result_parts = []
    for i, seg in enumerate(segments):
        result_parts.append(seg)
        if i < len(digits):
            if i == vpos:
                result_parts.append(f"{{{lo}..{hi}}}")
            else:
                result_parts.append(digits[i])
    return "".join(result_parts)


def _describe_ops(entry: AOAMappingEntry) -> str:
    """Describe the operations for a mapping entry."""
    ops = []
    if len(entry.slice_mappings) > 1:
        ops.append("concat")
    # Collect postprocess from first slice (representative)
    if entry.slice_mappings:
        pp = entry.slice_mappings[0].postprocess
        if pp:
            for p in pp:
                if p.startswith("["):
                    ops.append(f"permute({p})")
                else:
                    ops.append(f"cast({p})")
    return " + ".join(ops)


# ---------------------------------------------------------------------------
# Helpers: Key list formatting
# ---------------------------------------------------------------------------


def _format_key_list(keys: set[str]) -> list[str]:
    """Format a set of keys with prefix grouping and truncation."""
    if not keys:
        return []
    sorted_keys = sorted(keys)
    if len(sorted_keys) <= _MAX_KEYS_SHOWN:
        return [f"    {k}" for k in sorted_keys]

    # Adaptive grouping: find the prefix depth that gives reasonable group sizes
    groups = _group_keys_adaptive(sorted_keys)

    lines = []
    groups_shown = 0
    for prefix, group_keys in sorted(groups.items(), key=lambda x: -len(x[1])):
        if groups_shown >= _MAX_KEYS_SHOWN:
            remaining_groups = len(groups) - groups_shown
            remaining_keys = sum(
                len(v)
                for i, v in enumerate(
                    sorted(groups.values(), key=len, reverse=True)
                )
                if i >= groups_shown
            )
            lines.append(
                f"    ... and {remaining_groups} more groups ({remaining_keys} keys)"
            )
            break
        groups_shown += 1
        if len(group_keys) > 3:
            lines.append(f"    [{prefix}] ({len(group_keys)} keys):")
            for k in group_keys[:3]:
                lines.append(f"      {k}")
            lines.append(f"      ... +{len(group_keys) - 3} more")
        else:
            for k in group_keys:
                lines.append(f"    {k}")
    return lines


def _group_keys_adaptive(keys: list[str]) -> dict[str, list[str]]:
    """Group keys by normalized pattern (digits replaced with *)."""
    groups: dict[str, list[str]] = defaultdict(list)
    for k in keys:
        # Replace all digit segments with * to get the pattern
        pattern = re.sub(r"(?<=\.)\d+(?=\.)|(?<=\.)\d+$", "*", k)
        groups[pattern].append(k)
    return dict(groups)


# ---------------------------------------------------------------------------
# Helpers: Output
# ---------------------------------------------------------------------------


def _emit(lines: list[str]) -> None:
    """Output lines via logger, respecting total line limit."""
    for i, line in enumerate(lines):
        if i >= _MAX_TOTAL_LINES:
            logger.info(
                f"... output truncated ({len(lines) - i} lines omitted)"
            )
            break
        logger.info(line)