import ast
import json
import logging
from typing import Any, Callable, Dict, Iterator, List, Optional, Union

import pandas as pd
import pyarrow as pa
import pyarrow.compute as pc
import shapely.wkb
import shapely.wkt
from lark import Lark, Transformer, UnexpectedInput
from pydantic import BaseModel, PrivateAttr
from shapely.geometry.geo import shape


class Op(BaseModel):
    def __repr__(self):
        return f"〔{str(self)}〕"

    def bind(self, binds: Union[list, dict]):
        for op in self.all():
            if isinstance(op, BindVar):
                if op.key is None:
                    raise ValueError("incomplete bind variable")
                try:
                    op.val = binds[op.key]  # positional binds on list, named binds on dict
                    if not isinstance(op.val, Op):
                        op.val = op.val
                except KeyError:
                    raise ValueError(f"missing bind variable {op.key}")

    def all(self) -> Iterator["Op"]:
        """returns all the nodes in the tree"""
        yield self
        for attr in self.__dict__:
            value = getattr(self, attr)
            if isinstance(value, Op):
                yield from value.all()
            if isinstance(value, list):
                for v in value:
                    if isinstance(v, Op):
                        yield from v.all()

    # DEPRECATED: use op.all() instead
    def walk(self, fn: Callable[["Op"], "Op"]) -> "Op":
        """walks the tree and applies the given function to each node"""
        out = fn(self) or self

        for attr in out.__dict__:
            value = getattr(out, attr)
            if isinstance(value, Op):
                value = value.walk(fn)
                setattr(out, attr, value)
            if isinstance(value, list):
                value = [v.walk(fn) for v in value if isinstance(v, Op)]
                setattr(out, attr, value)

        return out

    def pyarrow(self) -> pc.Expression:
        raise NotImplementedError("pyarrow not implemented for %s" % self.__class__.__name__)

    def pandas(self, df: pd.DataFrame, schema: pa.Schema) -> pd.Series:
        raise NotImplementedError("pandas not implemented for %s" % self.__class__.__name__)


class Field(Op):
    _invocant: Optional[Op] = PrivateAttr(None)
    name: str

    def with_invocant(self, invocant: Op) -> "Field":
        self._invocant = invocant
        return self

    def __str__(self):
        if self._invocant:
            return f"{self._invocant}.{self.name}"
        else:
            return self.name

    def __repr__(self):
        if self._invocant:
            return f"〔{self._invocant.__repr__()} . {self.name}〕"
        else:
            return f"field({self.name})"

    def pyarrow(self) -> pc.Expression:
        if self._invocant is None:
            return pc.field(self.name)
        name = []
        cur = self
        while cur:
            if isinstance(cur, Field):
                name.append(cur.name)
                cur = cur._invocant
            else:
                raise ValueError(f"pyarrow unsupported {self}")
        return pc.field(".".join(reversed(name)))

    def pandas(self, df: pd.DataFrame, schema: pa.Schema) -> pd.Series:
        full = str(self)
        if full in df:
            return df[full]
        if self._invocant is not None:
            s = self._invocant.pandas(df, schema)
            return s.apply(lambda x: x[self.name])  # TODO(oha) is KeyError ok?
        raise ValueError(f"field not found: {self}")


class Func(Field):
    args: list[Op]

    def __str__(self):
        if self._invocant is None:
            return f"{self.name}({', '.join(map(str, self.args))})"
        return f"{super().__str__()}({', '.join(map(str, self.args))})"

    def __repr__(self):
        if self._invocant is None:
            return f"〔{self.name}({', '.join(map(repr, self.args))})〕"
        return f"〔{super().__repr__()}({', '.join(map(repr, self.args))})〕"

    def pyarrow(self) -> pc.Expression:
        args = [a.pyarrow() for a in self.args]
        if self._invocant is not None:
            raise ValueError(f"pyarrow unsupported {self}")
        return pc.field("")._call(self.name, args)

    def _arg_field(self, idx: int) -> Field:
        arg = self.args[idx]
        if isinstance(arg, Func):
            pass
        if isinstance(arg, Field):
            return arg
        raise ValueError(f"expected field, got {arg}")

    def _arg_scalar(self, idx: int) -> "Scalar":
        arg = self.args[idx]
        if isinstance(arg, Scalar):
            return arg
        raise ValueError(f"expected field, got {self.args[idx]}")

    def pandas(self, df: pd.DataFrame, schema: pa.Schema) -> pd.Series:
        # FIXME(oha): we need to create a context with the data, schema and registry, and add this there instead
        if self.name == "h3":
            import h3  # import late, so it doesn't need to be installed unless needed

            res = self._arg_scalar(0).to_int()
            lat, lon = "", ""
            for name in schema.names:
                field = schema.field(name)
                meta = field.metadata or {}
                if meta.get(b"geo") == b"latitude":
                    lat = name
                elif meta.get(b"geo") == b"longitude":
                    lon = name
                elif b"geo" in meta:
                    logging.warning("unexpected geo field: %s", meta)
            if not lat or not lon:
                raise ValueError('fields need to be marked as "geo":"latitude" and "geo":"longitude"')
            return df.apply(lambda row: h3.geo_to_h3(row[lat], row[lon], res), axis=1)
        else:
            raise NotImplementedError(f"not supported: {self}")


