import ast
import math
from collections.abc import Callable


def _new(expr: "_Expr", ast: ast.expr) -> "_Expr":
    if isinstance(expr, _ParamAccessor):
        return _Expr(ast)
    else:
        return type(expr)(ast)


def _ast(expr) -> ast.expr:
    return object.__getattribute__(expr, "_ast")


def _auto_ast(node) -> ast.expr:
    if isinstance(node, _Expr):
        return _ast(node)
    else:
        return ast.Constant(value=node)


def _unary(op: ast.unaryop) -> "_Expr":
    return lambda self: _new(self, ast.UnaryOp(op=op, operand=_ast(self)))


def _unary_fn(fn: Callable) -> "_Expr":
    return lambda self: _new(
        self,
        ast.Call(
            func=ast.Constant(value=fn),
            args=[_ast(self)],
            keywords=[],
        ),
    )


def _binary(op: ast.operator):
    def left(self, other) -> "_Expr":
        return _new(self, ast.BinOp(left=_ast(self), op=op, right=_auto_ast(other)))

    def right(self, other) -> "_Expr":
        return _new(self, ast.BinOp(left=_auto_ast(other), op=op, right=_ast(self)))

    return left, right


def _binary_fn(fn: Callable):
    def left(self, other) -> "_Expr":
        return _new(
            self,
            ast.Call(
                func=ast.Constant(value=fn),
                args=[_ast(self), _auto_ast(other)],
                keywords=[],
            ),
        )

    def right(self, other) -> "_Expr":
        return _new(
            self,
            ast.Call(
                func=ast.Constant(value=fn),
                args=[_auto_ast(other), _ast(self)],
                keywords=[],
            ),
        )

    return left, right


class _Expr:
    def __init__(self, ast: ast.expr):
        self._ast = ast

    def __getattribute__(self, attr: str) -> "_Expr":
        if attr.startswith("__") and attr.endswith("__"):
            return object.__getattribute__(self, attr)
        return _new(self, ast.Attribute(value=_ast(self), attr=attr, ctx=ast.Load()))

    def __getitem__(self, key) -> "_Expr":
        return _new(
            self,
            ast.Subscript(
                value=_ast(self),
                slice=_auto_ast(key),
                ctx=ast.Load(),
            ),
        )

    def __call__(self, *args, **kwargs) -> "_Expr":
        return _new(
            self,
            ast.Call(
                func=_ast(self),
                args=[_auto_ast(arg) for arg in args],
                keywords=[
                    ast.keyword(arg=key, value=_auto_ast(value))
                    for key, value in kwargs.items()
                ],
            ),
        )

    def __round__(self, ndigits=None) -> "_Expr":
        return _new(
            self,
            ast.Call(
                func=ast.Name(id="round", ctx=ast.Load()),
                args=[_ast(self), _auto_ast(ndigits)],
                keywords=[],
            ),
        )

    __neg__ = _unary(ast.USub())
    __pos__ = _unary(ast.UAdd())
    __abs__ = _unary_fn(abs)
    __invert__ = _unary(ast.Invert())

    __reversed__ = _unary_fn(reversed)

    __trunc__ = _unary_fn(math.trunc)
    __floor__ = _unary_fn(math.floor)
    __ceil__ = _unary_fn(math.ceil)

    __add__, __radd__ = _binary(ast.Add())
    __sub__, __rsub__ = _binary(ast.Sub())
    __mul__, __rmul__ = _binary(ast.Mult())
    __matmul__, __rmatmul__ = _binary(ast.MatMult())
    __truediv__, __rtruediv__ = _binary(ast.Div())
    __floordiv__, __rfloordiv__ = _binary(ast.FloorDiv())
    __mod__, __rmod__ = _binary(ast.Mod())
    __divmod__, __rdivmod__ = _binary_fn(divmod)
    __pow__, __rpow__ = _binary(ast.Pow())
    __lshift__, __rlshift__ = _binary(ast.LShift())
    __rshift__, __rrshift__ = _binary(ast.RShift())
    __and__, __rand__ = _binary(ast.BitAnd())
    __xor__, __rxor__ = _binary(ast.BitXor())
    __or__, __ror__ = _binary(ast.BitOr())

    __eq__ = _binary(ast.Eq())[0]
    __ne__ = _binary(ast.NotEq())[0]
    __lt__ = _binary(ast.Lt())[0]
    __gt__ = _binary(ast.Gt())[0]
    __le__ = _binary(ast.LtE())[0]
    __ge__ = _binary(ast.GtE())[0]


class _ParamAccessor(_Expr):
    def __init__(self, key: ast.expr | int | str):
        super().__init__(ast.Constant(value=self))
        self._key = key


class _ParamAccessorInterface:
    def __getitem__(self, key):
        if not isinstance(key, (int, str)):
            raise TypeError("Parameter key must be int or str")
        return _ParamAccessor(key)


P = _ParamAccessorInterface()


class _Compiler(ast.NodeTransformer):
    def __init__(self):
        self.arg_counter = 0
        self.max_args = 0
        self.kwargs = set()
        self.constants = []

    def visit_Name(self, node: ast.Name):
        if node.id == "$":
            name = f"${self.arg_counter}"
            self.arg_counter += 1
            return ast.Name(id=name, ctx=ast.Load())
        else:
            return node

    def visit_Constant(self, node: ast.Constant):
        if isinstance(node.value, _ParamAccessor):
            key = object.__getattribute__(node.value, "_key")
            if isinstance(key, int):
                idx = key
                self.max_args = max(self.max_args, idx + 1)
                name = f"${idx}"
            else:
                self.kwargs.add(str(key))
                name = str(key)
            return ast.Name(id=name, ctx=ast.Load())
        elif not isinstance(node.value, (str, int, float, bool, type(None))):
            const_idx = len(self.constants)
            self.constants.append(node.value)
            return ast.Name(id=f"#{const_idx}", ctx=ast.Load())
        return node


def _compile(expr: _Expr) -> Callable:
    compiler = _Compiler()
    expr = compiler.visit(_ast(expr))
    arg_count = max(compiler.max_args, compiler.arg_counter)
    lmbd = ast.Lambda(
        args=ast.arguments(
            posonlyargs=[ast.arg(arg=f"${i}") for i in range(arg_count)],
            args=[],
            kwonlyargs=[ast.arg(arg=key) for key in compiler.kwargs],
            kw_defaults=[None] * len(compiler.kwargs),
            vararg=None,
            kwarg=None,
            defaults=[],
        ),
        body=expr,
    )
    output = ast.Expression(body=lmbd)
    ast.fix_missing_locations(output)
    code = compile(output, filename="<lmbd>", mode="eval")
    return eval(
        code, {f"#{i}": value for i, value in enumerate(compiler.constants)}, {}
    )


def F(expr) -> Callable:
    if isinstance(expr, _Expr):
        return _compile(expr)
    elif callable(expr):
        return expr
    else:
        raise TypeError("F() argument must be an expression or a callable")


_ = _Expr(ast.Name(id="$", ctx=ast.Load()))

__all__ = ["F", "P", "_"]
