from __future__ import annotations

import re
from abc import abstractmethod
from datetime import datetime, timedelta
from functools import lru_cache
from typing import TYPE_CHECKING, Any

from RestrictedPython import (
    compile_restricted,
    limited_builtins,
    safe_builtins,
    utility_builtins,
)
from RestrictedPython.Eval import default_guarded_getitem
from RestrictedPython.Guards import (
    full_write_guard,
    guarded_iter_unpack_sequence,
    guarded_unpack_sequence,
)

if TYPE_CHECKING:
    from collections.abc import Callable


def safe_import(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
    if args and args[0] != "time":
        msg = f"{args[0]} package cannot be imported."
        raise ImportError(msg)
    return __import__(*args, **kwargs)


extra_builtins = {
    "datetime": datetime,
    "timedelta": timedelta,
    "sorted": sorted,
    "min": min,
    "max": max,
    "sum": sum,
    "any": any,
    "all": all,
    "__import__": safe_import,
}

builtins = safe_builtins.copy()
builtins.update(utility_builtins)
builtins.update(limited_builtins)
builtins.update(extra_builtins)


def safer_getattr(
    obj: object,
    name: str,
    default: object = None,
    getattr: Callable[..., object] = getattr,  # noqa: A002
) -> object:
    """Getattr implementation which prevents using format on string objects.

    format() is considered harmful:
    http://lucumr.pocoo.org/2016/12/29/careful-with-str-format/

    """
    if name == "format" and isinstance(obj, str):
        msg = f"Using format() on a {obj.__class__.__name__} is not safe."
        raise NotImplementedError(msg)
    if name[0] == "_":
        msg = f'"{name}" is an invalid attribute name because it starts with "_"'
        raise AttributeError(msg)
    return getattr(obj, name, default)


restricted_globals = {
    "__builtins__": builtins,
    "_getattr_": safer_getattr,
    "_write_": full_write_guard,
    "_getiter_": iter,
    "_getitem_": default_guarded_getitem,
    "_iter_unpack_sequence_": guarded_iter_unpack_sequence,
    "_unpack_sequence_": guarded_unpack_sequence,
}

from loguru import logger  # noqa: E402

if TYPE_CHECKING:
    from types import CodeType

# CODE_FORMAT = """{expr})"""  # noqa: ERA001


class Query:
    __slots__ = ("__weakref__", "expr")

    def __init__(self, expr: str) -> None:
        self.expr = expr

    @abstractmethod
    def check(self, **params: Any) -> bool: ...  # noqa: ANN401

    def __call__(self, **params: Any) -> bool:  # noqa: ANN401
        return self.check(**params)


@lru_cache(maxsize=100)
def compile_checker(expr: str) -> CodeType:
    source_code = expr
    return compile_restricted(source_code, filename="<inline code>", mode="eval")


def syntax_error_check(expr: str) -> None:
    if not expr:
        return
    expr = strip_query(expr)
    try:
        compile_restricted(expr, filename="<inline code>", mode="eval")
    except SyntaxError:
        compile(expr, filename="<inline code>", mode="eval")


@lru_cache(maxsize=100)
def strip_query(query: str) -> str:
    stripped_query = query.strip()
    # cut the hardcoded part (SELECT something IF)
    if query.lower().startswith("select"):
        try:
            stripped_query = re.split("if", query, maxsplit=1, flags=re.IGNORECASE)[1]
        except IndexError:
            stripped_query = ""

    if stripped_query:
        stripped_query = f"({stripped_query.strip()})"

    return stripped_query


@lru_cache(maxsize=100)
def query_add_default_expr(query: str) -> str:
    default_expression = "run.is_archived == False"
    if not query:
        return default_expression
    if "run.is_archived" not in query and "run.archived" not in query:
        return f"{default_expression} and {query}"
    return query


class RestrictedPythonQuery(Query):
    __slots__ = ("_checker", "run_metadata_cache")

    allowed_params = frozenset(
        (
            "run",
            "metric",
            "images",
            "audios",
            "distributions",
            "figures",
            "texts",
        ),
    )

    def __init__(self, query: str) -> None:
        stripped_query = strip_query(query)
        expr = query_add_default_expr(stripped_query)
        super().__init__(expr=expr)
        self._checker = compile_checker(expr)
        self.run_metadata_cache = None

    def __bool__(self) -> bool:
        return bool(self.expr)

    def check(self, **params: Any) -> bool:  # noqa: ANN401
        # prevent possible messing with globals
        assert set(params.keys()).issubset(self.allowed_params)

        try:
            namespace = dict(**params, **restricted_globals)
            return eval(self._checker, restricted_globals, namespace)  # noqa: S307
        except Exception:  # noqa: BLE001
            logger.debug("AimQL query evaluation failed", expr=self.expr)
            return False
