"""Download a model checkpoint, load its state.

Supports optional pretrained embeddings, and save the pretrained model to a directory.
"""

import argparse
import json
import os
from collections import OrderedDict
from logging import INFO
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import TYPE_CHECKING, Any, cast

import torch
from composer.loggers import RemoteUploaderDownloader
from composer.utils.checkpoint import _torch_load_with_validation  # noqa: PLC2701
from flwr.common import log
from llmfoundry.utils.builders import (
    build_composer_model,
    build_tokenizer,
)
from llmfoundry.utils.config_utils import (
    process_init_device,
)
from omegaconf import DictConfig

from mt_dao.conf.constants import ModelConfig, Tokenizers
from mt_dao.constants import NUMPY_ARRAY_SUFFIX, PYTORCH_SUFFIX
from mt_dao.file_utils import (
    create_remote_up_down,
    download_file_from_s3,
    load_model_states_from_file,
)
from mt_dao.initialization_utils.helper import (
    retrieve_checkpoint_states,
    retrieve_model_states_and_types,
)
from mt_dao.masks_utils import ModelStateNames
from mt_dao.moe_utils.utils import (
    merge_experts_layers_tensors,
    split_experts_layers_tensors,
)
from mt_dao.pytorch_utils.model_parameters_getters import get_tensors_dict
from mt_dao.pytorch_utils.model_parameters_setters import set_model_tensors_dict
from mt_dao.utils import create_tensors_dict, set_frozen_layers_tensors_dict

if TYPE_CHECKING:
    from transformers.modeling_utils import PreTrainedModel

    from mt_dao.conf.base_schema import S3CommConfig


def load_model_state(  # noqa: PLR0913, PLR0917
    model_state: dict[str, torch.Tensor],
    local_path: Path,
    parameter_names: list[str],
    temp_dir: str,
    up_down: RemoteUploaderDownloader,
    wte_checkpoint: str | None = None,
) -> dict[str, torch.Tensor]:
    """Load the model state from a local file or a checkpoint.

    Parameters
    ----------
    model_state : dict[str, torch.Tensor]
        The initial model state.
    local_path : Path
        The local path to the checkpoint.
    parameter_names : list[str]
        The names of the parameters.
    temp_dir : str
        The temporary directory.
    up_down : RemoteUploaderDownloader
        The remote uploader and downloader.
    wte_checkpoint : str, optional
        The checkpoint to load the word token embedding from, by default None.

    Returns
    -------
    dict[str, torch.Tensor]
        The loaded model state.

    """
    if local_path.suffix == PYTORCH_SUFFIX:
        model_state = _torch_load_with_validation(
            local_path,
            map_location="cpu",
        )["state"]["model"]
    elif local_path.suffix == NUMPY_ARRAY_SUFFIX:
        server_model = load_model_states_from_file(local_path)
        for name, array in zip(
            parameter_names,
            server_model,
            strict=False,
        ):
            model_state[name] = torch.tensor(
                array,
                dtype=model_state[name].dtype,
                device=model_state[name].device,
            )

    if wte_checkpoint is not None:
        wte_temp_dir_path = Path(temp_dir) / wte_checkpoint
        wte_local_path = wte_temp_dir_path
        download_file_from_s3(
            up_down,
            wte_checkpoint,
            wte_local_path,
        )
        if wte_local_path.suffix == PYTORCH_SUFFIX:
            wte_state: dict[str, torch.Tensor] = _torch_load_with_validation(
                wte_local_path,
                map_location="cpu",
            )["state"]["model"]
            # For any key containing "wte" in the name, load from wte_state
            for key in model_state:
                if "wte" in key:
                    model_state[key] = wte_state[key]
                    log(INFO, f"Loaded {key} from {wte_local_path}")
        else:
            wte_server_model = load_model_states_from_file(wte_local_path)
            for name in parameter_names:
                if "wte" in name:
                    model_state[name] = torch.tensor(
                        wte_server_model[0],
                        dtype=model_state[name].dtype,
                        device=model_state[name].device,
                    )
                    log(INFO, f"Loaded {name} from {wte_local_path}")
                    break
    return model_state


