# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: E722, F403
"""TVM Script Parser doc AST"""

import ast
import inspect
import typing
from collections import defaultdict

from . import doc_core as doc
from .doc_core import *  # pylint: disable=unused-import,wildcard-import,redefined-builtin,W0614

FnToDoc = typing.Callable[[ast.AST], doc.AST]
FnFromDoc = typing.Callable[[doc.AST], ast.AST]


class Entry:
    """Mapping entry between python AST node type str and doc AST.

    Parameters
    ----------
    to_doc : typing.Optional[FnToDoc]
        The callable methods for converting python AST node to doc AST.

    from_doc : typing.Optional[FnFromDoc]
        The callable methods for converting doc AST to python AST node.
    """

    to_doc: FnToDoc | None
    from_doc: FnFromDoc | None

    def __init__(self):
        self.to_doc = None
        self.from_doc = None


class Registry:
    """Registration map from python AST node type str to methods of conversion
    between python AST node and doc AST node.

    Parameters
    ----------
    _inst : typing.Optional["Registry"]
        The instance of Registry.

    table : typing.Dict[str, Entry]
        The registration map from python AST node type str to methods of conversion
        between python AST node and doc AST node.
    """

    _inst: typing.Optional["Registry"] = None
    table: dict[str, Entry]

    def __init__(self):
        self.table = defaultdict(Entry)


def register_to_doc(name: str):
    """Register the to_doc method for python AST node type.

    Parameters
    ----------
    name : str
        The type of python AST node.

    Returns
    -------
    f : Callable[[FnToDoc], None]
        The function of registering the to_doc method for python AST node type.
    """

    def f(to_doc: FnToDoc):  # pylint: disable=redefined-outer-name
        reg = Registry._inst  # pylint: disable=protected-access
        reg.table[name].to_doc = to_doc

    return f


def register_from_doc(name: str):
    """Register the from_doc method for python AST node type.

    Parameters
    ----------
    name : str
        The type of python AST node.

    Returns
    -------
    f : Callable[[FnFromDoc], None]
        The function of registering the from_doc method for python AST node type.
    """

    def f(to_doc: FnFromDoc):  # pylint: disable=redefined-outer-name
        reg = Registry._inst  # pylint: disable=protected-access
        reg.table[name].from_doc = to_doc

    return f


def _is_atomic_type(node):
    return (
        node is None
        or node in [..., True, False]
        or isinstance(
            node,
            int | float | str | bool | bytes | complex,
        )
    )


def _get_registry_entry(cls_name, attr):
    cls_name = cls_name.split(".")[-1]
    reg = Registry._inst  # pylint: disable=protected-access
    if cls_name in reg.table:
        entry = reg.table[cls_name]
        return getattr(entry, attr, None)
    return None


def from_doc(node):
    """Get original python AST node from doc AST node.

    Parameters
    ----------
    node : doc.AST
        The doc AST node.

    Returns
    -------
    res : ast.AST
        The corresponding AST node.
    """
    if _is_atomic_type(node):
        return node
    if isinstance(node, tuple):
        return tuple(from_doc(n) for n in node)
    if isinstance(node, list):
        return [from_doc(n) for n in node]
    func = _get_registry_entry(node.__class__.__name__, "from_doc")
    if not func:
        raise NotImplementedError(f"from_doc is not implemented for: {node.__class__.__name__}")
    return func(node)


def to_doc(node):
    """Get doc AST node from python AST node.

    Parameters
    ----------
    node : ast.AST
        The AST node.

    Returns
    -------
    res : doc.AST
        The corresponding doc AST node.
    """
    if _is_atomic_type(node):
        return node
    if isinstance(node, tuple):
        return tuple(to_doc(n) for n in node)
    if isinstance(node, list):
        return [to_doc(n) for n in node]
    func = _get_registry_entry(node.__class__.__name__, "to_doc")
    if not func:
        raise NotImplementedError(f"to_doc is not implemented for: {node.__class__.__name__}")
    return func(node)


def parse(
    source: str,
    filename: str = "<unknown>",
    mode: str = "exec",
) -> doc.AST:
    """Parse TVMScript source code str to doc AST.

    Its interface is consistent with python built-in ast.parse.
    And it will parse by python 3.8 first if possible,
    or it will parse with python version in current environment.

    Parameters
    ----------
    source : str
        The TVMScript source code.

    filename : str
        The optional filename of the file where source code locates.

    mode : str
        The parsing mode for ast.parse.

    Returns
    -------
    res : doc.AST
        The parsed doc AST.
    """
    try:
        program = ast.parse(  # pylint: disable=unexpected-keyword-arg
            source=source,
            filename=filename,
            mode=mode,
            feature_version=(3, 8),
        )
    except:  # pylint: disable=bare-except
        program = ast.parse(
            source=source,
            filename=filename,
            mode=mode,
        )
    return to_doc(program)


class NodeVisitor:
    """Node visitor for doc AST"""

    def visit(self, node: doc.AST) -> None:
        if isinstance(node, list | tuple):
            for item in node:
                self.visit(item)
            return
        if not isinstance(node, doc.AST):
            return
        getattr(
            self,
            "visit_" + node.__class__.__name__.split(".")[-1],
            self.generic_visit,
        )(node)

    def generic_visit(self, node: doc.AST) -> None:
        for field in node.__class__._FIELDS:  # pylint: disable=protected-access
            value = getattr(node, field, None)
            if value is None:
                pass
            elif isinstance(value, doc.AST | list | tuple):
                self.visit(value)


class NodeTransformer:
    """Node transformer for doc AST"""

    def visit(self, node: doc.AST) -> doc.AST:
        if isinstance(node, list):
            return [self.visit(item) for item in node]
        if isinstance(node, tuple):
            return tuple(self.visit(item) for item in node)
        if not isinstance(node, doc.AST):
            return node
        return getattr(
            self,
            "visit_" + node.__class__.__name__.split(".")[-1],
            self.generic_visit,
        )(node)

    def generic_visit(self, node: doc.AST) -> doc.AST:
        kv: dict[str, typing.Any] = {}
        for field in node.__class__._FIELDS:  # pylint: disable=protected-access
            value = getattr(node, field, None)
            if value is None:
                pass
            elif isinstance(value, doc.AST | list | tuple):
                value = self.visit(value)
            kv[field] = value
        return node.__class__(**kv)


def _register_default():
    class DefaultTranslator:
        def __init__(self, doc_cls, func, fields):
            self.doc_cls = doc_cls  # getattr(doc, name)
            self.func = func
            self.fields = fields

        def __call__(self, node):
            kv = {attr: self.func(getattr(node, attr, None)) for attr in self.fields}
            return self.doc_cls(**kv)

    Registry._inst = Registry()  # pylint: disable=protected-access
    for cls_name in dir(doc):
        doc_cls = getattr(doc, cls_name)
        if not hasattr(ast, cls_name):
            continue
        if inspect.isclass(doc_cls) and issubclass(doc_cls, doc.AST):
            assert "." not in cls_name
            register_to_doc(cls_name)(
                DefaultTranslator(
                    getattr(doc, cls_name),
                    to_doc,
                    doc_cls._FIELDS,  # pylint: disable=protected-access
                )
            )
            register_from_doc(cls_name)(
                DefaultTranslator(
                    getattr(ast, cls_name),
                    from_doc,
                    doc_cls._FIELDS,  # pylint: disable=protected-access
                )
            )


_register_default()