# coding: utf-8
# Copyright (c) Max-Planck-Institut für Eisenforschung GmbH - Computational Materials Design (CM) Department
# Distributed under the terms of "New BSD License", see the LICENSE file.

import ast
import inspect
import re
import string
import sys
import textwrap
import warnings
from dataclasses import dataclass
from functools import wraps
from typing import (
    Annotated,
    Any,
    Callable,
    get_args,
    get_origin,
    get_type_hints,
)

from flowrep.tools import get_function_metadata
from pint import Quantity, UnitRegistry
from pint.registry_helpers import (
    _apply_defaults,
    _parse_wrap_args,
    _replace_units,
    _to_units_container,
)

from semantikon.datastructure import TypeMetadata

__author__ = "Sam Waseda"
__copyright__ = (
    "Copyright 2021, Max-Planck-Institut für Eisenforschung GmbH "
    "- Computational Materials Design (CM) Department"
)
__version__ = "1.0"
__maintainer__ = "Sam Waseda"
__email__ = "waseda@mpie.de"
__status__ = "development"
__date__ = "Aug 21, 2021"


class NotAstNameError(TypeError): ...


def _get_ureg(args: Any, kwargs: dict[str, Any]) -> UnitRegistry | None:
    for arg in args + tuple(kwargs.values()):
        if isinstance(arg, Quantity):
            return arg._REGISTRY
    return None


def parse_metadata(value: Any) -> TypeMetadata:
    """
    Parse the metadata of a Quantity object.

    Args:
        value: Quantity object

    Returns:
        dictionary of the metadata. Available keys are `units`, `label`,
        `triples`, `uri` and `shape`. See `semantikon.dataclasses.TypeMetadata` for more details.
    """
    if all(isinstance(d, dict) for d in value.__metadata__):
        merged_dict = {key: val for d in value.__metadata__ for key, val in d.items()}
        return TypeMetadata(**merged_dict)
    metadata = value.__metadata__[0]
    return TypeMetadata(**{k: v for k, v in zip(metadata[::2], metadata[1::2])})


def meta_to_dict(
    value: Any, default=inspect.Parameter.empty, flatten_metadata: bool = True
) -> dict[str, Any]:
    semantikon_was_used = hasattr(value, "__metadata__")
    type_hint_was_present = value is not inspect.Parameter.empty
    default_is_defined = default is not inspect.Parameter.empty
    if semantikon_was_used:
        if flatten_metadata:
            result = parse_metadata(value).to_dictionary()
        else:
            result = {"metadata": parse_metadata(value)}
        if hasattr(value.__args__[0], "__forward_arg__"):
            result["dtype"] = value.__args__[0].__forward_arg__
        else:
            result["dtype"] = value.__args__[0]
    else:
        result = {}
        if type_hint_was_present:
            result["dtype"] = value
    if default_is_defined:
        result["default"] = default
    return result


def extract_undefined_name(error_message: str) -> str:
    match = re.search(r"name '(.+?)' is not defined", error_message)
    if match:
        return match.group(1)
    raise ValueError(
        "No undefined name found in the error message: {}".format(error_message)
    )


def _resolve_annotation(annotation, func_globals=None):
    if func_globals is None:
        func_globals = globals()
    if not isinstance(annotation, str):
        return annotation
    # Lazy annotations: evaluate manually
    try:
        return eval(annotation, func_globals)
    except NameError as e:
        # Handle undefined names in lazy annotations
        undefined_name = extract_undefined_name(str(e))
        if undefined_name == annotation:
            return annotation
        new_annotations = eval(annotation, func_globals | {undefined_name: object})
        args = get_args(new_annotations)
        assert len(args) == 2, "Invalid annotation format"
        return Annotated[undefined_name, args[1]]


def _to_tag(item: Any, count=None, must_be_named: bool = False) -> str:
    if isinstance(item, ast.Name):
        return item.id
    elif must_be_named:
        raise NotAstNameError(
            "With `must_be_named=True`, item must be captured in an `ast.Name` "
            "variables, i.e only simple variable(-s) not containing any operation or "
            "other protected character can be returned."
        )
    elif count is None:
        return "output"
    else:
        return f"output_{count}"


def get_return_expressions(
    func: Callable, separate_tuple: bool = True, strict: bool = False
) -> str | tuple[str, ...] | None:
    source = inspect.getsource(func)
    source = textwrap.dedent(source)
    parsed = ast.parse(source)

    func_node = next(n for n in parsed.body if isinstance(n, ast.FunctionDef))

    ret_list: list[str | tuple[str, ...]] = []

    for node in ast.walk(func_node):
        if isinstance(node, ast.Return):
            value = node.value
            if value is None:
                ret_list.append("None")
            elif isinstance(value, ast.Tuple):
                ret_list.append(
                    tuple(
                        [
                            _to_tag(elt, ii, must_be_named=strict)
                            for ii, elt in enumerate(value.elts)
                        ]
                    )
                )
            else:
                ret_list.append(_to_tag(value, must_be_named=strict))

    if len(ret_list) == 0 and not strict:
        return None
    elif len(set(ret_list)) == 1 and (
        separate_tuple or not isinstance(ret_list[0], tuple)
    ):
        return ret_list[0]
    elif (
        all(isinstance(exp, tuple) for exp in ret_list)
        and len(set(len(r) for r in ret_list)) == 1
        and separate_tuple
        and not strict
    ):
        return tuple([f"output_{i}" for i in range(len(ret_list[0]))])
    elif strict:
        raise NotAstNameError(
            "With `strict=True`, all returns must be captured in independent "
            "variables."
        )
    return "output"


