from typing import Callable, Dict
from bergen.models import Node
from bergen.schema import NodeType
from bergen.types.node.ports.arg import *
from bergen.types.node.ports.kwarg import *
from bergen.types.node.ports.returns import *
import logging
from inspect import signature, Parameter
import inspect
from django.db import models
from docstring_parser import parse

logger = logging.getLogger(__name__)



def parseFunctionToDict(function: Callable, widgets: dict = {}, allow_empty_doc=False, interface=None):
    is_generator = inspect.isasyncgenfunction(function) or inspect.isgenerator(function)
    logger.info(f"Node is {'Generator' if is_generator else 'Function'}")

    sig = signature(function)
    # Generate Args and Kwargs from the Annotation
    args = []
    kwargs = []
    function_inputs = sig.parameters
    for key, value in function_inputs.items():
        widget = widgets.get(key, None)
        the_class = value.annotation

        print(the_class)

        if value.default == Parameter.empty: # Non default parameters are args
            if issubclass(the_class, models.Model):
                args.append(ModelArgPort.fromParameter(value, the_class._meta, widget=widget))
            if issubclass(the_class, int):
                args.append(IntArgPort.fromParameter(value, widget=widget))
            if issubclass(the_class, str):
                args.append(StringArgPort.fromParameter(value, widget=widget))
        else:
            if issubclass(the_class, models.Model):
                kwargs.append(ModelKwargPort.fromParameter(value, the_class._meta, widget=widget))
            if issubclass(the_class, int):
                kwargs.append(IntKwargPort.fromParameter(value, widget=widget))
            if issubclass(the_class, str):
                kwargs.append(StringKwargPort.fromParameter(value, widget=widget))



    # Generate types from the Output
    returns = []

    function_output = sig.return_annotation
    try:
        # Raises type error if we use it with a class but needed here because typing is actually not a class but an Generic Alias :rolling_eyes::
        if function_output._name == "Tuple":
            for type in function_output.__args__:
                if issubclass(type, models.Model):
                    returns.append(ModelReturnPort(value._meta, key=type.__name__))
                if issubclass(type, int):
                    returns.append(IntReturnPort(key=type.__name__))

    except AttributeError:
        # Once here we should have only classes... lets see about that
        if issubclass(function_output, models.Model):
                returns.append(ModelReturnPort(function_output._meta, key=function_output.__name__))
        if issubclass(function_output, int):
                returns.append(IntReturnPort(key=function_output.__name__))

    # Docstring Parser to help with descriptions
    docstring = parse(function.__doc__)
    if docstring.long_description is None:
        assert allow_empty_doc is not False, f"We don't allow empty documentation for function {function.__name__}. Please Provide"
        logger.warn(f"Allowing empty Documentatoin. Please consider providing a documentation for function {function.__name__}")


    name = interface or docstring.short_description or function.__name__
    description = docstring.long_description or "No Description"

    doc_param_map = {param.arg_name: {
        "required": not param.is_optional,
        "description": param.description if not param.description.startswith("[") else None,
    } for param in docstring.params}


    if docstring.returns:
        return_description = docstring.returns.description
        seperated_list = return_description.split(",")
        assert len(returns) == len(seperated_list), "Length of Description and Returns not Equal: If you provide a Return Annotation make sure you seperate the description for each port with ','"
        for index, doc in enumerate(seperated_list):
            returns[index].description = doc

    # TODO: Update with documentatoin.... (Set description for portexample)
    for port in args:
        if port.key in doc_param_map:
            updates = doc_param_map[port.key]
            port.description = updates["description"] or port.description

    for port in kwargs:
        if port.key in doc_param_map:
            updates = doc_param_map[port.key]
            port.description = updates["description"] or port.description


    logger.info(f"Creating Arg Ports: {[str(port.key) for  port in args]}")
    logger.info(f"Creating Kwargs Ports: {[str(port.key) for  port in kwargs]}")
    logger.info(f"Creating Returns Ports: {[str(port.key) for  port in returns]}")

    return {
        "name": name,
        "interface": interface or function.__name__,
        "description" : description,
        "args" : [arg.serialize() for arg in args],
        "kwargs" : [kwarg.serialize() for kwarg in kwargs],
        "returns" : [re.serialize() for re in returns],
        "type" : NodeType.GENERATOR if is_generator else NodeType.FUNCTION
    }


class NodesRegistry:


    def __init__(self):
        self.node_dicts: Dict[str, Dict] = {}


    def register(self):
        for dictio in self.node_dicts.values():
            print(Node.objects.update_or_create(**dictio))


    def func(self, widgets: dict = {}):

        def decorating(func):
            node_dict = parseFunctionToDict(func, widgets=widgets)
            self.node_dicts[f"{node_dict['interface']}"] = node_dict


        return decorating





registry = NodesRegistry()