# Copyright (c) ModelScope Contributors. All rights reserved.

import os
import platform
import re
import torch
from abc import ABC, abstractmethod
from copy import deepcopy
from dataclasses import asdict, dataclass, field
from transformers import AutoConfig, PretrainedConfig, PreTrainedModel, PreTrainedTokenizerBase
from transformers.utils.versions import require_version
from typing import Any, Dict, List, Literal, Optional, Tuple, Type

from swift.utils import HfConfigFactory, get_logger, safe_snapshot_download
from .utils import get_default_torch_dtype

logger = get_logger()


@dataclass
class Model:
    ms_model_id: Optional[str] = None
    hf_model_id: Optional[str] = None
    model_path: Optional[str] = None

    ms_revision: Optional[str] = None
    hf_revision: Optional[str] = None


@dataclass
class ModelGroup:
    models: List[Model]

    # Higher priority. If set to None, the attributes of the ModelMeta will be used.
    template: Optional[str] = None
    ignore_patterns: Optional[List[str]] = None
    requires: Optional[List[str]] = None
    tags: List[str] = field(default_factory=list)

    def __post_init__(self):
        assert not isinstance(self.template, (list, tuple))  # check ms-swift4.0
        assert isinstance(self.models, (tuple, list)), f'self.models: {self.models}'


class BaseModelLoader(ABC):

    @abstractmethod
    def __init__(self, model_info, model_meta, *args, **kwargs):
        pass

    @abstractmethod
    def load(self) -> Tuple[Optional[PreTrainedModel], PreTrainedTokenizerBase]:
        pass


@dataclass
class ModelMeta:
    model_type: Optional[str]
    # Used to list the model_ids from modelscope/huggingface,
    # which participate in the automatic inference of the model_type.
    model_groups: List[ModelGroup]
    loader: Optional[Type[BaseModelLoader]] = None

    template: Optional[str] = None
    model_arch: Optional[str] = None
    mcore_model_type: Optional[str] = None
    architectures: List[str] = field(default_factory=list)
    # Additional files that need to be saved for full parameter training/merge-lora.
    additional_saved_files: List[str] = field(default_factory=list)
    torch_dtype: Optional[torch.dtype] = None

    is_multimodal: bool = False
    is_reward: bool = False
    task_type: Optional[str] = None

    # File patterns to ignore when downloading the model.
    ignore_patterns: Optional[List[str]] = None
    # Usually specifies the version limits of transformers.
    requires: List[str] = field(default_factory=list)
    tags: List[str] = field(default_factory=list)

    def __post_init__(self):
        from .constant import MLLMModelType, RMModelType
        from .register import ModelLoader
        assert not isinstance(self.loader, str)  # check ms-swift4.0
        if self.loader is None:
            self.loader = ModelLoader
        if not isinstance(self.model_groups, (list, tuple)):
            self.model_groups = [self.model_groups]
        self.candidate_templates = list(
            dict.fromkeys(t for t in [self.template] + [mg.template for mg in self.model_groups] if t is not None))
        if self.model_type in MLLMModelType.__dict__:
            self.is_multimodal = True
        if self.model_type in RMModelType.__dict__:
            self.is_reward = True

    def get_matched_model_group(self, model_name: str) -> Optional[ModelGroup]:
        for model_group in self.model_groups:
            for model in model_group.models:
                for key in ['ms_model_id', 'hf_model_id', 'model_path']:
                    value = getattr(model, key)

                    if isinstance(value, str) and model_name == value.rsplit('/', 1)[-1].lower():
                        return model_group

    def check_requires(self, model_info=None):
        extra_requires = []
        if model_info and model_info.quant_method:
            mapping = {'bnb': ['bitsandbytes'], 'awq': ['autoawq'], 'gptq': ['auto_gptq'], 'aqlm': ['aqlm']}
            extra_requires += mapping.get(model_info.quant_method, [])
        requires = []
        for require in self.requires + extra_requires:
            try:
                require_version(require)
            except ImportError:
                requires.append(f'"{require}"')
        if requires:
            requires = ' '.join(requires)
            logger.warning(f'Please install the package: `pip install {requires} -U`.')


