import json
import math
from pathlib import Path
from typing import Union

import mlx.core as mx
import mlx.nn as nn
from mlx.utils import tree_flatten

from .lora import LoRaLayer

DEFAULT_LORA_NUM_LAYERS = -1


class Colors:
    HEADER = "\033[95m"
    OKBLUE = "\033[94m"
    OKCYAN = "\033[96m"
    OKGREEN = "\033[92m"
    WARNING = "\033[93m"
    FAIL = "\033[91m"
    ENDC = "\033[0m"
    BOLD = "\033[1m"
    UNDERLINE = "\033[4m"


not_supported_for_training = {"gemma3n", "qwen3_omni"}


def grad_checkpoint(layer):
    """
    Update all instances of type(layer) to use gradient checkpointing.
    """
    fn = type(layer).__call__

    def checkpointed_fn(model, *args, **kwargs):
        def inner_fn(params, *args, **kwargs):
            model.update(params)
            return fn(model, *args, **kwargs)

        return mx.checkpoint(inner_fn)(model.trainable_parameters(), *args, **kwargs)

    type(layer).__call__ = checkpointed_fn


def get_learning_rate(
    iters: int,
    step: int,
    warmup_steps: int,
    learning_rate: float,
    min_learning_rate: float,
):
    if step < warmup_steps:
        return learning_rate * (step / warmup_steps)

    progress = (step - warmup_steps) / (iters - warmup_steps)
    cosine_decay = 0.5 * (1 + math.cos(math.pi * progress))
    return min_learning_rate + (learning_rate - min_learning_rate) * cosine_decay


def get_module_by_name(model, name):
    parts = name.split(".")
    module = model
    for part in parts:
        if part.isdigit():
            module = module[int(part)]
        else:
            module = getattr(module, part)
    return module


def set_module_by_name(model, name, new_module):
    parts = name.split(".")
    module = model
    for part in parts[:-1]:
        if part.isdigit():
            module = module[int(part)]
        else:
            module = getattr(module, part)
    if parts[-1].isdigit():
        module[int(parts[-1])] = new_module
    else:
        setattr(module, parts[-1], new_module)


def _lora_scale(alpha: float, rank: int) -> float:
    return alpha / rank


def _linear_layer_key(model_key: str, language_key: str) -> str:
    return f"{model_key}.{language_key}" if language_key else model_key


def _to_lora(layer, lora_parameters, use_dora=False):
    if isinstance(layer, (nn.Linear, nn.QuantizedLinear)):
        if use_dora:
            from mlx_lm.tuner.dora import DoRALinear

            return DoRALinear.from_base(
                layer,
                r=lora_parameters["rank"],
                scale=lora_parameters["scale"],
                dropout=lora_parameters["dropout"],
            )

        from mlx_lm.tuner.lora import LoRALinear

        return LoRALinear.from_base(
            layer,
            r=lora_parameters["rank"],
            scale=lora_parameters["scale"],
            dropout=lora_parameters["dropout"],
        )

    raise ValueError(f"Can't convert layer of type {type(layer).__name__} to LoRA")


def _apply_language_lora_layers(model, linear_layers, lora_parameters):
    targets = set(linear_layers)
    keys = []
    for name, module in model.language_model.named_modules():
        if isinstance(module, (nn.Linear, nn.QuantizedLinear)):
            if name.split(".")[-1] in targets:
                set_module_by_name(
                    model.language_model,
                    name,
                    _to_lora(module, lora_parameters),
                )
                keys.append(_linear_layer_key("language_model", name))
    return keys


def _apply_lora_layers(model, config):
    lora_parameters = dict(config["lora_parameters"])
    fine_tune_type = config.get("fine_tune_type", "lora")
    use_dora = fine_tune_type == "dora"

    if fine_tune_type == "full":
        return model

    if "keys" not in lora_parameters:
        from mlx_lm.tuner.utils import linear_to_lora_layers

        linear_to_lora_layers(
            model,
            config.get("num_layers", DEFAULT_LORA_NUM_LAYERS),
            lora_parameters,
            use_dora=use_dora,
        )
        return model

    for name in lora_parameters["keys"]:
        module = get_module_by_name(model, name)
        set_module_by_name(
            model,
            name,
            _to_lora(module, lora_parameters, use_dora=use_dora),
        )
    return model


def _lora_config(rank, alpha, dropout):
    return {
        "fine_tune_type": "lora",
        "num_layers": DEFAULT_LORA_NUM_LAYERS,
        "lora_parameters": {
            "rank": rank,
            "dropout": dropout,
            "scale": _lora_scale(alpha, rank),
        },
    }


def _apply_legacy_lora_layers(model, config):
    list_of_modules = find_all_linear_names(model.language_model)
    return get_peft_model(
        model,
        list_of_modules,
        rank=config["rank"],
        alpha=config.get("alpha", 0.1),
        dropout=config.get("dropout", 0.1),
        legacy=True,
    )


def get_peft_model(
    model,
    linear_layers,
    rank=10,
    alpha=0.1,
    dropout=0.1,
    freeze=True,
    verbose=True,
    legacy=False,
):
    if freeze:
        freeze_model(model)

    if legacy:
        for name, module in model.language_model.named_modules():
            if isinstance(module, nn.Linear) or isinstance(module, nn.QuantizedLinear):
                if name.split(".")[-1] in linear_layers:
                    lora_layer = LoRaLayer(module, rank, alpha, dropout)
                    set_module_by_name(model.language_model, name, lora_layer)

        model.config.lora = {}
        model.config.lora["rank"] = rank
        model.config.lora["alpha"] = alpha
        model.config.lora["dropout"] = dropout
    else:
        config = _lora_config(rank, alpha, dropout)
        config["lora_parameters"]["keys"] = _apply_language_lora_layers(
            model, linear_layers, config["lora_parameters"]
        )
        model.config.lora = config

    if verbose:
        print_trainable_parameters(model.language_model)

    return model