_op_map = {
    "AND": "and",
    "&": "and",
    "OR": "or",
    "|": "or",
    "=": "==",
}

_op_rev = {
    "==": "==",
    "!=": "!=",
    "<": ">=",
    "<=": ">",
    ">": "<=",
    ">=": "<",
}


class BinOp(Op):
    left: Op
    op: str
    right: Op

    def __init__(self, left: Op, op: str, right: Op):
        super().__init__(
            left=left,
            op=_op_map.get(op, op),
            right=right,
        )

    def flip(self) -> "BinOp":
        if self.op in _op_rev:
            return BinOp(self.right, _op_rev[self.op], self.left)
        raise NotImplementedError(f"Can't flip {self.op}")

    def __str__(self):
        return f"{self.left} {self.op} {self.right}"

    def __repr__(self):
        return f"〔{self.left.__repr__()} {self.op} {self.right.__repr__()}〕"

    def pyarrow(self) -> pc.Expression:
        left = self.left.pyarrow()
        right = self.right.pyarrow()
        if self.op == "<":
            return left < right
        if self.op == "<=":
            return left <= right
        if self.op == ">":
            return left > right
        if self.op == ">=":
            return left >= right
        if self.op == "==":
            return left == right
        if self.op == "!=":
            return left != right
        if self.op == "+":
            return left + right
        if self.op == "-":
            return left - right
        if self.op == "*":
            return left * right
        if self.op == "/":
            return left / right
        if self.op == "and":
            return pc.and_kleene(left, right)
        if self.op == "or":
            return pc.or_kleene(left, right)
        if self.op == "~":
            return pc.scalar(True)
        raise NotImplementedError(f"pyarrow not implemented for {self.op}")

    def pandas(self, df: pd.DataFrame, schema: pa.Schema) -> pd.Series:
        left = self.left.pandas(df, schema)
        right = self.right.pandas(df, schema)
        df = pd.DataFrame({"left": left, "right": right})
        if self.op == "==":
            return df.apply(lambda row: row["left"] == row["right"], axis=1)
        if self.op == "<":
            return df.apply(lambda row: row["left"] < row["right"], axis=1)
        if self.op == "<=":
            return df.apply(lambda row: row["left"] <= row["right"], axis=1)
        if self.op == ">":
            return df.apply(lambda row: row["left"] > row["right"], axis=1)
        if self.op == ">=":
            return df.apply(lambda row: row["left"] >= row["right"], axis=1)
        if self.op == "!=":
            return df.apply(lambda row: row["left"] != row["right"], axis=1)
        if self.op == "+":
            return df.apply(lambda row: row["left"] + row["right"], axis=1)
        if self.op == "-":
            return df.apply(lambda row: row["left"] - row["right"], axis=1)
        if self.op == "*":
            return df.apply(lambda row: row["left"] * row["right"], axis=1)
        if self.op == "/":
            return df.apply(lambda row: row["left"] / row["right"], axis=1)
        if self.op == "and":
            return df.apply(lambda row: row["left"] and row["right"], axis=1)
        if self.op == "or":
            return df.apply(lambda row: row["left"] or row["right"], axis=1)
        if self.op == "~":
            return df.apply(lambda row: True, axis=1)
        if self.op == "intersect" or self.op == "intersects":
            return df.apply(lambda row: _load_geo(row["left"]).intersects(_load_geo(row["right"])), axis=1)
        if self.op == "contain" or self.op == "contains":
            return df.apply(lambda row: _load_geo(row["left"]).contains(_load_geo(row["right"])), axis=1)
        if self.op == "within":
            return df.apply(lambda row: _load_geo(row["left"]).within(_load_geo(row["right"])), axis=1)

        raise NotImplementedError(f"pandas not implemented for {self.op}")


