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

import astunparse
from griffe import (
    DocstringSectionParameters,
    Expr,
    Extension,
    Object,
    ObjectNode,
    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,
        agent: Any = None,
    ) -> 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] + "..."
