import argparse
import logging
import os
import re
import shutil
import tempfile
from hashlib import md5
from typing import Iterable, Optional
from urllib.parse import urlparse

import torch
from omegaconf import OmegaConf as om
from tqdm import tqdm

from hf_instella.modeling_instella import InstellaForCausalLM
from instella import ModelConfig, TrainConfig
from instella.checkpoint import build_sharded_checkpointer
from instella.util import _get_s3_client
from hf_instella.configuration_instella import InstellaConfig
from tokenizers import Tokenizer
from transformers.models.gpt_neox.tokenization_gpt_neox_fast import GPTNeoXTokenizerFast

logger = logging.getLogger(__name__)

HF_FILENAMES = {
    "config.json",
    "pytorch_model.bin",
    "special_tokens_map.json",
    "tokenizer_config.json",
    "tokenizer.json",
}


def longest_common_prefix(strs: Iterable[str]) -> str:
    """
    Finds the longest common prefix among a list of strings.
    """
    if not strs:
        return ""

    # Find the shortest string in the list
    shortest_str = min(strs, key=len)

    for i, char in enumerate(shortest_str):
        for other_str in strs:
            if other_str[i] != char:
                return shortest_str[:i]

    return shortest_str


def write_config(checkpoint_dir: str):
    # save config as HF config

    logger.info(f"Loading checkpoint from {checkpoint_dir}")

    config_path = os.path.join(checkpoint_dir, "config.yaml")
    model_config = ModelConfig.load(config_path, key="model")
    config_kwargs = model_config.asdict()
    config_kwargs["use_cache"] = True

    config = InstellaConfig()
    set_kwargs = {}
    for k in config_kwargs:
        if k in config:
            set_kwargs[k] = config_kwargs[k]
    
    set_kwargs["architectures"] =  ["InstellaForCausalLM"]
    set_kwargs["hidden_size"] = config_kwargs["d_model"]
    set_kwargs["intermediate_size"] = config_kwargs["mlp_hidden_size"] // 2
    set_kwargs["max_position_embeddings"] = config_kwargs["max_sequence_length"]
    set_kwargs["num_attention_heads"] = config_kwargs["n_heads"]
    set_kwargs["num_hidden_layers"] = config_kwargs["n_layers"]
    set_kwargs["num_key_value_heads"] = config_kwargs["n_heads"]
    set_kwargs["tie_word_embeddings"] = config_kwargs["weight_tying"]
    set_kwargs["pad_token_id"] = config_kwargs["pad_token_id"]
    set_kwargs["eos_token_id"] = config_kwargs["eos_token_id"]
    set_kwargs["bos_token_id"] = config_kwargs["eos_token_id"]
    set_kwargs["vocab_size"] = config_kwargs.get("embedding_size", config_kwargs["vocab_size"])

    if config_kwargs.get("layer_norm_eps"):
        set_kwargs["rms_norm_eps"] = config_kwargs["layer_norm_eps"]

    if config_kwargs.get("rope_theta"):
        set_kwargs["rope_theta"] = config_kwargs["rope_theta"]

    config = InstellaConfig(**set_kwargs)
    logger.info(f"Saving HF-compatible config to {os.path.join(checkpoint_dir, 'config.json')}")
    print(f"Saving HF-compatible config to {os.path.join(checkpoint_dir, 'config.json')}")
    config.save_pretrained(checkpoint_dir)