def _load_geo(val) -> shapely.geometry:
    if isinstance(val, str):
        return shapely.wkt.loads(val)
    if isinstance(val, bytes):
        return shapely.wkb.loads(val)
    raise ValueError(f"can't convert {type(val)} to geometry")


class UnaryOp(Op):
    prefix: str
    exp: Op
    suffix: str

    def __str__(self):
        return f"{self.prefix}{self.exp}{self.suffix}"

    def pyarrow(self):
        return ~(self.exp.pyarrow())

    @staticmethod
    def negate(op: Op) -> "UnaryOp":
        return UnaryOp(prefix="~", exp=op, suffix="")

    def pandas(self, df: pd.DataFrame, schema: pa.Schema) -> pd.Series:
        exp = self.exp.pandas(df, schema)
        if self.prefix == "~":
            return exp.apply(lambda x: not x)
        if self.prefix == "not":
            return exp.apply(lambda x: not x)
        raise NotImplementedError(f"pandas not implemented for {self.prefix}")


class BindVar(Op):
    key: Optional[Union[int, str]]  # we can bind by position or name
    val: Any

    def __str__(self):
        if isinstance(self.key, str):
            return f"${self.key}"  # by name
        return "?"  # positional

    def pyarrow(self) -> pc.Expression:
        if self.val is not None:
            return self.val
        raise ValueError("bind var not bound")

    def pandas(self, df: pd.DataFrame, schema: pa.Schema) -> pd.Series:
        if self.val is not None:
            return pd.Series([self.val] * len(df))
        raise ValueError("bind var not bound")


class Scalar(Op):
    src: str  # what was in the source code, e.g. '"123"' or 'True'
    type: str

    def __str__(self):
        return self.src

    @classmethod
    def from_py(cls, val):
        if val is None:
            return cls(src="null", type="null")
        if isinstance(val, bool):
            if val:
                return cls(src="true", type="bool")
            else:
                return cls(src="false", type="bool")
        if isinstance(val, (float, int)):
            return cls(src=str(val), type="SIGNED_NUMBER")
        if isinstance(val, str):
            return cls(src=json.dumps(val), type="ESCAPED_STRING")
        raise NotImplementedError(f"Scalar.from_py not implemented for {type(val)}")

    def to_int(self) -> int:
        if self.type in ["SIGNED_NUMBER"]:
            return int(self.src)
        raise ValueError(f"expected int, got: {self.src}")

    def to_string(self) -> str:
        if self.type in ["ESCAPED_STRING"]:
            return self.src
        raise ValueError(f"expected string, got: {self.src}")

    def to_py(self):
        if self.type in ["SIGNED_NUMBER"]:
            return float(self.src)
        if self.type in ["ESCAPED_STRING"]:
            return json.loads(self.src)
        if self.type == "bool":
            if self.src.lower() == "true":
                return True
            if self.src.lower() == "false":
                return False
            raise ValueError("unexpected bool value: %s" % self.src)
        if self.type == "null":
            return None
        try:
            return ast.literal_eval(self.src)  # raises ValueError
        except Exception:
            pass
        raise ValueError(f"unexpected scalar value: {self.src}")

    def pyarrow(self) -> pc.Expression:
        if self.type in ["SIGNED_NUMBER"]:
            if "." in self.src:
                return pc.scalar(float(self.src))
            else:
                return pc.scalar(int(self.src))
        if self.type in ["ESCAPED_STRING"]:
            return pc.scalar(json.loads(self.src))
        if self.type == "bool":
            if self.src.lower() == "true":
                return pc.scalar(True)
            if self.src.lower() == "false":
                return pc.scalar(False)
            raise ValueError("unexpected bool value: %s" % self.src)
        if self.type == "null":
            return pc.scalar(None)
        try:
            return pc.scalar(ast.literal_eval(self.src))  # raises ValueError
        except ValueError:
            raise ValueError(f"unexpected scalar value: {self.src}")

    def pandas(self, df: pd.DataFrame, schema: pa.Schema) -> pd.Series:
        return pd.Series([self.to_py()] * len(df))  # TODO(oha) is this smart? or should we return a scalar?


class Parens(Op):
    exp: Op

    def __init__(self, exp: Op):
        super().__init__(exp=exp)

    def __str__(self):
        return f"({self.exp})"

    def pyarrow(self) -> pc.Expression:
        return self.exp.pyarrow()

    def pandas(self, df: pd.DataFrame, schema: pa.Schema) -> pd.Series:
        return self.exp.pandas(df, schema)


