# Partly Copyright (C) 2025 Arcee AI
# Partly Copyright (C) 2025-2026 Allura-org
#
# This software is free software: you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This software is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program. If not, see http://www.gnu.org/licenses/.

import importlib
import importlib.resources
import logging
import os
from collections import Counter
from typing import Dict, List, Optional

import tqdm
import transformers

from mergekitty._data import chat_templates
from mergekitty.architecture import (
    ArchitectureInfo,
    get_config_value,
    get_architecture_info,
    set_config_value,
)
from mergekitty.card import generate_card
from mergekitty.config import MergeConfiguration, OutputSliceDefinition
from mergekitty.executor import build_executor
from mergekitty.io.tasks import LoaderCache
from mergekitty.options import MergeOptions
from mergekitty.plan import MergePlanner
from mergekitty.tokenizer import TokenizerInfo


def run_merge(
    merge_config: MergeConfiguration,
    out_path: str,
    options: MergeOptions,
    config_source: Optional[str] = None,
):
    if options.random_seed is not None:
        transformers.trainer_utils.set_seed(options.random_seed)

    if not merge_config.models and not merge_config.slices:
        raise RuntimeError("No output requested")

    model_arch_info = [
        get_architecture_info(m.config(trust_remote_code=options.trust_remote_code))
        for m in merge_config.referenced_models()
    ]
    if not options.allow_crimes:
        if not all(a == model_arch_info[0] for a in model_arch_info[1:]):
            raise RuntimeError(
                "Must specify --allow-crimes to attempt to mix different architectures"
            )
    arch_info = model_arch_info[0]

    # initialize loader cache and set options
    loader_cache = LoaderCache()
    loader_cache.setup(options=options)

    # create config for output model
    cfg_out = _model_out_config(
        merge_config, arch_info, trust_remote_code=options.trust_remote_code
    )

    # warm up loader cache
    for model in (
        pbar := tqdm.tqdm(
            merge_config.referenced_models(),
            desc="Warmup loader cache",
            disable=options.quiet,
        )
    ):
        loader_cache.get(model)
    del pbar

    logging.info("Planning operations")
    targets = MergePlanner(
        merge_config,
        arch_info,
        options=options,
        out_model_config=cfg_out,
    ).plan_to_disk(out_path=out_path)

    exec = build_executor(
        tasks=targets,
        math_device="cuda" if options.cuda else "cpu",
        storage_device="cuda" if options.low_cpu_memory else "cpu",
        executor=options.executor,
    )

    tokenizer = None
    for _task, value in exec.run(quiet=options.quiet):
        if isinstance(value, TokenizerInfo):
            tokenizer = value.tokenizer

    if tokenizer:
        pad_to_multiple_of = None
        if merge_config.tokenizer and merge_config.tokenizer.pad_to_multiple_of:
            pad_to_multiple_of = merge_config.tokenizer.pad_to_multiple_of
        _update_config_vocab(cfg_out, tokenizer, pad_to_multiple_of=pad_to_multiple_of)

    logging.info("Saving config")
    cfg_out.save_pretrained(out_path)

    if options.write_model_card:
        if not config_source:
            config_source = merge_config.to_yaml()

        card_md = generate_card(
            config=merge_config,
            config_yaml=config_source,
            name=os.path.basename(out_path),
        )
        with open(os.path.join(out_path, "README.md"), "w", encoding="utf-8") as fp:
            fp.write(card_md)

        with open(
            os.path.join(out_path, "mergekitty_config.yml"), "w", encoding="utf-8"
        ) as fp:
            fp.write(config_source)

    if tokenizer is None:
        logging.error(
            "Somehow, the tokenizer was not created during the merging process -- did you manually set `tokenizer_source` to None without setting `tokenizer`?"
        )
    else:
        logging.info("Saving tokenizer")
        _set_chat_template(tokenizer, merge_config)
        tokenizer.save_pretrained(out_path, safe_serialization=True)