def write_model(checkpoint_dir: str, ignore_instella_compatibility: bool = False):
    # For device_map = "auto", etc. the models are loaded in a way that start_prefix is not computed correctly.
    # So, we explicitly store the model with the expected prefix.

    old_model_path = os.path.join(checkpoint_dir, "model.pt")
    new_model_path = os.path.join(checkpoint_dir, "pytorch_model.bin")

    state_dict = torch.load(old_model_path, map_location="cpu")

    # this takes care of the case where the model was saved with a different prefix,
    # typically due to unsharding.
    common_prefix = longest_common_prefix(state_dict.keys())
    new_state_dict = {
        key.replace(common_prefix, f"{InstellaForCausalLM.base_model_prefix}.transformer."): val
        for key, val in state_dict.items()
    }

    new_state_dict = {}

    for k in state_dict:
        if k == "transformer.wte.weight":
            new_state_dict["model.embed_tokens.weight"] = state_dict[k]
        elif "k_norm.weight" in k:
            new_state_dict[k.replace("transformer.blocks", "model.layers").replace("k_norm.weight", "self_attn.k_norm.weight")] = state_dict[k]
        elif "q_norm.weight" in k:
            new_state_dict[k.replace("transformer.blocks", "model.layers").replace("q_norm.weight", "self_attn.q_norm.weight")] = state_dict[k]
        elif "att_proj.weight" in k:
            part1, part2, part3 = torch.chunk(state_dict[k], 3)            
            new_state_dict[k.replace("transformer.blocks", "model.layers").replace("att_proj.weight", "self_attn.q_proj.weight")] = part1
            new_state_dict[k.replace("transformer.blocks", "model.layers").replace("att_proj.weight", "self_attn.k_proj.weight")] = part2
            new_state_dict[k.replace("transformer.blocks", "model.layers").replace("att_proj.weight", "self_attn.v_proj.weight")] = part3
        elif "ff_proj.weight" in k:
            part1, part2 = torch.chunk(state_dict[k], 2)            
            new_state_dict[k.replace("transformer.blocks", "model.layers").replace("ff_proj.weight", "mlp.up_proj.weight")] = part1
            new_state_dict[k.replace("transformer.blocks", "model.layers").replace("ff_proj.weight", "mlp.gate_proj.weight")] = part2
        elif "attn_out.weight" in k:
            new_state_dict[k.replace("transformer.blocks", "model.layers").replace("attn_out.weight", "self_attn.o_proj.weight")] = state_dict[k]
        elif "ff_out.weight" in k and "blocks" in k:
            new_state_dict[k.replace("transformer.blocks", "model.layers").replace("ff_out.weight", "mlp.down_proj.weight")] = state_dict[k]
        elif "attn_norm.weight" in k:
            new_state_dict[k.replace("transformer.blocks", "model.layers").replace("attn_norm.weight", "pre_attention_layernorm.weight")] = state_dict[k]
        elif "ff_norm.weight" in k:
            new_state_dict[k.replace("transformer.blocks", "model.layers").replace("ff_norm.weight", "pre_feedforward_layernorm.weight")] = state_dict[k]
        elif "ff_out.weight" in k and "blocks" not in k:
            new_state_dict[k.replace("transformer", "").replace(".ff_out.weight", "lm_head.weight")] = state_dict[k]

    new_state_dict["model.norm.weight"] = state_dict["transformer.ln_f.weight"]

    torch.save(new_state_dict, new_model_path)

    if ignore_instella_compatibility:
        os.remove(old_model_path)

def write_tokenizer(checkpoint_dir: str, input_tokenizer_path):
    base_tokenizer = Tokenizer.from_file(str(input_tokenizer_path))

    eos_token_id = 0
    pad_token_id = 1

    tokenizer = GPTNeoXTokenizerFast(
        tokenizer_object=base_tokenizer,
        eos_token=base_tokenizer.decode([eos_token_id], skip_special_tokens=False),
        pad_token=base_tokenizer.decode([pad_token_id], skip_special_tokens=False),
    )
    
    tokenizer.pad_token_id = pad_token_id
    tokenizer.eos_token_id = eos_token_id
    tokenizer.bos_token_id = eos_token_id
    tokenizer.chat_template =  "{{ bos_token }}{% for message in messages %}{% if message['role'] == 'system' %}{{ '<|system|>\n' + message['content'] + '\n' }}{% elif message['role'] == 'user' %}{{ '<|user|>\n' + message['content'] + '\n' }}{% elif message['role'] == 'assistant' %}{% if not loop.last %}{{ '<|assistant|>\n'  + message['content'] + eos_token + '\n' }}{% else %}{{ '<|assistant|>\n'  + message['content'] + eos_token }}{% endif %}{% endif %}{% if loop.last and add_generation_prompt %}{{ '<|assistant|>\n' }}{% endif %}{% endfor %}" 

    
    tokenizer.save_pretrained(checkpoint_dir)