def download_and_save_pretrained_model(  # noqa: PLR0913, PLR0914, PLR0915
    checkpoints: tuple[str, str | None],
    safe_tensors_dir: str,
    checkpoints_dir: str,
    model_type: str,
    tokenizer: str,
    *,
    random_model: bool = False,
    run_uuid: str | None = None,
    desired_steps: int | None = None,
    n_experts: int | None = None,
) -> None:
    """Build the model, download its checkpoint, save.

    Parameters
    ----------
    checkpoints : tuple[str, str | None]
        The checkpoint to download and
        the checkpoint to load the word token embedding from.
    safe_tensors_dir : str
        The directory in which to save the model's safetensors.
    checkpoints_dir : str
        The directory in which the checkpoints are stored.
    model_type : str
        The type of the model.
    tokenizer : str
        The type of the tokenizer.
    random_model : bool, optional
        Whether to request a random model, by default False.
    run_uuid : str | None, optional
        The run UUID for when using the new checkpoint, by default None.
    desired_steps : int | None, optional
        The desired steps for the new checkpoint, by default None.
    n_experts : int | None, optional
        The number of experts for the new checkpoint, by default None.

    """
    checkpoint, wte_checkpoint = checkpoints
    # Build remote up_down client.
    up_down = create_remote_up_down(
        bucket_name=checkpoints_dir,
        prefix="",
        run_uuid="download",
        num_attempts=3,
        client_config={"connect_timeout": 3600, "read_timeout": 3600},
    )

    # Build model.
    model_config: dict[str, Any] = dict(ModelConfig[model_type].value)
    tokenizer_conf: dict[str, Any] = dict(Tokenizers[tokenizer].value)

    model = build_composer_model(
        name=str(model_config["name"]),
        cfg=model_config,
        tokenizer=build_tokenizer(
            str(tokenizer_conf["name"]),
            tokenizer_conf["kwargs"],
        ),
        init_context=process_init_device(model_config, None),
        master_weights_dtype=None,
    )

    if not random_model:
        with TemporaryDirectory() as general_temp_dir:
            os.environ["TMPDIR"] = general_temp_dir

            if run_uuid is None or desired_steps is None or n_experts is None:
                with TemporaryDirectory() as temp_dir:
                    log(INFO, "Temp dir: %s", temp_dir)

                    temp_dir_path = Path(temp_dir) / checkpoint
                    local_path = temp_dir_path

                    # The below 'download_and_save_model' routine is inlined here:
                    model_state = model.state_dict()
                    parameter_names = sorted(model_state.keys())

                    # Download the checkpoint.
                    download_file_from_s3(
                        up_down,
                        checkpoint,
                        local_path,
                    )

                    # Load state from the checkpoint.
                    new_model_state = load_model_state(
                        model_state=model_state,
                        local_path=local_path,
                        parameter_names=parameter_names,
                        temp_dir=temp_dir,
                        up_down=up_down,
                        wte_checkpoint=wte_checkpoint,
                    )
                    model.load_state_dict(new_model_state)
            else:
                tensors_dict = get_tensors_dict(
                    model,
                    only_requires_grad=True,
                    sort_dict=True,
                )

                tensors_dict = split_experts_layers_tensors(
                    tensors_dict=tensors_dict,
                    n_experts=n_experts,
                )
                param_names = list(
                    tensors_dict.keys(),
                )
                param_names = [
                    param_name.replace("model.", "", 1) for param_name in param_names
                ]

                conf = cast(
                    "S3CommConfig",
                    DictConfig(
                        {
                            "num_attempts": 3,
                            "bucket_name": checkpoints_dir,
                            "backend_kwargs": {
                                "client_config": {
                                    "connect_timeout": 3600,
                                    "read_timeout": 3600,
                                },
                            },
                        },
                    ),
                )
                checkpoint_information = retrieve_checkpoint_states(
                    checkpoint_dir=f"s3://{checkpoints_dir}/{run_uuid}/server/{desired_steps}",
                    s3_comm_config=conf,
                    run_uuid=run_uuid,
                )
                (
                    local_checkpoint_directory,
                    tmp_dir,
                    remote_or_local,
                    checkpoint_suffixes,
                ) = checkpoint_information

                model_states, model_names_and_types = retrieve_model_states_and_types(
                    local_checkpoint_directory=local_checkpoint_directory,
                    checkpoint_suffixes=checkpoint_suffixes,
                    preserve_optimizer_states=True,
                    desired_steps=-1,
                    n_experts=n_experts,
                    param_names=param_names,
                )
                server_model = []

                for state, (_layer_name, layer_type) in zip(
                    model_states,
                    model_names_and_types,
                    strict=True,
                ):
                    if layer_type == ModelStateNames.PARAMETERS:
                        server_model.append(state)

                log(INFO, f"Length of server model: {len(server_model)}")
                log(INFO, f"Length of param names: {len(param_names)}")
                log(INFO, f"Length of initial tensors dict: {len(model_states)}")

                pre_merge_tensors_dict = set_frozen_layers_tensors_dict(
                    create_tensors_dict(
                        parameters=server_model,
                        parameters_names=param_names,
                    ),
                )

                merged_tensors_dict = OrderedDict(
                    merge_experts_layers_tensors(
                        tensors_dict=dict(pre_merge_tensors_dict),
                        n_experts=n_experts,
                    ),
                )

                set_model_tensors_dict(
                    model=model,
                    tensors_dict=merged_tensors_dict,
                )

                if remote_or_local and tmp_dir:
                    tmp_dir.cleanup()

    # 1. Get the Hugging Face model
    hf_model = model.model

    # 2. Get its configuration
    model_config = hf_model.config  # pyright: ignore[reportAssignmentType, reportAttributeAccessIssue]

    # 3. Modify the architectures attribute 💡
    # so we do not need to add anything to huggingface
    model_config.architectures = ["MPTForCausalLM"]  # type: ignore[reportAttributeAccessIssue,attr-defined]

    cast("PreTrainedModel", model.model).save_pretrained(safe_tensors_dir)

    # 1. Define the path to the config file
    config_path = Path(safe_tensors_dir) / "config.json"

    # 2. Load the freshly saved (but incorrect) config
    with config_path.open("r", encoding="utf-8") as f:
        config_data = json.load(f)

    # 3. Overwrite the architecture
    config_data["architectures"] = ["MPTForCausalLM"]

    # 4. Write the corrected data back to the file
    with config_path.open("w", encoding="utf-8") as f:
        json.dump(config_data, f, indent=2)

    # Save tokenizer.
    tokenizer_object = build_tokenizer(
        str(tokenizer_conf["name"]),
        tokenizer_conf["kwargs"],
    )
    tokenizer_object.save_pretrained(safe_tensors_dir)


