#! /usr/bin/env python
# _*_ coding: utf-8 _*_
"""
author:     52s
date:       2024/9/26
file:       main.py
version:
doc:
"""
import argparse
import importlib
import sys
from typing import Dict, Type

from corekit.utils import get_debug_logger

_logger = get_debug_logger(__file__)

ModeName = str
CategoryName = str


class BaseTaskConfig:
    """任务配置基类"""

    def __init__(self, engine: str = 'demo', mode: str = 'train', category: str = 'tune'):
        self.engine = engine
        self.mode = mode
        self.category = category

        self.config = {}

    def start(self, *args, **kwargs):
        assert self.config, f'配置为空, 无法开始任务, 请先通过 from_dict() 方法上传任务配置'
        raise NotImplemented

    def get_all_config(self) -> dict:
        # 可设置项更新, 设置优先级 命令行参数 > 前端参数 > 默认参数 ================================
        args = {}  # 初始化
        # 默认参数
        default_config = self.get_default_config()  # 参数的默认值
        args.update(default_config)
        # 前端参数
        frontend_config = self.get_frontend_config()
        args.update(frontend_config)
        # 命令行参数
        command_config = self.get_command_config()
        for key, value in command_config.items():
            if value is None:
                continue
            args[key] = value
        # 可设置项更新, 设置优先级 命令行参数 > 前端参数 > 默认参数 ================================

        args = self.update_special_config(args)
        return args

    def from_dict(self, config: dict = None, **kwargs):
        if config is None:
            config = {}
        self.config.update(config)
        if kwargs:
            self.config.update(kwargs)

    @classmethod
    def get_default_config(cls) -> dict:
        """获得参数的默认值"""
        return {}

    @classmethod
    def get_frontend_config(cls) -> dict:
        """获取用户在前端配置的值"""
        return {}

    @classmethod
    def get_command_config(cls) -> dict:
        """为配置项增加命令行设置功能"""
        _logger.warning(f'你应该在自定义类中为自己的参数添加命令行设置功能, 而不是调用基类方法')
        # 举例
        parser = argparse.ArgumentParser()
        parser.add_argument("--work_dir", type=str, help='工程路径', default=None)
        parser.add_argument("--ip_address", type=str, help='ip地址', default=None)

        command_parse = parser.parse_args()
        commandt_args = vars(command_parse)
        return commandt_args

    @classmethod
    def update_special_config(cls, config: dict) -> dict:
        """更新一些特殊参数"""
        return config


class TaskConfigAPI:
    """任务配置"""
    _register: Dict[ModeName, Dict[CategoryName, Type["BaseTaskConfig"]]] = {}

    def __init__(self, engine: str = 'demo'):
        """
        :param engine: 引擎名
        """
        self.engine = engine

    @classmethod
    def register(cls, mode, category):
        """通过装饰器来注册一个子类策略, 但问题是子类策略文件必须被导入才能实现注册, 故不可直接使用"""

        def decorator(subclass: Type["BaseTaskConfig"]):
            if not cls._register.get(mode):
                cls._register[mode] = {}
            cls._register[mode][category] = subclass
            return subclass

        return decorator

    def build(self, mode: str = 'train', category: str = 'tune') -> BaseTaskConfig:
        """
        :param mode: 运行模式名
        :param category: 类别名
        :return:
        """
        engine = self.engine
        if mode == 'train':
            engine_path: str = f'envs.{engine}'
        elif mode == 'deduce':
            engine_path = f'deduce.{engine}'
        else:
            sys.exit(1)

        importlib.import_module(f'{engine_path}')  # 导入引擎相关模块
        # except (ImportError, ModuleNotFoundError) as _e:
        #     _logger.error(f'导入 {engine_path} 模块失败请检查, 错误信息: {_e}')
        #     sys.exit(1)
        task_config_class = self._register.get(mode, {}).get(category)
        if not task_config_class:
            _logger.error(f'引擎 {self.engine} 没有{mode}模式-{category}类别, 请检查')
            sys.exit(1)
        return task_config_class(self.engine, mode, category)