def get_return_labels(
    func: Callable, separate_tuple: bool = True, strict: bool = False
) -> tuple[str, ...]:
    return_vars = get_return_expressions(
        func, separate_tuple=separate_tuple, strict=strict
    )
    if return_vars is None:
        return ("None",)
    elif isinstance(return_vars, str):
        return (return_vars,)
    elif isinstance(return_vars, tuple) and all(
        isinstance(v, str) for v in return_vars
    ):
        return return_vars
    raise TypeError(
        f"{get_return_labels.__module__}.{get_return_labels.__qualname__} expected "
        f"None, a string, or a tuple of strings, but got {return_vars}"
    )


def get_annotated_type_hints(func: Callable) -> dict[str, Any]:
    """
    Get the type hints of a function, including lazy annotations. The function
    practically does the same as `get_type_hints` for Python 3.11 and later,

    Args:
        func: function to be parsed

    Returns:
        dictionary of the type hints. The keys are the names of the arguments
        and the values are the type hints. The return type is stored under the
        key "return".
    """
    try:
        if sys.version_info >= (3, 11):
            # Use the official, public API
            return get_type_hints(func, include_extras=True)
        else:
            # Manually inspect __annotations__ and resolve them
            hints = {}
            sig = inspect.signature(func)
            for name, param in sig.parameters.items():
                hints[name] = _resolve_annotation(param.annotation, func.__globals__)
            if sig.return_annotation is not inspect.Signature.empty:
                hints["return"] = _resolve_annotation(
                    sig.return_annotation, func.__globals__
                )
            return hints
    except NameError:
        hints = {}
        for key, value in func.__annotations__.items():
            hints[key] = _resolve_annotation(value, func.__globals__)
        if hasattr(func, "__return_annotation__"):
            hints["return"] = _resolve_annotation(
                func.__return_annotation__, func.__globals__
            )
        return hints


def parse_input_args(func: Callable) -> dict[str, dict]:
    """
    Parse the input arguments of a function.

    Args:
        func: function to be parsed

    Returns:
        dictionary of the input arguments. Available keys are `units`, `label`,
        `triples`, `uri` and `shape`. See `semantikon.typing.u` for more details.
    """
    type_hints = get_annotated_type_hints(func)
    return {
        key: meta_to_dict(type_hints.get(key, value.annotation), value.default)
        for key, value in inspect.signature(func).parameters.items()
    }


def parse_output_args(
    func: Callable, separate_tuple: bool = True
) -> dict | tuple[dict, ...]:
    """
    Parse the output arguments of a function.

    Args:
        func: function to be parsed

    Returns:
        dictionary of the output arguments if there is only one output. Otherwise,
        a list of dictionaries is returned. Available keys are `units`,
        `label`, `triples`, `uri` and `shape`. See `semantikon.typing.u` for
        more details.
    """
    ret = get_annotated_type_hints(func).get("return", inspect.Parameter.empty)
    if get_origin(ret) is tuple and separate_tuple:
        return tuple([meta_to_dict(ann) for ann in get_args(ret)])
    else:
        return meta_to_dict(ret)


def _get_converter(func: Callable) -> Callable | None:
    args = []
    for value in parse_input_args(func).values():
        if value is not None:
            args.append(value.get("units", value.get("unit")))
        else:
            args.append(None)
    if any([arg is not None for arg in args]):
        return _parse_wrap_args(args)
    else:
        return None


def _get_ret_units(
    output: dict, ureg: UnitRegistry, names: dict[str, Any]
) -> Quantity | None:
    if output == {}:
        return None
    ret = _to_units_container(output.get("units", output.get("unit")), ureg)
    names = {key: 1.0 * value.units for key, value in names.items()}
    return ureg.Quantity(1, _replace_units(ret[0], names) if ret[1] else ret[0])


def _get_output_units(
    output: dict | tuple, ureg: UnitRegistry, names: dict[str, Any]
) -> Quantity | tuple[Quantity | None, ...] | None:
    if isinstance(output, tuple):
        return tuple([_get_ret_units(oo, ureg, names) for oo in output])
    else:
        return _get_ret_units(output, ureg, names)


def _is_dimensionless(output: Quantity | tuple[Quantity, ...] | None) -> bool:
    if output is None:
        return True
    if isinstance(output, tuple):
        return all([_is_dimensionless(oo) for oo in output])
    if output.to_base_units().magnitude == 1.0 and output.dimensionless:
        return True
    return False