if __name__ == "__main__":
    default_parser: argparse.ArgumentParser = argparse.ArgumentParser()
    default_parser.add_argument(
        "--checkpoint_dir",
        type=str,
        required=True,
        help="Checkpoint directory",
    )

    default_parser.add_argument(
        "--checkpoint",
        type=str,
        required=False,
        help="Name of checkpoint to download",
    )

    default_parser.add_argument(
        "--run_uuid",
        type=str,
        required=False,
        help="Run UUID for when using the new checkpoint",
    )
    default_parser.add_argument(
        "--desired_steps",
        type=int,
        required=False,
        help="Desired steps for the new checkpoint",
    )
    default_parser.add_argument(
        "--n_experts",
        type=int,
        required=False,
        help="Number of experts for the new checkpoint",
    )
    default_parser.add_argument(
        "--wte_checkpoint",
        type=str,
        required=False,
        help="Checkpoint to load an embedding from",
        default=None,
    )
    default_parser.add_argument(
        "--tokenizer",
        type=str,
        required=False,
        help="Checkpoint to load an embedding from",
        default=None,
    )
    default_parser.add_argument(
        "--model_type",
        type=str,
        required=False,
        help="Checkpoint to load an embedding from",
        default="SMOLLM_1B",
    )
    # add parameter to request a random model
    default_parser.add_argument(
        "--random_model",
        action="store_true",
        help="Request a random model",
    )

    # Add a default directory in which to save the model safetensors
    default_parser.add_argument(
        "--safe_tensors_dir",
        type=str,
        default="safe_tensors_test",
        help="Directory in which to save the model safetensors",
    )
    args = default_parser.parse_args()

    # Print args
    log(INFO, "Augments: %s", args)

    download_and_save_pretrained_model(
        checkpoints=(args.checkpoint, args.wte_checkpoint),
        safe_tensors_dir=args.safe_tensors_dir,
        checkpoints_dir=args.checkpoint_dir,
        model_type=args.model_type,
        tokenizer=args.tokenizer,
        random_model=args.random_model,
        run_uuid=args.run_uuid,
        desired_steps=args.desired_steps,
        n_experts=args.n_experts,
    )