class _T(Transformer):
    def exp(self, tok):
        return tok[0]

    def func(self, tok):
        name = tok[0].value
        args = tok[1:]
        if args == [None]:  # no args
            args = []
        if name == "add_checked" or name == "add":
            assert len(args) == 2
            return Parens(BinOp(left=args[0], op="+", right=args[1]))
        if name == "subtract_checked":
            assert len(args) == 2
            return Parens(BinOp(left=args[0], op="-", right=args[1]))
        return Func(name=name, args=args or [])

    def scalar(self, tok):
        return Scalar(src=tok[0].value, type=tok[0].type)

    def bind_var(self, tok):
        if len(tok) == 1:
            return BindVar(key=tok[0].value)
        return BindVar(val=None)

    def field(self, tok):
        name = tok[0].value
        if name in ["True", "False", "true", "false"]:
            return Scalar(src=name, type="bool")
        if name in ["NULL", "null", "None"]:
            return Scalar(src="null", type="null")
        return Field(name=name)

    def member(self, tok):
        member: Field = tok[1]
        member = member.with_invocant(tok[0])
        return member

    def parens(self, tok):
        return Parens(exp=tok[0])

    def mul(self, tok):
        return BinOp(left=tok[0], op=tok[1], right=tok[2])

    def add(self, tok):
        return BinOp(left=tok[0], op=tok[1], right=tok[2])

    def cmp(self, tok):
        return BinOp(left=tok[0], op=tok[1], right=tok[2])

    def cmp2(self, tok):
        return BinOp(left=tok[0], op=tok[1], right=tok[2])

    def bool(self, tok):
        return BinOp(left=tok[0], op=tok[1], right=tok[2])

    def notop(self, tok: Any):
        return UnaryOp(prefix="~", exp=tok[0], suffix="")

    def is_null(self, tok):
        return Func(name="is_null", args=[tok[0]])

    def is_not_null(self, tok):
        return UnaryOp(prefix="~", exp=self.is_null(tok), suffix="")


_parser = Lark(
    """

?exp: bool

parens: "(" bool ")"

?bool: (bool /(and|or|AND|OR)/)? cmp

?cmp: (cmp /(==|!=|~|!~|intersect|within|contains)/)? cmp2
    | is_null
    | is_not_null

is_null: add "is" "null"
       | add "is" "None"
is_not_null: add "is" "not" "null"
           | add "is" "not" "None"

?cmp2: (cmp2 /(<=|<|>=|>)/)? add

?add: (add /(\\+|-)/)? mul

?mul: (mul /(\\*|\\/|%)/)? term

?pow: term ("**" pow)?

notop: "~" term
     | "not" term

?term: notop
     | func
     | bind_var
     | scalar
     | member
     | field
     | parens

scalar: SIGNED_NUMBER
      | /(True|False|None|null|true|false)/
      | ESCAPED_STRING
      | /'[^']*'/

bind_var: "?"
      | "$" NAME

member: (parens | member | field | func) "." (func | field)

field: NAME

func: NAME "(" [exp ("," exp)*] ")"

NAME: /[a-zA-Z_][a-zA-Z0-9_]*/

%import common.ESCAPED_STRING
%import common.SIGNED_NUMBER
%import common.WS

%ignore WS
""",
    start="exp",
    transformer=_T(),
    parser="lalr",
    debug=True,
)


class ExpressionError(Exception):
    def __init__(self, msg: str, at: str):
        super().__init__(msg)
        self.at = at


def parse(s: str) -> Op:
    """parses the given string into an Op tree"""
    try:
        op: Op = _parser.parse(s)  # noqa
        i = 0
        for node in op.all():
            if isinstance(node, BindVar):
                if node.key is None:
                    node.key = i  # set the bind_var position index
                    i += 1
        return op
    except UnexpectedInput as e:
        at = e.get_context(s)
        raise ExpressionError("Can't parse:\n" + at + "\n" + str(e), at) from e


