import ast
import importlib
import inspect
import logging
import sys
from typing import Union

import astunparse
from griffe import Extension, Object, ObjectNode
from griffe.docstrings.dataclasses import DocstringSectionParameters
from griffe.expressions import Expr
from griffe.logger import patch_loggers


def get_logger(name):
    new_logger = logging.getLogger(name)
    new_logger.setLevel("ERROR")
    return new_logger


patch_loggers(get_logger)

logger = get_logger(__name__)


class EDSNLPDocstrings(Extension):
    def __init__(self):
        super().__init__()

        self.PIPE_OBJ = {}
        self.FACT_MEM = {}
        self.PIPE_TO_FACT = {}

    def on_instance(self, node: Union[ast.AST, ObjectNode], obj: Object) -> None:
        if (
            isinstance(node, ast.Assign)
            and obj.name == "create_component"
            and isinstance(node.value, ast.Call)
            and isinstance(node.value.func, ast.Call)
        ):
            module_name = obj.path.rsplit(".", 1)[0]
            for name, mod in list(sys.modules.items()):
                if name.startswith("edspdf"):
                    importlib.reload(mod)
            module = importlib.reload(importlib.import_module(module_name))

            config_node = node.value.func
            config_node = next(
                (kw.value for kw in config_node.keywords if kw.arg == "default_config"),
                None,
            )
            try:
                default_config = eval(astunparse.unparse(config_node), module.__dict__)
            except Exception:
                default_config = {}

            # import object to get its evaluated docstring
            try:
                runtime_obj = getattr(module, obj.name)
                source = inspect.getsource(runtime_obj)
                self.visit(ast.parse(source))
            except ImportError:
                logger.debug(f"Could not get dynamic docstring for {obj.path}")
                return
            except AttributeError:
                logger.debug(f"Object {obj.path} does not have a __doc__ attribute")
                return

            callee = (
                runtime_obj.__init__
                if hasattr(runtime_obj, "__init__")
                else runtime_obj
            )
            spec = inspect.getfullargspec(callee)
            func_defaults = dict(
                zip(spec.args[-len(callee.__defaults__) :], callee.__defaults__)
                if callee.__defaults__
                else (),
                **(callee.__kwdefaults__ or {}),
            )
            defaults = {**func_defaults, **default_config}
            self.FACT_MEM[obj.path] = (node, obj, defaults)
            pipe_path = runtime_obj.__module__ + "." + runtime_obj.__name__
            self.PIPE_TO_FACT[pipe_path] = obj.path

            if pipe_path in self.PIPE_OBJ:
                pipe = self.PIPE_OBJ[pipe_path]
                obj.docstring = pipe.docstring
            else:
                return
        elif obj.is_class or obj.is_function:
            self.PIPE_OBJ[obj.path] = obj
            if obj.path in self.PIPE_TO_FACT:
                node, fact_obj, defaults = self.FACT_MEM[self.PIPE_TO_FACT[obj.path]]
                fact_obj.docstring = obj.docstring
                obj = fact_obj
            else:
                return
        else:
            return

        if obj.docstring is None:
            return

        param_section: DocstringSectionParameters = None
        obj.docstring.parser = "numpy"
        for section in obj.docstring.parsed:
            if isinstance(section, DocstringSectionParameters):
                param_section = section  # type: ignore

        if param_section is None:
            return

        for param in param_section.value:
            if param.name in defaults:
                param.default = str(defaults[param.name])
            if isinstance(param.default, Expr):
                continue
            if param.default is not None and len(param.default) > 50:
                param.default = param.default[: 50 - 3] + "..."