def units(func: Callable) -> Callable:
    """
    Decorator to convert the output of a function to a Quantity object with
    the specified units.

    Args:
        func: function to be decorated

    Returns:
        decorated function
    """
    sig = inspect.signature(func)
    converter = _get_converter(func)

    @wraps(func)
    def wrapper(*args, **kwargs):
        ureg = _get_ureg(args, kwargs)
        if converter is None or ureg is None:
            return func(*args, **kwargs)
        args, kwargs = _apply_defaults(sig, args, kwargs)

        # Extend kwargs to account for **kwargs
        ext_kwargs = {
            key: kwargs.get(key, 0) for key in list(sig.parameters.keys())[len(args) :]
        }

        args, new_kwargs, names = converter(ureg, sig, args, ext_kwargs, strict=False)
        for key in list(new_kwargs.keys()):
            if key not in kwargs:
                new_kwargs.pop(key)

        try:
            output_units = _get_output_units(parse_output_args(func), ureg, names)
        except AttributeError:
            output_units = None

        if _is_dimensionless(output_units):
            return func(*args, **new_kwargs)
        elif isinstance(output_units, tuple):
            return tuple(
                [oo * ff for oo, ff in zip(output_units, func(*args, **new_kwargs))]
            )
        else:
            return output_units * func(*args, **new_kwargs)

    return wrapper


def get_function_dict(function: Callable) -> dict[str, Any]:
    result = get_function_metadata(function, full_metadata=True)
    if hasattr(function, "_semantikon_metadata"):
        result.update(function._semantikon_metadata)
    return result


def semantikon_dataclass(cls: type) -> type:
    """
    A class decorator to append type hints to class attributes.

    Args:
        cls: class to be decorated

    Returns:
        The modified class with type hints appended to its attributes.

    Comments:

    >>> from typing import Annotated
    >>> from semantikon.converter import semantikon_dataclass

    >>> @semantikon_dataclass
    >>> class Pizza:
    ...     price: Annotated[float, "money"]
    ...     size: Annotated[float, "dimension"]

    >>>     class Topping:
    ...         sauce: Annotated[str, "matter"]

    >>> append_types(Pizza)
    >>> print(Pizza)
    >>> print(Pizza.Topping)
    >>> print(Pizza.size)
    >>> print(Pizza.price)
    >>> print(Pizza.Topping.sauce)

    Outputs:

    <class '__main__.Pizza'>
    <class '__main__.Pizza.Topping'>
    Annotated[float, 'dimension']
    Annotated[float, 'money']
    Annotated[str, 'matter']
    """
    cls = dataclass(cls)
    for key, value in cls.__dict__.items():
        if isinstance(value, type):
            setattr(cls, key, semantikon_dataclass(getattr(cls, key)))
    try:
        for key, value in cls.__annotations__.items():
            setattr(cls, key, value)  # Append type hints to attributes
    except AttributeError:
        pass
    setattr(cls, "_is_semantikon_class", True)
    return cls


def with_explicit_defaults(**messages) -> Callable:
    """
    Decorator to marks a value as an explicit default, which can be used to
    indicate that a value should be replaced with a default value in the
    context of serialization or processing.

    Args:
        **messages: keyword arguments where the key is the name of the argument
        to be checked for an explicit default, and the value is the warning
        message to be issued if the default is used. If the value is `None`,
        a generic warning message will be issued.

    Returns:
        decorated function that replaces explicit defaults with the actual default
        value and issues a warning if the default is used.

    Example:

    >>> @with_explicit_defaults(x=None)
    >>> def f(x=3):
    ...     return x

    >>> f()  # This will return 3, and a warning will be issued.

    >>> f(3)  # This will also return 3 but without any warning.
    """

    if len(messages) == 0:
        raise ValueError(
            "At least one argument must be provided to the decorator."
            "Read the docstring for more details."
        )

    def decorator(func):
        sig = inspect.signature(func)

        @wraps(func)
        def wrapper(*args, **kwargs):
            bound = sig.bind_partial(*args, **kwargs)

            for name in messages:
                if name not in bound.arguments:
                    if messages[name] is not None:
                        warnings.warn(messages[name])
                    else:
                        warnings.warn(
                            f"'{name}' not provided,"
                            f" using default: {sig.parameters['x'].default}"
                        )

            return func(*args, **kwargs)

        return wrapper

    return decorator


def to_identifier(s: str) -> str:
    """
    Convert an arbitrary string into a valid Python identifier by
    replacing non-identifier characters with underscores.

    Rules applied:
    - ASCII letters, digits, and underscores are preserved
    - All other characters are replaced with '_'
    - If the identifier would start with a digit, prefix it with '_'

    Args:
        s (str): Input string.

    Returns:
        str: A valid Python identifier.
    """
    allowed = string.ascii_letters + string.digits + "_"
    result = "".join(c if c in allowed else "_" for c in s)

    if not result or result[0].isdigit():
        result = "_" + result

    return result