def convert_checkpoint(checkpoint_dir: str, ignore_instella_compatibility: bool = False, input_tokenizer_path = None):
    print("Converting checkpoint to HF format...")
    write_config(checkpoint_dir)

    print("Saving model to checkpoint...")
    write_model(checkpoint_dir, ignore_instella_compatibility=ignore_instella_compatibility)

    print("Saving tokenizer to checkpoint...")
    write_tokenizer(checkpoint_dir, input_tokenizer_path)

    # Cannot remove it before writing the tokenizer
    if ignore_instella_compatibility:
        os.remove(os.path.join(checkpoint_dir, "config.yaml"))


def validate_tokenizer(checkpoint_dir: str, tokenizer_name_or_path: Optional[str] = None, fix_tokenizer=False):
    path = os.path.join(checkpoint_dir, "config.yaml")
    conf = om.load(path)

    print("Saving tokenizer to checkpoint...")

    tokenizer_name_or_path = str(tokenizer_name_or_path or conf["tokenizer"]["identifier"])  # pyright: ignore

    try:
        if os.path.exists(tokenizer_name_or_path):
            Tokenizer.from_file(tokenizer_name_or_path)
        else:
            Tokenizer.from_pretrained(tokenizer_name_or_path)
    except Exception as e:
        # the tokenizer is not valid
        logger.error(f"Invalid tokenizer: {tokenizer_name_or_path}. Error: {e}")
        raise e

    conf["tokenizer"]["identifier"] = tokenizer_name_or_path  # pyright: ignore

    if tokenizer_name_or_path == "allenai/gpt-neox-instella-dolma-v1_5" or tokenizer_name_or_path.endswith(
        "allenai_eleuther-ai-gpt-neox-20b-pii-special.json"
    ):
        if fix_tokenizer:
            conf["model"]["eos_token_id"] = 50279  # pyright: ignore

    om.save(conf, path)


def download_s3_directory(bucket_name: str, prefix: str, local_dir: str, ignore: str | None = None):
    # Create S3 client
    s3_client = _get_s3_client("s3")

    re_ignore = re.compile(ignore) if ignore else None

    # List objects within the given prefix
    paginator = s3_client.get_paginator("list_objects_v2")
    pages = paginator.paginate(Bucket=bucket_name, Prefix=prefix)

    # Create a list to hold all the files to download
    files_to_download = []
    for page in pages:
        for obj in page.get("Contents", []):
            if re_ignore and re_ignore.search(obj["Key"]):
                continue
            files_to_download.append(obj["Key"])

    # Initialize the progress bar
    for s3_key in tqdm(files_to_download, desc="Downloading files"):
        # Construct the full local path
        local_file_path = os.path.join(local_dir, os.path.relpath(s3_key, prefix))
        local_file_dir = os.path.dirname(local_file_path)

        # Ensure local directory exists
        if not os.path.exists(local_file_dir):
            os.makedirs(local_file_dir)

        # Download the file
        s3_client.download_file(bucket_name, s3_key, local_file_path)


def make_local_checkpoint(checkpoint_dir: str) -> str:
    parsed_dir = urlparse(checkpoint_dir)
    print(parsed_dir)

    assert parsed_dir.scheme in ["s3", ""], "Only s3 and local paths are supported."

    if os.path.exists(checkpoint_dir):
        return checkpoint_dir

    temp_dir = os.path.join(tempfile.gettempdir(), md5(checkpoint_dir.encode()).hexdigest())
    if os.path.exists(temp_dir):
        return temp_dir
    try:
        os.makedirs(temp_dir, exist_ok=True)
        print(f"Downloading checkpoint to {temp_dir}...")
        download_s3_directory(
            bucket_name=parsed_dir.netloc,
            prefix=parsed_dir.path.lstrip("/"),
            local_dir=temp_dir,
            ignore=r"/(optim|train)/",
        )
    except Exception as e:
        logger.error(f"Error downloading checkpoint: {e}")
        shutil.rmtree(temp_dir)
        raise e

    return temp_dir


