import inspect

from PytorchBoot.component import Component
from PytorchBoot.utils.log_util import Log
import PytorchBoot.namespace as namespace


def ensure_component_subclass(cls, type_name, name):
    if not issubclass(cls, Component):
        new_cls  = type(cls.__name__, (Component, cls), {
            **cls.__dict__,
            "TYPE": type_name,
            "NAME": name
        })
        new_cls.__original_class__ = cls
    else:
        new_cls = cls
    for method_name, method in inspect.getmembers(cls, predicate=inspect.isfunction):
        if getattr(method, "__isabstractmethod__", False):
            Log.error(f"Component <{name}> contains abstract method <{method_name}>.", True)
        
    return cls


# --- Classes --- #
dataset_classes = {}
dataset_comments = {}
def dataset(dataset_name, comment=None):
    def decorator(cls):
        if not hasattr(cls, "get_loader") or not callable(getattr(cls, "get_loader")):
            Log.error(f"dataset <{cls.__name__}> must implement a 'get_loader' method", True)
        cls = ensure_component_subclass(cls, namespace.Stereotype.DATASET, dataset_name)
        dataset_comments[dataset_name] = comment
        dataset_classes[dataset_name] = cls
        return cls
    return decorator

module_classes = {}
module_comments = {}
def module(module_name, comment=None):
    def decorator(cls):
        cls = ensure_component_subclass(cls, namespace.Stereotype.MODULE, module_name)
        module_comments[module_name] = comment
        module_classes[module_name] = cls
        return cls
    return decorator

pipeline_classes = {}
pipline_comments = {}
def pipeline(pipeline_name, comment=None):
    def decorator(cls):
        if not hasattr(cls, 'forward') or not callable(getattr(cls, 'forward')):
            Log.error(f"pipeline <{cls.__name__}> must implement a 'forward' method", True)
        cls = ensure_component_subclass(cls, namespace.Stereotype.PIPELINE, pipeline_name)
        pipeline_classes[pipeline_name] = cls
        pipline_comments[pipeline_name] = comment
        return cls
    return decorator

runner_classes = {}
runner_comments = {}
def runner(runner_name, comment=None):
    def decorator(cls):
        if not hasattr(cls, 'run') or not callable(getattr(cls, 'run')):
            Log.error(f"runner <{cls.__name__}> must implement a 'run' method", True)
        cls = ensure_component_subclass(cls, namespace.Stereotype.RUNNER, runner_name)
        runner_classes[runner_name] = cls
        runner_comments[runner_name] = comment
        return cls
    return decorator

factory_classes = {}
factory_comments = {}
def factory(factory_name, comment=None):
    def decorator(cls):
        if not hasattr(cls, 'create') or not callable(getattr(cls, 'create')):
            Log.error(f"factory <{cls.__name__}> must implement a 'create' method", True)
        
        cls = ensure_component_subclass(cls, namespace.Stereotype.FACTORY, factory_name)
        factory_classes[factory_name] = cls
        factory_comments[factory_name] = comment
        return cls
    return decorator

loss_classes = {}
loss_comments = {}
def loss_function(loss_name, comment=None):
    def decorator(cls):
        
        if not hasattr(cls, 'compute') or not callable(getattr(cls, 'compute')):
            Log.error(f"loss function <{cls.__name__}> must implement a 'compute' method", True)
        cls = ensure_component_subclass(cls, namespace.Stereotype.LOSS_FUNCTION, loss_name)
        loss_classes[loss_name] = cls
        loss_comments[loss_name] = comment
        return cls
    return decorator

evaluation_classes = {}
evaluation_comments = {}
def evaluation_method(evaluation_name, comment=None):
    def decorator(cls):
        
        if not hasattr(cls, 'evaluate') or not callable(getattr(cls, 'evaluate')):
            Log.error(f"evaluation method <{cls.__name__}> must implement a 'evaluate' method", True)
        cls = ensure_component_subclass(cls, namespace.Stereotype.EVALUATION_METHOD, evaluation_name)
        evaluation_classes[evaluation_name] = cls
        evaluation_comments[evaluation_name] = comment
        return cls
    return decorator

# --- Others --- #
EXTERNAL_FRONZEN_MODULES = set()

def external_frozen_module(cls):
    if not hasattr(cls, 'load') or not callable(getattr(cls, 'load')):
        Log.error(f"external module <{cls.__name__}> must implement a 'load' method", True)
    EXTERNAL_FRONZEN_MODULES.add(cls)
    return cls

# --- Utils --- #

all_component_classes = {
        namespace.Stereotype.DATASET: dataset_classes,
        namespace.Stereotype.MODULE: module_classes,
        namespace.Stereotype.PIPELINE: pipeline_classes,
        namespace.Stereotype.RUNNER: runner_classes,
        namespace.Stereotype.LOSS_FUNCTION: loss_classes,
        namespace.Stereotype.EVALUATION_METHOD: evaluation_classes,
        namespace.Stereotype.FACTORY: factory_classes
}

all_component_comments = {
        namespace.Stereotype.DATASET: dataset_comments,
        namespace.Stereotype.MODULE: module_comments,
        namespace.Stereotype.PIPELINE: pipline_comments,
        namespace.Stereotype.RUNNER: runner_comments,
        namespace.Stereotype.LOSS_FUNCTION: loss_comments,
        namespace.Stereotype.EVALUATION_METHOD: evaluation_comments,
        namespace.Stereotype.FACTORY: factory_comments
    }

def get_all_component_classes():
    return all_component_classes

def get_all_component_comments():
    return all_component_comments

def get_component_classes(component_type):
    return all_component_classes.get(component_type, None)