def freeze_model(model):
    top_level_to_freeze = {
        "language_model",
        "vision_model",
        "vision_tower",
        "aligner",
        "connector",
        "multi_modal_projector",
        "mm_projector",
        "audio_tower",
        "embed_audio",
        "embed_vision",
    }
    for name, module in model.named_modules():
        name = name.split(".")[0]
        if name in top_level_to_freeze and hasattr(model, name):
            try:
                model[f"{name}"].freeze()
            except Exception:
                # Fallback for towers whose .freeze() errors on non-Module
                # sub-objects (e.g. Gemma 4 audio_tower).
                try:
                    from mlx.utils import tree_flatten

                    top = model[f"{name}"]
                    leaves = tree_flatten(
                        top.leaf_modules(), is_leaf=lambda m: isinstance(m, nn.Module)
                    )
                    for _, m in leaves:
                        m.freeze(recurse=False)
                except Exception:
                    pass


def find_all_linear_names(model):
    cls = nn.Linear
    quantized_cls = nn.QuantizedLinear
    lora_module_names = set()
    multimodal_keywords = [
        "mm_projector",
        "vision_tower",
        "vision_resampler",
        "aligner",
    ]
    for name, module in model.named_modules():
        if any(mm_keyword in name for mm_keyword in multimodal_keywords):
            continue
        if isinstance(module, cls) or isinstance(module, quantized_cls):
            names = name.split(".")
            lora_module_names.add(names[0] if len(names) == 1 else names[-1])

    if "lm_head" in lora_module_names:  # needed for 16-bit
        lora_module_names.remove("lm_head")
    return list(lora_module_names)


def count_parameters(model):
    def nparams(m):
        if isinstance(m, (nn.QuantizedLinear, nn.QuantizedEmbedding)):
            return m.weight.size * (32 // m.bits)
        return sum(v.size for _, v in tree_flatten(m.parameters()))

    leaf_modules = tree_flatten(
        model.leaf_modules(), is_leaf=lambda m: isinstance(m, nn.Module)
    )
    total_p = sum(nparams(m) for _, m in leaf_modules) / 10**6

    return total_p


def print_trainable_parameters(model):
    def nparams(m):
        if isinstance(m, (nn.QuantizedLinear, nn.QuantizedEmbedding)):
            return m.weight.size * (32 // m.bits)
        return sum(v.size for _, v in tree_flatten(m.parameters()))

    leaf_modules = tree_flatten(
        model.leaf_modules(), is_leaf=lambda m: isinstance(m, nn.Module)
    )
    total_p = sum(nparams(m) for _, m in leaf_modules) / 10**6
    trainable_p = (
        sum(v.size for _, v in tree_flatten(model.trainable_parameters())) / 10**6
    )

    print(
        f"#trainable params: {trainable_p} M || all params: {total_p} M || trainable%: {(trainable_p * 100 / total_p):.3f}%"
    )


def apply_lora_layers(model: nn.Module, adapter_path: str) -> nn.Module:
    """
    Apply LoRA layers to the model.

    Args:
        model (nn.Module): The neural network model.
        adapter_path (str): Path to the adapter configuration file.

    Returns:
        nn.Module: The updated model with LoRA layers applied.
    """
    if getattr(model, "_is_text_model", False):
        from mlx_lm.utils import load_adapters

        model.language_model._model = load_adapters(
            model.language_model._model, adapter_path
        )
        return model

    adapter_path = Path(adapter_path)

    if not adapter_path.exists():
        raise FileNotFoundError(f"The adapter path does not exist: {adapter_path}")

    with open(adapter_path / "adapter_config.json", "r") as f:
        config = json.load(f)
        if "rank" not in config and "lora_parameters" not in config:
            raise ValueError("The adapter does not have lora params in the config")

    if "lora_parameters" in config:
        model = _apply_lora_layers(model, config)
    else:
        model = _apply_legacy_lora_layers(model, config)

    model.load_weights(str(adapter_path / "adapters.safetensors"), strict=False)

    return model


def unfreeze_modules(model: nn.Module, module_names):
    """Unfreeze modules whose qualified names match any of the given patterns.

    This scans model.named_modules() so nested components like
    "vision_tower.layers.0" are handled as well.
    """
    targets = set(module_names)
    found = set()
    for full_name, sub in model.named_modules():
        top = full_name.split(".")[0] if full_name else ""
        if any((name == top) or (name in full_name) for name in targets):
            if hasattr(sub, "unfreeze"):
                sub.unfreeze()
                found.add(top or full_name)
    if not found:
        print(
            "[warn] unfreeze_modules: no matching modules found for patterns:",
            ", ".join(module_names),
        )


def save_adapter(model: nn.Module, adapter_file: Union[str, Path]):
    """Save adapter weights and config."""
    path = Path(adapter_file)
    path.parent.mkdir(parents=True, exist_ok=True)

    # Save adapter config if available
    if hasattr(model, "config") and hasattr(model.config, "lora"):
        with open(path.parent / "adapter_config.json", "w") as f:
            json.dump(model.config.lora, f, indent=2)

    # Save weights
    flattened_tree = tree_flatten(model.trainable_parameters())
    mx.save_safetensors(str(adapter_file), dict(flattened_tree))