def _set_chat_template(
    tokenizer: transformers.PreTrainedTokenizerBase,
    merge_config: MergeConfiguration,
    trust_remote_code: bool = False,
):
    chat_template = merge_config.chat_template
    if not chat_template:
        return

    if chat_template == "auto":
        # see if there is a plurality chat template among the input models
        model_templates = []
        for model in merge_config.referenced_models():
            try:
                tok = transformers.AutoTokenizer.from_pretrained(
                    model.model.path,
                    revision=model.model.revision,
                    trust_remote_code=trust_remote_code,
                )
                template = tok.chat_template
                if isinstance(template, dict):
                    template = template.get("default", None)
                if template:
                    model_templates.append(template.strip())
            except Exception as e:
                logging.warning(f"Unable to load tokenizer for {model}", exc_info=e)

        if not model_templates:
            return

        chat_template = Counter(model_templates).most_common(1)[0][0]
        logging.info(f"Auto-selected chat template: {chat_template}")

    elif importlib.resources.is_resource(chat_templates, chat_template + ".jinja"):
        with importlib.resources.open_text(
            chat_templates, chat_template + ".jinja"
        ) as fp:
            chat_template = fp.read()

    elif len(chat_template) < 20 or "{" not in chat_template:
        raise RuntimeError(f"Invalid chat template: {chat_template}")

    tokenizer.chat_template = chat_template


def _model_out_config(
    config: MergeConfiguration,
    arch_info: ArchitectureInfo,
    trust_remote_code: bool = False,
) -> transformers.PretrainedConfig:
    """Return a configuration for the resulting model."""
    if config.base_model:
        res = config.base_model.config(trust_remote_code=trust_remote_code)
    else:
        res = config.referenced_models()[0].config(trust_remote_code=trust_remote_code)
    if config.out_dtype:
        res.torch_dtype = config.out_dtype
    elif config.dtype:
        res.torch_dtype = config.dtype

    if config.slices:
        try:
            num_layers = sum(
                s.sources[0].layer_range[1] - s.sources[0].layer_range[0]
                for s in config.slices
            )
            num_layers_key = arch_info.num_layers_config_key()
            set_config_value(res, num_layers_key, num_layers)
            _resize_layer_types(
                res,
                num_layers_key,
                num_layers,
                config.slices,
                trust_remote_code=trust_remote_code,
            )
        except Exception as e:
            logging.warning(
                "Unable to set number of layers in output config - you may need to manually correct it.",
                exc_info=e,
            )

    return res


def _update_config_vocab(
    config: transformers.PretrainedConfig,
    tokenizer: transformers.PreTrainedTokenizerBase,
    pad_to_multiple_of: Optional[int] = None,
):
    vocab_size = len(tokenizer.get_vocab())
    if pad_to_multiple_of and vocab_size % pad_to_multiple_of:
        vocab_size = vocab_size + pad_to_multiple_of - (vocab_size % pad_to_multiple_of)
    try:
        if hasattr(config, "vocab_size"):
            config.vocab_size = vocab_size
        elif hasattr(config, "text_config") and hasattr(
            config.text_config, "vocab_size"
        ):
            config.text_config.vocab_size = vocab_size
        else:
            raise AttributeError("config has no vocab_size field")
    except Exception as e:
        logging.warning(
            "Unable to set vocabulary size in output config - you may need to manually correct it.",
            exc_info=e,
        )


def _resize_layer_types(
    config: transformers.PretrainedConfig,
    num_layers_key: str,
    num_layers: int,
    slices: List[OutputSliceDefinition],
    trust_remote_code: bool = False,
) -> None:
    layer_types_key = (
        f"{num_layers_key.rsplit('.', 1)[0]}.layer_types"
        if "." in num_layers_key
        else "layer_types"
    )
    try:
        layer_types = get_config_value(config, layer_types_key)
    except (AttributeError, KeyError):
        return

    if not isinstance(layer_types, (list, tuple)) or not layer_types:
        return

    source_config_cache: Dict[str, transformers.PretrainedConfig] = {}
    resized_layer_types = []
    warned_on_mismatch = False
    for out_slice in slices:
        selected_layer_types = None
        for source in out_slice.sources:
            source_key = str(source.model)
            if source_key not in source_config_cache:
                source_config_cache[source_key] = source.model.config(
                    trust_remote_code=trust_remote_code
                )

            try:
                source_layer_types = get_config_value(
                    source_config_cache[source_key], layer_types_key
                )
            except (AttributeError, KeyError):
                return

            if not isinstance(source_layer_types, (list, tuple)):
                return

            start, end = source.layer_range
            layer_type_segment = list(source_layer_types[start:end])
            if len(layer_type_segment) != end - start:
                return

            if selected_layer_types is None:
                selected_layer_types = layer_type_segment
            elif selected_layer_types != layer_type_segment and not warned_on_mismatch:
                logging.warning(
                    "Input slice sources have different layer_types metadata; "
                    "using the first source model values in the output config."
                )
                warned_on_mismatch = True

        if selected_layer_types is not None:
            resized_layer_types.extend(selected_layer_types)

    if len(resized_layer_types) != num_layers:
        return

    set_config_value(config, layer_types_key, resized_layer_types)


__all__ = ["MergeOptions", "run_merge"]