MODEL_MAPPING: Dict[str, ModelMeta] = {}


@dataclass
class ModelInfo:
    model_type: str
    model_dir: str
    torch_dtype: torch.dtype
    max_model_len: int
    quant_method: Literal['gptq', 'awq', 'bnb', 'aqlm', 'hqq', None]
    quant_bits: int

    # extra
    rope_scaling: Optional[Dict[str, Any]] = None
    is_moe_model: bool = False
    is_multimodal: bool = False
    config: Optional[PretrainedConfig] = None
    task_type: Optional[str] = None
    num_labels: Optional[int] = None

    def __post_init__(self):
        self.model_name = get_model_name(self.model_dir)


def get_model_name(model_id_or_path: str) -> Optional[str]:
    assert isinstance(model_id_or_path, str), f'model_id_or_path: {model_id_or_path}'
    # compat hf hub
    model_id_or_path = model_id_or_path.rstrip('/')
    match_ = re.search('/models--.+?--(.+?)/snapshots/', model_id_or_path)
    if match_ is not None:
        return match_.group(1)

    model_name = model_id_or_path.rsplit('/', 1)[-1]
    if platform.system().lower() == 'windows':
        model_name = model_name.rsplit('\\', 1)[-1]
    # compat modelscope snapshot_download
    model_name = model_name.replace('___', '.')
    return model_name


def get_matched_model_meta(model_id_or_path: str) -> Optional[ModelMeta]:
    model_name = get_model_name(model_id_or_path).lower()
    for model_type, model_meta in MODEL_MAPPING.items():
        model_group = ModelMeta.get_matched_model_group(model_meta, model_name)
        if model_group is not None:
            model_meta = deepcopy(model_meta)
            for k, v in asdict(model_group).items():
                if v is not None and k in model_meta.__dict__:
                    setattr(model_meta, k, v)
            return model_meta


def _get_arch_mapping():
    res = {}
    for model_type, model_meta in MODEL_MAPPING.items():
        architectures = model_meta.architectures
        if not architectures:
            architectures.append('null')
        for arch in architectures:
            if arch not in res:
                res[arch] = []
            res[arch].append(model_type)
    return res


def get_matched_model_types(architectures: Optional[List[str]]) -> List[str]:
    """Get possible model_type."""
    architectures = architectures or ['null']
    if architectures:
        architectures = architectures[0]
    arch_mapping = _get_arch_mapping()
    return arch_mapping.get(architectures) or []


def _read_args_json_model_type(model_dir):
    if not os.path.exists(os.path.join(model_dir, 'args.json')):
        return
    from swift.arguments import BaseArguments
    args = BaseArguments.from_pretrained(model_dir)
    return args.model_type


def _get_model_info(model_dir: str, model_type: Optional[str], quantization_config) -> ModelInfo:
    try:
        config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
    except Exception:
        config = PretrainedConfig.get_config_dict(model_dir)[0]
    if quantization_config is not None:
        HfConfigFactory.set_config_attr(config, 'quantization_config', quantization_config)
    quant_info = HfConfigFactory.get_quant_info(config) or {}
    torch_dtype = HfConfigFactory.get_torch_dtype(config, quant_info)
    max_model_len = HfConfigFactory.get_max_model_len(config)
    rope_scaling = HfConfigFactory.get_config_attr(config, 'rope_scaling')
    is_moe_model = HfConfigFactory.is_moe_model(config)
    is_multimodal = HfConfigFactory.is_multimodal(config)

    if model_type is None:
        model_type = _read_args_json_model_type(model_dir)
    if model_type is None:
        architectures = HfConfigFactory.get_config_attr(config, 'architectures')
        model_types = get_matched_model_types(architectures)
        if len(model_types) > 1:
            raise ValueError(f'Failed to automatically match `model_type` for `{model_dir}`. '
                             f'Multiple possible types found: {model_types}. '
                             'Please specify `model_type` manually. See documentation: '
                             'https://swift.readthedocs.io/en/latest/Instruction/Supported-models-and-datasets.html')
        elif len(model_types) == 1:
            model_type = model_types[0]
    elif model_type not in MODEL_MAPPING:
        raise ValueError(f"model_type: '{model_type}' not in {list(MODEL_MAPPING.keys())}")

    res = ModelInfo(
        model_type,
        model_dir,
        torch_dtype,
        max_model_len,
        quant_info.get('quant_method'),
        quant_info.get('quant_bits'),
        rope_scaling=rope_scaling,
        is_moe_model=is_moe_model,
        is_multimodal=is_multimodal,
    )
    return res