def parse_oqs_dict(oqs: Dict[str, Union[str, List[Any]]]) -> Op:
    """parses the given JSON string into an Op tree"""
    # assert len(oqs) == 1
    if not oqs:
        return Scalar.from_py(None)
    if isinstance(oqs, dict):
        operation, operands = next(iter(oqs.items()))
        operation = operation.lower()
        if not operation.startswith("#"):
            return Scalar.from_py(shape(oqs).wkt)
        if operation == "#constant":
            return Scalar.from_py(operands[0])
        if operation == "#ref":
            return Field(name=operands[0])
        if operation == "#list":
            return [parse_oqs_dict(v) for v in operands]
        if operation == "#equals":
            return BinOp(left=parse_oqs_dict(operands[0]), op="==", right=parse_oqs_dict(operands[1]))
        if operation == "#not_equals":
            return BinOp(left=parse_oqs_dict(operands[0]), op="!=", right=parse_oqs_dict(operands[1]))
        if operation == "#greater_than":
            return BinOp(left=parse_oqs_dict(operands[0]), op=">", right=parse_oqs_dict(operands[1]))
        if operation == "#greater_than_or_equals":
            return BinOp(left=parse_oqs_dict(operands[0]), op=">=", right=parse_oqs_dict(operands[1]))
        if operation == "#less_than":
            return BinOp(left=parse_oqs_dict(operands[0]), op="<", right=parse_oqs_dict(operands[1]))
        if operation == "#less_than_or_equals":
            return BinOp(left=parse_oqs_dict(operands[0]), op="<=", right=parse_oqs_dict(operands[1]))
        if operation == "#and":
            op_list = [parse_oqs_dict(v) for v in operands]
            return list_to_tree_expression(op_list, "and")
        if operation == "#or":
            op_list = [parse_oqs_dict(v) for v in operands]
            return list_to_tree_expression(op_list, "or")
        if operation == "#xor":
            raise ValueError("xor not supported in tabular")
        if operation == "#within":
            right_op = parse_oqs_dict(operands[1])
            if isinstance(right_op, list):
                op_list = [BinOp(left=parse_oqs_dict(operands[0]), op="==", right=operand) for operand in right_op]
                return list_to_tree_expression(op_list, "or")
            return BinOp(left=parse_oqs_dict(operands[0]), op="==", right=right_op)
        if operation == "#negate":
            return UnaryOp.negate(parse_oqs_dict(operands[0]))
        if operation == "#non_null":
            return BinOp(left=parse_oqs_dict(operands[0]), op="!=", right=Scalar.from_py(None))
        if operation == "#null":
            return BinOp(left=parse_oqs_dict(operands[0]), op="==", right=Scalar.from_py(None))
        if operation == "#true":
            return Scalar.from_py(True)
        if operation == "#false":
            return Scalar.from_py(False)
        if operation == "#sum":
            return Func(name="sum", args=[parse_oqs_dict(operands[0])])
        if operation == "#difference":
            return Func(name="difference", args=[parse_oqs_dict(operands[0])])
        if operation == "#product":
            return Func(name="product", args=[parse_oqs_dict(operands[0])])
        if operation == "#quotient" or operation == "#floor_quotient":
            return BinOp(left=parse_oqs_dict(operands[0]), op="/", right=parse_oqs_dict(operands[1]))
        if operation == "#modulo":
            return BinOp(left=parse_oqs_dict(operands[0]), op="%", right=parse_oqs_dict(operands[1]))
        if operation == "#exponentiation":
            return BinOp(left=parse_oqs_dict(operands[0]), op="**", right=parse_oqs_dict(operands[1]))
        if operation == "#st_within":
            return Func(name="ktable_st_within", args=[parse_oqs_dict(operands[0]), parse_oqs_dict(operands[1])])
        if operation == "#st_equals":
            return Func(name="ktable_st_equals", args=[parse_oqs_dict(operands[0]), parse_oqs_dict(operands[1])])
        if operation == "#st_intersects":
            return Func(name="ktable_st_intersects", args=[parse_oqs_dict(operands[0]), parse_oqs_dict(operands[1])])
        if operation == "#st_contains":
            return Func(name="ktable_st_contains", args=[parse_oqs_dict(operands[0]), parse_oqs_dict(operands[1])])
        # TODO: Decide how to handle times, timestamps and dates
        raise ValueError(f"unexpected key {operation}")

    if isinstance(oqs, str):
        if oqs.startswith("$"):
            return Field(name=oqs.removeprefix("$"))

    try:
        return Scalar.from_py(oqs)
    except NotImplementedError:
        raise ValueError(f"unexpected value {oqs} with type {type(oqs)}")


def list_to_tree_expression(list_expressions: List[Op], operation: str) -> Op:
    if len(list_expressions) == 1:
        return list_expressions[0]
    else:
        temp = []
        for i in range(0, len(list_expressions), 2):
            if i + 1 == len(list_expressions):
                temp.append(list_expressions[i])
            else:
                temp.append(BinOp(left=list_expressions[i], op=operation, right=list_expressions[i + 1]))
        return list_to_tree_expression(temp, operation)
