from pathlib import Path
from typing import TYPE_CHECKING

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

from mflux.dreambooth.lora_layers.linear_lora_layer import LoRALinear
from mflux.dreambooth.state.training_spec import SingleTransformerBlocks, TrainingSpec, TransformerBlocks
from mflux.dreambooth.state.zip_util import ZipUtil
from mflux.models.transformer.joint_transformer_block import JointTransformerBlock
from mflux.models.transformer.single_transformer_block import SingleTransformerBlock
from mflux.utils.version_util import VersionUtil
from mflux.weights.weight_handler import MetaData, WeightHandler

if TYPE_CHECKING:
    from mflux.flux.flux import Flux1


class LoRALayers:
    def __init__(self, weights: "WeightHandler"):
        self.layers = weights

    @staticmethod
    def from_spec(flux: "Flux1", training_spec: TrainingSpec) -> "LoRALayers":
        if training_spec.lora_layers.state_path is not None:
            # Load from state if present in the spec
            from mflux.weights.weight_handler_lora import WeightHandlerLoRA

            weights = ZipUtil.unzip(
                zip_path=training_spec.checkpoint_path,
                filename=training_spec.lora_layers.state_path,
                loader=lambda x: WeightHandlerLoRA.load_lora_weights(
                    transformer=flux.transformer, lora_files=[x], lora_scales=[1.0]
                ),
            )
            return LoRALayers(weights=weights[0])
        else:
            # Construct the LoRA weights from the spec
            transformer_lora_layers = {}
            single_transformer_lora_layers = {}

            if training_spec.lora_layers.transformer_blocks:
                transformer_lora_layers = LoRALayers._construct_layers(
                    blocks=flux.transformer.transformer_blocks,
                    block_spec=training_spec.lora_layers.transformer_blocks,
                    block_prefix="transformer.transformer_blocks",
                )

            if training_spec.lora_layers.single_transformer_blocks:
                single_transformer_lora_layers = LoRALayers._construct_layers(
                    blocks=flux.transformer.single_transformer_blocks,
                    block_spec=training_spec.lora_layers.single_transformer_blocks,
                    block_prefix="transformer.single_transformer_blocks",
                )

            lora_layers = {**transformer_lora_layers, **single_transformer_lora_layers}

            weights = WeightHandler(
                meta_data=MetaData(mflux_version=VersionUtil.get_mflux_version()),
                transformer=mlx.utils.tree_unflatten(list(lora_layers.items()))["transformer"],
            )

            return LoRALayers(weights=weights)

    @staticmethod
    def _construct_layers(
        block_spec: TransformerBlocks | SingleTransformerBlocks,
        blocks: list[JointTransformerBlock] | list[SingleTransformerBlock],
        block_prefix: str,
    ) -> dict:
        block_indices = block_spec.block_range.get_blocks()
        lora_layers = {}
        for idx in block_indices:
            if idx >= len(blocks):
                raise IndexError(f"Index {idx} over range")

            block = blocks[idx]
            for layer_type in block_spec.layer_types:
                original_layer = LoRALayers._get_nested_attr(block, layer_type)
                is_list = isinstance(original_layer, list)

                lora_layer = LoRALinear.from_linear(
                    linear=original_layer[0] if is_list else original_layer,
                    r=block_spec.lora_rank,
                )
                layer_path = f"{block_prefix}.{idx}.{layer_type}"

                lora_layers[layer_path] = [lora_layer] if is_list else lora_layer

        return lora_layers

    @staticmethod
    def transformer_dict_from_template(weights: dict, transformer: nn.Module, scale: float) -> dict:
        lora_layers = {}
        for key in weights.keys():
            if key.endswith(".lora_A"):
                base_path = key[: -len(".lora_A")]
                parts = base_path.split(".")

                if parts[1] == "transformer_blocks":
                    LoRALayers._handle_transformer_blocks(
                        weights=weights,
                        scale=scale,
                        transformer=transformer,
                        lora_layers=lora_layers,
                        base_path=base_path,
                    )

                elif parts[1] == "single_transformer_blocks":
                    LoRALayers._handle_single_transformer_blocks(
                        weights=weights,
                        scale=scale,
                        transformer=transformer,
                        lora_layers=lora_layers,
                        base_path=base_path,
                    )

                # Handle top-level transformer components like x_embedder, context_embedder, proj_out
                elif len(parts) == 2 and parts[0] == "transformer":
                    LoRALayers._handle_top_level_component(
                        weights=weights,
                        scale=scale,
                        transformer=transformer,
                        lora_layers=lora_layers,
                        base_path=base_path,
                        component_name=parts[1],
                    )

        return lora_layers

    @staticmethod
    def _resolve_legacy_paths(module, module_name: str, attr_name: str, parts: list, block_idx: int) -> tuple:
        # Handle legacy naming: ff.net.2 -> ff.linear2
        if module_name == "ff" and attr_name == "net" and len(parts) >= 6 and parts[5] == "2":
            return getattr(module, "linear2"), f"transformer.transformer_blocks.{block_idx}.ff.linear2"

        # Handle attn.to_out.0 -> attn.to_out[0] (list access)
        if module_name == "attn" and attr_name == "to_out" and len(parts) >= 6 and parts[5] == "0":
            return module.to_out[0], f"transformer.transformer_blocks.{block_idx}.attn.to_out"

        # Default case
        original_layer = getattr(module, attr_name)
        if len(parts) == 6:
            original_layer = original_layer[0]
        return original_layer, ".".join(parts)

    @staticmethod
    def _create_lora_layer(weights: dict, base_path: str, original_layer, scale: float):
        lora_A = weights[f"{base_path}.lora_A"]
        rank = lora_A.shape[1]
        lora_layer = LoRALinear.from_linear(linear=original_layer, r=rank, scale=scale)
        lora_layer.lora_A = lora_A
        lora_layer.lora_B = weights[f"{base_path}.lora_B"]
        return lora_layer

    @staticmethod
    def _handle_transformer_blocks(
        weights: dict, scale: float, transformer: nn.Module, lora_layers: dict, base_path: str
    ):
        parts = base_path.split(".")
        block_idx = int(parts[2])
        module_name = parts[3]
        attr_name = parts[4]
        block = transformer.transformer_blocks[block_idx]
        module = getattr(block, module_name)

        # Resolve legacy naming and get the actual layer and storage path
        original_layer, storage_path = LoRALayers._resolve_legacy_paths(
            module=module,
            module_name=module_name,
            attr_name=attr_name,
            parts=parts,
            block_idx=block_idx,
        )

        # Create and store LoRA layer
        lora_layer = LoRALayers._create_lora_layer(weights, base_path, original_layer, scale)
        lora_layers[storage_path] = lora_layer

    @staticmethod
    def _handle_single_transformer_blocks(
        weights: dict, scale: float, transformer: nn.Module, lora_layers: dict, base_path: str
    ):
        parts = base_path.split(".")
        block_idx = int(parts[2])
        module_name = parts[3]

        if len(parts) == 4:
            original_layer = getattr(transformer.single_transformer_blocks[block_idx], module_name)
        else:
            attr_name = parts[4]
            block = transformer.single_transformer_blocks[block_idx]
            module = getattr(block, module_name)
            original_layer = getattr(module, attr_name)

        # Create and store LoRA layer
        lora_layer = LoRALayers._create_lora_layer(weights, base_path, original_layer, scale)
        lora_layers[base_path] = lora_layer

    @staticmethod
    def _handle_top_level_component(
        weights: dict, scale: float, transformer: nn.Module, lora_layers: dict, base_path: str, component_name: str
    ):
        original_layer = getattr(transformer, component_name)
        lora_layer = LoRALayers._create_lora_layer(weights, base_path, original_layer, scale)
        lora_layers[base_path] = lora_layer

    @staticmethod
    def set_transformer_block(transformer_block, dictionary: dict):
        for key, val in dictionary.items():
            if key == "attn":
                LoRALayers._set_attribute(transformer_block, key, val, "to_q")
                LoRALayers._set_attribute(transformer_block, key, val, "to_k")
                LoRALayers._set_attribute(transformer_block, key, val, "to_v")
                LoRALayers._set_attribute(transformer_block, key, val, "to_out")
                LoRALayers._set_attribute(transformer_block, key, val, "add_q_proj")
                LoRALayers._set_attribute(transformer_block, key, val, "add_k_proj")
                LoRALayers._set_attribute(transformer_block, key, val, "add_v_proj")
                LoRALayers._set_attribute(transformer_block, key, val, "to_add_out")
            elif key == "ff" or key == "ff_context":
                LoRALayers._set_attribute(transformer_block, key, val, "linear1")
                LoRALayers._set_attribute(transformer_block, key, val, "linear2")
            elif key == "norm1" or key == "norm1_context":
                LoRALayers._set_attribute(transformer_block, key, val, "linear")
            else:
                raise Exception("Could not set LoRA weights")

    @staticmethod
    def set_single_transformer_block(single_transformer_block, dictionary: dict):
        for key, val in dictionary.items():
            if key == "attn":
                LoRALayers._set_attribute(single_transformer_block, key, val, "to_q")
                LoRALayers._set_attribute(single_transformer_block, key, val, "to_k")
                LoRALayers._set_attribute(single_transformer_block, key, val, "to_v")
            elif key == "norm":
                LoRALayers._set_attribute(single_transformer_block, key, val, "linear")
            elif key == "proj_mlp" or key == "proj_out":
                single_transformer_block[key] = val
            else:
                raise Exception("Could not set LoRA weights")

    @staticmethod
    def _set_attribute(block, key: str, val: dict, name: str):
        if block[key].get(name, False) and val.get(name, False):
            if name == "to_out" and isinstance(block[key][name], list):
                if len(block[key][name]) > 0:
                    lora_layer = val[name][0] if isinstance(val[name], list) else val[name]
                    block[key][name][0] = lora_layer
            else:
                block[key][name] = val[name]

    @staticmethod
    def _get_nested_attr(obj, attr_path):
        attrs = attr_path.split(".")
        for attr in attrs:
            obj = getattr(obj, attr)
        return obj

    def save(self, path: Path, training_spec: TrainingSpec) -> None:
        weights = {}
        for entry in tree_flatten(self.layers.transformer):
            name = entry[0]
            weight = entry[1]
            if name.endswith(".lora_A") or name.endswith(".lora_B"):
                weights[name] = weight

        weights = {key: mx.transpose(val) for key, val in weights.items()}
        weights = {"transformer": weights}
        mx.save_safetensors(
            str(path),
            dict(tree_flatten(weights)),
            metadata={
                "mflux_version": VersionUtil.get_mflux_version(),
                "transformer_blocks": str(training_spec.lora_layers.transformer_blocks),
                "single_transformer_blocks": str(training_spec.lora_layers.single_transformer_blocks),
            },
        )