def get_model_info_meta(
        model_id_or_path: str,
        *,
        torch_dtype: Optional[torch.dtype] = None,
        # hub
        use_hf: Optional[bool] = None,
        hub_token: Optional[str] = None,
        revision: Optional[str] = None,
        download_model: bool = False,
        # model kwargs
        model_type: Optional[str] = None,
        quantization_config=None,
        task_type=None,
        num_labels=None,
        problem_type=None,
        **kwargs) -> Tuple[ModelInfo, ModelMeta]:
    from .register import ModelLoader
    model_meta = get_matched_model_meta(model_id_or_path)
    model_dir = safe_snapshot_download(
        model_id_or_path,
        revision=revision,
        download_model=download_model,
        use_hf=use_hf,
        ignore_patterns=getattr(model_meta, 'ignore_patterns', None),
        hub_token=hub_token)

    model_type = model_type or getattr(model_meta, 'model_type', None)
    model_info = _get_model_info(model_dir, model_type, quantization_config=quantization_config)
    if model_type is None and model_info.model_type is not None:
        model_type = model_info.model_type
        logger.info(f'Setting model_type: {model_type}')
    if model_type is not None and (model_meta is None or model_meta.model_type != model_type):
        model_meta = MODEL_MAPPING[model_type]
    if model_meta is None:  # not found
        if model_info.is_multimodal:
            raise ValueError(f'Model "{model_id_or_path}" is not supported because no suitable `model_type` was found. '
                             'Please refer to the documentation and specify an appropriate `model_type` manually: '
                             'https://swift.readthedocs.io/en/latest/Instruction/Supported-models-and-datasets.html')
        else:
            model_meta = ModelMeta(None, [], ModelLoader, template='dummy', model_arch=None)
            logger.info(f'Temporarily create model_meta: {model_meta}')
    if torch_dtype is None:
        torch_dtype = model_meta.torch_dtype or get_default_torch_dtype(model_info.torch_dtype)
        logger.info(f'Setting torch_dtype: {torch_dtype}')
    model_info.torch_dtype = torch_dtype
    if task_type is None:
        if model_meta.is_reward:
            num_labels = 1
        if num_labels is None:
            task_type = 'causal_lm'
        else:
            task_type = 'seq_cls'
        if model_meta.task_type is not None:
            task_type = model_meta.task_type

    # Handle reranker task type
    if task_type == 'reranker':
        if num_labels is None:
            num_labels = 1  # Default to 1 for reranker tasks
        logger.info(f'Setting reranker task with num_labels={num_labels}')
    elif task_type == 'generative_reranker':
        # Generative reranker doesn't need num_labels as it uses CausalLM structure
        num_labels = None
        logger.info('Setting generative_reranker task (no num_labels needed)')
    elif task_type == 'seq_cls':
        assert num_labels is not None, 'Please pass the parameter `num_labels`.'
        if problem_type is None:
            if num_labels == 1:
                problem_type = 'regression'
            else:
                problem_type = 'single_label_classification'

    model_info.task_type = task_type
    model_info.num_labels = num_labels
    model_info.problem_type = problem_type
    if model_meta.is_multimodal:
        model_info.is_multimodal = True
    model_meta.check_requires(model_info)
    return model_info, model_meta
