import json
import os
from pathlib import Path
from typing import Any, Dict, Generator, Optional, cast

import torch

from xlm.harness import Harness
from xlm.utils.rich_utils import print_config_tree


if "PROJECT_ROOT" not in os.environ:
    os.environ["PROJECT_ROOT"] = "."
os.environ["HYDRA_FULL_ERROR"] = "1"

# region: Import necessary modules
import hydra
from omegaconf import DictConfig, OmegaConf
from lightning import seed_everything
from xlm.utils.rank_zero import RankedLogger

logger = RankedLogger(__name__, rank_zero_only=True)

# endregion


def get_generation_output_file_path(cfg: DictConfig) -> Optional[Path]:
    generation_output_dir: Optional[Path] = cfg.generation.get(
        "output_dir", None
    )
    generation_output_dir = (
        Path(generation_output_dir)
        if generation_output_dir is not None
        else None
    )
    generation_output_file_name: Optional[str] = cfg.generation.get(
        "output_file_name", None
    )
    if (
        generation_output_file_name is None
        and generation_output_dir is not None
    ):
        raise ValueError(
            "output_file_name is required when output_dir is provided"
        )
    generation_output_file_path: Optional[Path] = None
    if generation_output_file_name is not None:
        assert generation_output_dir is not None
        generation_output_file_path = (
            generation_output_dir / generation_output_file_name
        )
    return generation_output_file_path


def instantiate_model(
    cfg: DictConfig,
    datamodule: Any,
    tokenizer: Any,
) -> Harness:
    """

    Supports two modes:
        1. Load a model from full training checkpoint using `lightning_module.load_from_checkpoint(cfg.generation.ckpt_path)`
        2. Load a model from model only checkpoint using `lightning_module.model.load_state_dict(torch.load(cfg.generation.model_only_checkpoint_path))`

    Args:
        cfg: Hydra config
        generation_ckpt_path: Path to the generation checkpoint
        datamodule: Datamodule
    """
    generation_ckpt_path = cfg.generation.ckpt_path
    torch.set_float32_matmul_precision("medium")
    if generation_ckpt_path is not None:
        module_cls = hydra.utils.get_class(cfg.lightning_module._target_)
        lightning_module = module_cls.load_from_checkpoint(
            checkpoint_path=generation_ckpt_path,
            tokenizer=tokenizer,
            datamodule=datamodule,
            cfg=cfg,  # chance to override the config of the checkpoint
            manual_ema_restore=True,
        )
    else:
        lightning_module = hydra.utils.instantiate(
            cfg.lightning_module,
            tokenizer=tokenizer,
            datamodule=datamodule,
            cfg=cfg,
            _recursive_=False,
            manual_ema_restore=True,
        )
    lightning_module = cast(Harness, lightning_module)
    lightning_module = lightning_module.to("cuda")
    lightning_module.eval()

    # check if we have model only checkpoint (replicating train functionality)
    model_only_ckpt_path = None
    if cfg.generation.get("model_only_checkpoint_path", None) is not None:
        if generation_ckpt_path is not None:
            logger.error(
                "generation.model_only_checkpoint_path and generation.ckpt_path cannot both be provided. "
                "We will use generation.ckpt_path for the model weights as well."
            )
        else:
            if not os.path.isfile(cfg.generation.model_only_checkpoint_path):
                raise ValueError(
                    f"The model only checkpoint path {cfg.generation.model_only_checkpoint_path} does not exist."
                )
            model_only_ckpt_path = cfg.generation.model_only_checkpoint_path

    if model_only_ckpt_path is not None:
        message = lightning_module.model.load_state_dict(
            torch.load(model_only_ckpt_path)
        )
        logger.warning(
            f"Loading weights for `model` from a pretrained model at {model_only_ckpt_path} before generation. "
            "Make sure that the model weights were saved with EMA applied."
        )
        logger.warning(message)
    return lightning_module


def generate(cfg: DictConfig):
    print_config_tree(cfg, resolve=True, save_to_file=False)
    if cfg.get("seed"):
        logger.info(f"Seed everything with seed {cfg.seed}")
        seed_everything(cfg.seed)

    # Always create the global components first.
    global_components: Dict[str, Any] = hydra.utils.instantiate(
        cfg.global_components
    )
    OmegaConf.clear_resolver("global_components")
    OmegaConf.register_new_resolver(
        "global_components", lambda x: global_components[x]
    )
    # instantiate the datamodule
    datamodule = hydra.utils.instantiate(cfg.datamodule)
    tokenizer = datamodule.tokenizer

    # update the omegaconf resolvers
    OmegaConf.clear_resolver("tokenizer")
    OmegaConf.register_new_resolver(
        "tokenizer", lambda x: getattr(tokenizer, x)
    )
    OmegaConf.clear_resolver("datamodule")
    OmegaConf.register_new_resolver(
        "datamodule", lambda x: getattr(datamodule, x)
    )
    datamodule.no_trainer_mode = True
    datamodule.prepare_data()
    datamodule.setup("predict")

    # instantiate the model
    lightning_module = instantiate_model(cfg, datamodule, tokenizer)

    # output file path
    generation_output_file_path = get_generation_output_file_path(cfg)

    if generation_output_file_path is not None:
        logger.info(
            f"Writing generation output to {generation_output_file_path}"
        )
    predict_dataloaders = datamodule.predict_dataloader()
    dataloader_names = []
    if isinstance(predict_dataloaders, list):
        for i, dl in enumerate(predict_dataloaders):
            dataloader_names.append(datamodule.dataloader_names["predict"][i])
    else:
        dataloader_names.append(datamodule.dataloader_names["predict"][0])
        predict_dataloaders = [predict_dataloaders]

    for dataloader_idx, (dataloader_name, dl) in enumerate(
        zip(dataloader_names, predict_dataloaders)
    ):
        for batch_idx, batch in enumerate(dl):
            # hardcode the move to cuda for now
            batch = {
                k: (v.to("cuda") if isinstance(v, torch.Tensor) else v)
                for k, v in batch.items()
            }
            preds = lightning_module.predict_step(
                batch, batch_idx, dataloader_idx
            )
            lightning_module._call_log_predictions(
                batch, preds, "predict", dataloader_name, no_trainer_mode=True
            )