def upload_local_checkpoint(local_checkpoint_dir: str, destination_dir: str):
    if destination_dir == local_checkpoint_dir:
        return
    elif (parsed_url := urlparse(destination_dir)).scheme == "s3":
        s3_bucket_name = parsed_url.netloc
        s3_prefix = parsed_url.path[1:]

        local_paths = [
            os.path.join(root, post_fn)
            for root, _, files in os.walk(local_checkpoint_dir)
            for post_fn in files
            if os.path.basename(post_fn) in HF_FILENAMES
        ]
        dest_paths = [
            os.path.join(s3_prefix, os.path.relpath(local_path, local_checkpoint_dir))
            for local_path in local_paths
        ]

        s3_client = _get_s3_client("s3")
        for local_path, dest_path in tqdm(
            zip(local_paths, dest_paths), desc="Uploading files", total=len(local_paths)
        ):
            s3_client.upload_file(local_path, s3_bucket_name, dest_path)
    elif parsed_url.scheme == "":
        shutil.copytree(local_checkpoint_dir, destination_dir)
    else:
        raise ValueError(f"Unsupported destination: {destination_dir}. Only s3 and local paths are supported.")


def maybe_unshard(checkpoint_dir: str):
    if os.path.exists(os.path.join(checkpoint_dir, "model.pt")):
        return

    print(f"Unsharding {checkpoint_dir}...")
    train_config = TrainConfig.load(os.path.join(checkpoint_dir, "config.yaml"))
    checkpointer = build_sharded_checkpointer(train_config)
    model_state, _, _ = checkpointer.unshard_checkpoint(
        load_path=checkpoint_dir, load_optimizer_state=False, load_trainer_state=False
    )
    torch.save(model_state, os.path.join(checkpoint_dir, "model.pt"))


def main():
    parser = argparse.ArgumentParser(
        description="Adds a config.json to the checkpoint directory, and creates pytorch_model.bin, "
        "making it easier to load weights as HF models."
    )
    parser.add_argument(
        "--checkpoint-dir",
        help="Location of Instella checkpoint.",
        required=True,
    )

    parser.add_argument(
        "--destination-dir",
        help="Location to save the converted checkpoint; default is the same as the checkpoint-dir.",
        default=None,
    )

    parser.add_argument(
        "--ignore-instella-compatibility",
        action="store_true",
        help="Ignore compatibility with the instella codebase. "
        "This will remove files that are needed specifically for instella codebase, eg. config.yaml, etc.",
    )
    parser.add_argument(
        "--logger-level",
        default="warning",
        help="Set the logger level.",
    )

    parser.add_argument(
        "--tokenizer",
        help="Override the tokenizer to use for the checkpoint.",
    )

    parser.add_argument(
        "--fix-tokenizer",
        help="Change eos token to 50279",
    )

    parser.add_argument(
        "--keep-instella-artifacts",
        action="store_true",
        help="Keep instella-specific writ in the checkpoint.",
    )

    args = parser.parse_args()

    args.destination_dir = args.destination_dir or args.checkpoint_dir
    logging.basicConfig()
    logger.setLevel(logging.getLevelName(args.logger_level.upper()))

    local_checkpoint_dir = make_local_checkpoint(args.checkpoint_dir)
    args.checkpoint_dir = local_checkpoint_dir
    maybe_unshard(local_checkpoint_dir)

    validate_tokenizer(checkpoint_dir=local_checkpoint_dir, tokenizer_name_or_path=args.tokenizer, fix_tokenizer=args.fix_tokenizer)
    convert_checkpoint(args.checkpoint_dir, args.ignore_instella_compatibility, args.tokenizer,)

    if not args.keep_instella_artifacts:
        print("Removing non-HF artifacts...")
        os.remove(os.path.join(local_checkpoint_dir, "config.yaml"))
        os.remove(os.path.join(local_checkpoint_dir, "model.pt"))
        os.remove(os.path.join(local_checkpoint_dir, "train.pt"))
        os.remove(os.path.join(local_checkpoint_dir, "optim.pt"))
        shutil.rmtree(os.path.join(local_checkpoint_dir, "optim"), ignore_errors=True)
        shutil.rmtree(os.path.join(local_checkpoint_dir, "model"), ignore_errors=True)
        shutil.rmtree(os.path.join(local_checkpoint_dir, "train"), ignore_errors=True)

    upload_local_checkpoint(local_checkpoint_dir, args.destination_dir)

    print(f"Converted checkpoint saved to {args.destination_dir}")


if __name__ == "__main__":
    main()