"""Expression serialization and deserialization.

This module provides functionality to serialize expression objects to strings
and deserialize strings back into expression objects using a grammar-based parser.
"""

from __future__ import annotations

import re
from collections.abc import Callable
from typing import Any
from urllib.parse import quote, unquote

import tatsu  # type: ignore[import-untyped]
from tatsu.ast import AST  # type: ignore[import-untyped]

from therismos.expr._expr import (
    AllExpr,
    AnyExpr,
    Eq,
    Expr,
    ExprVisitor,
    F,
    FalseExpr,
    Field,
    Ge,
    Gt,
    In,
    IsNull,
    Le,
    Lt,
    Ne,
    NotExpr,
    Regex,
    TrueExpr,
)

# TatSu-compatible grammar
_GRAMMAR = r"""
@@grammar::CQL
@@ignorecase::False
@@whitespace :: /[ \t\r\n]+/
@@comments :: /#.*?$/

start
    =
    expr
    $
    ;

# -------------------------
# Precedence: NOT > AND > OR
# -------------------------

expr
    =
    or_expr
    ;

or_expr
    =
    and_expr {',' and_expr}*
    ;

and_expr
    =
    not_expr {';' not_expr}*
    ;

not_expr
    =
      '!' not_expr
    | primary
    ;

primary
    =
      atom
    | '(' expr ')'
    ;

# -------------------------
# Atoms
# -------------------------

atom
    =
      boolean_const
    | null_check
    | regex_expr
    | in_expr
    | comparison
    ;

boolean_const
    =
      'true()'
    | 'false()'
    ;

null_check
    =
    field:field null_op:null_op 'null'
    ;

null_op
    =
      '=='
    | '!='
    ;

comparison
    =
    field:field comp_op:comp_op value:value
    ;

comp_op
    =
      '=='
    | '!='
    | '<='
    | '>='
    | '<'
    | '>'
    ;

in_expr
    =
    field:field '=in=' '(' values:value_list ')'
    ;

value_list
    =
    value {',' value}*
    ;

regex_expr
    =
    field:field '~regex' '(' strings:string_list ')'
    ;

string_list
    =
    string [',' string]
    ;

# -------------------------
# Field with optional type annotation
# -------------------------

field
    =
    ident:ident [type_annot:type_annot]
    ;

type_annot
    =
    '{' type_name:type_name '}'
    ;

type_name
    =
    ident
    ;

# -------------------------
# Values
# -------------------------

value
    =
      string
    | number
    | boolean
    | 'null'
    | ident
    ;

boolean
    =
      'true'
    | 'false'
    ;

# -------------------------
# Lexical rules
# -------------------------

ident
    =
    /[A-Za-z_][A-Za-z0-9_\.]*/
    ;

number
    =
    /-?\d+(?:\.\d+)?/
    ;

# Double-quoted string with JSON-like escapes:
# \" \\ \n \t \r
string
    =
    /"(?:\\.|[^"\\])*"/
    ;
"""

# Compile the grammar once at module load time
_PARSER = tatsu.compile(_GRAMMAR)


class _SerializerVisitor(ExprVisitor[str]):
    """A visitor that converts expression objects into grammar-compliant strings."""

    def __init__(
        self,
        type_registry: dict[type | Callable[..., Any], str],
        include_all_types: bool,
        url_encode: bool,
    ) -> None:
        """Initialize the serializer visitor.

        :param type_registry: A mapping from Python types to their string representations.
        :type type_registry: dict[type | Callable[..., Any], str]
        :param include_all_types: If ``True``, include type annotations for all fields.
        :type include_all_types: bool
        :param url_encode: If ``True``, URL-encode the output.
        :type url_encode: bool
        """
        self.type_registry = type_registry
        self.include_all_types = include_all_types
        self.url_encode = url_encode

    def _escape_string(self, s: str) -> str:
        """Escape a string for grammar output.

        :param s: The string to be escaped.
        :type s: str
        :return: The escaped and quoted string.
        :rtype: str
        """
        s = (
            s.replace("\\", "\\\\")
            .replace('"', '\\"')
            .replace("\n", "\\n")
            .replace("\t", "\\t")
            .replace("\r", "\\r")
        )
        return f'"{s}"'

    def _serialize_value(self, value: Any) -> str:
        """Serialize a value into its grammar format.

        :param value: The value to be serialized.
        :type value: Any
        :return: The string representation of the value.
        :rtype: str
        """
        if value is None:
            return "null"
        if isinstance(value, bool):
            return "true" if value else "false"
        if isinstance(value, str):
            return self._escape_string(value)
        if isinstance(value, (int, float)):
            return str(value)
        return self._escape_string(str(value))

    def _format_field(self, field: Any) -> str:
        """Format a field with an optional type annotation.

        :param field: The field object.
        :type field: Any
        :return: The formatted field string.
        :rtype: str
        """
        field_name = field.name
        type_annot = ""

        if field.type_ is not None and self.include_all_types:
            type_name = self.type_registry.get(field.type_, str(field.type_))
            type_annot = f"{{{type_name}}}"

        return f"{field_name}{type_annot}"

    def visit_eq(self, expr: Eq) -> str:
        """Visit an equality expression.

        :param expr: The equality expression.
        :type expr: Eq
        :return: The string representation.
        :rtype: str
        """
        field = self._format_field(expr.field)
        value = self._serialize_value(expr.casted_value())
        return f"{field}=={value}"

    def visit_ne(self, expr: Ne) -> str:
        """Visit an inequality expression.

        :param expr: The inequality expression.
        :type expr: Ne
        :return: The string representation.
        :rtype: str
        """
        field = self._format_field(expr.field)
        value = self._serialize_value(expr.casted_value())
        return f"{field}!={value}"

    def visit_lt(self, expr: Lt) -> str:
        """Visit a less-than expression.

        :param expr: The less-than expression.
        :type expr: Lt
        :return: The string representation.
        :rtype: str
        """
        field = self._format_field(expr.field)
        value = self._serialize_value(expr.casted_value())
        return f"{field}<{value}"

    def visit_le(self, expr: Le) -> str:
        """Visit a less-than-or-equal expression.

        :param expr: The less-than-or-equal expression.
        :type expr: Le
        :return: The string representation.
        :rtype: str
        """
        field = self._format_field(expr.field)
        value = self._serialize_value(expr.casted_value())
        return f"{field}<={value}"

    def visit_gt(self, expr: Gt) -> str:
        """Visit a greater-than expression.

        :param expr: The greater-than expression.
        :type expr: Gt
        :return: The string representation.
        :rtype: str
        """
        field = self._format_field(expr.field)
        value = self._serialize_value(expr.casted_value())
        return f"{field}>{value}"

    def visit_ge(self, expr: Ge) -> str:
        """Visit a greater-than-or-equal expression.

        :param expr: The greater-than-or-equal expression.
        :type expr: Ge
        :return: The string representation.
        :rtype: str
        """
        field = self._format_field(expr.field)
        value = self._serialize_value(expr.casted_value())
        return f"{field}>={value}"

    def visit_in(self, expr: In) -> str:
        """Visit an IN expression.

        :param expr: The IN expression.
        :type expr: In
        :return: The string representation.
        :rtype: str
        """
        field = self._format_field(expr.field)
        values = ",".join(self._serialize_value(v) for v in expr.casted_values())
        return f"{field}=in=({values})"

    def visit_regex(self, expr: Regex) -> str:
        """Visit a regex expression.

        :param expr: The regex expression.
        :type expr: Regex
        :return: The string representation.
        :rtype: str
        """
        field = self._format_field(expr.field)
        pattern = self._escape_string(expr.value)
        if expr.flags is not None:
            flags = self._escape_string(str(expr.flags))
            return f"{field}~regex({pattern},{flags})"
        return f"{field}~regex({pattern})"

    def visit_is_null(self, expr: IsNull) -> str:
        """Visit an IS NULL expression.

        :param expr: The IS NULL expression.
        :type expr: IsNull
        :return: The string representation.
        :rtype: str
        """
        field = self._format_field(expr.field)
        op = "==" if expr.is_null else "!="
        return f"{field}{op}null"

    def visit_all(self, expr: AllExpr) -> str:
        """Visit an AND expression.

        :param expr: The AND expression.
        :type expr: AllExpr
        :return: The string representation.
        :rtype: str
        """
        if not expr.exprs:
            return "true()"
        parts = [e.accept(self) for e in expr.exprs]
        return f"({';'.join(parts)})"

    def visit_any(self, expr: AnyExpr) -> str:
        """Visit an OR expression.

        :param expr: The OR expression.
        :type expr: AnyExpr
        :return: The string representation.
        :rtype: str
        """
        if not expr.exprs:
            return "false()"
        parts = [e.accept(self) for e in expr.exprs]
        return f"({','.join(parts)})"

    def visit_not(self, expr: NotExpr) -> str:
        """Visit a NOT expression.

        :param expr: The NOT expression.
        :type expr: NotExpr
        :return: The string representation.
        :rtype: str
        """
        inner = expr.expr.accept(self)
        return f"!{inner}" if isinstance(expr.expr, (AllExpr, AnyExpr)) else f"!({inner})"

    def visit_true(self, expr: TrueExpr) -> str:
        """Visit a TRUE constant expression.

        :param expr: The TRUE expression.
        :type expr: TrueExpr
        :return: The string representation.
        :rtype: str
        """
        return "true()"

    def visit_false(self, expr: FalseExpr) -> str:
        """Visit a FALSE constant expression.

        :param expr: The FALSE expression.
        :type expr: FalseExpr
        :return: The string representation.
        :rtype: str
        """
        return "false()"


class _DeserializerSemantics:
    """A TatSu semantics class to convert a parsed AST into expression objects."""

    def __init__(
        self,
        type_registry: dict[str, type | Callable[..., Any]],
        implicit_field_types: dict[str, type | Callable[..., Any]],
        sanitizer: Callable[[str, Any, type | Callable[..., Any] | None], Any] | None = None,
    ) -> None:
        """Initialize the deserializer semantics.

        :param type_registry: A mapping from type name strings to Python types.
        :type type_registry: dict[str, type | Callable[..., Any]]
        :param implicit_field_types: A mapping from field names to their implicit types.
        :type implicit_field_types: dict[str, type | Callable[..., Any]]
        :param sanitizer: An optional callback to sanitize values.
                          It receives the field name, value, and expected type,
                          and should return the sanitized value.
        :type sanitizer: Callable[[str, Any, type | Callable[..., Any] | None], Any] | None
        """
        self.type_registry = type_registry
        self.implicit_field_types = implicit_field_types
        self.sanitizer = sanitizer

    def _parse_field(self, ast: AST) -> Field:
        """Parse a field from an AST node.

        :param ast: The AST node for the field.
        :type ast: AST
        :return: A :class:`Field` object.
        :rtype: Any
        """
        field_name = ast.ident if hasattr(ast, "ident") else ast["ident"]
        field_type = None

        type_annot = ast.type_annot if hasattr(ast, "type_annot") else ast.get("type_annot")
        if type_annot:
            type_name = (
                type_annot.type_name
                if hasattr(type_annot, "type_name")
                else type_annot["type_name"]
            )
            field_type = self.type_registry.get(type_name)
        elif field_name in self.implicit_field_types:
            field_type = self.implicit_field_types[field_name]

        return F(field_name, field_type)

    def _parse_value(self, val: Any) -> Any:
        """Parse a value from the AST.

        :param val: The value from the AST.
        :type val: Any
        :return: The parsed Python value.
        :rtype: Any
        """
        if isinstance(val, str):
            if val.startswith('"'):
                s = val[1:-1]
                s = (
                    s.replace("\\n", "\n")
                    .replace("\\t", "\t")
                    .replace("\\r", "\r")
                    .replace('\\"', '"')
                    .replace("\\\\", "\\")
                )
                return s
            if val == "null":
                return None
            if val == "true":
                return True
            if val == "false":
                return False
            if re.match(r"^-?\d+(?:\.\d+)?$", val):
                return float(val) if "." in val else int(val)
            return val
        return val

    def boolean_const(self, ast: str) -> TrueExpr | FalseExpr:
        """Parse a boolean constant.

        :param ast: The AST node.
        :type ast: str
        :return: A :class:`TrueExpr` or :class:`FalseExpr`.
        :rtype: TrueExpr | FalseExpr
        """
        return TrueExpr() if ast == "true()" else FalseExpr()

    def null_check(self, ast: AST) -> IsNull:
        """Parse a null check expression.

        :param ast: The AST node.
        :type ast: AST
        :return: An :class:`IsNull` expression.
        :rtype: IsNull
        """
        field = self._parse_field(ast.field if hasattr(ast, "field") else ast["field"])
        null_op = ast.null_op if hasattr(ast, "null_op") else ast["null_op"]
        return IsNull(field, is_null=null_op == "==")

    def comparison(self, ast: AST) -> Eq | Ne | Lt | Le | Gt | Ge:
        """Parse a comparison expression.

        :param ast: The AST node.
        :type ast: AST
        :return: A comparison expression object.
        :rtype: Eq | Ne | Lt | Le | Gt | Ge
        """
        field = self._parse_field(ast.field if hasattr(ast, "field") else ast["field"])
        op = ast.comp_op if hasattr(ast, "comp_op") else ast["comp_op"]
        value = self._parse_value(ast.value if hasattr(ast, "value") else ast["value"])

        if self.sanitizer:
            value = self.sanitizer(field.name, value, field.type_)
        if field.type_:
            value = field.cast(value)

        if op == "==":
            return Eq(field, value)
        elif op == "!=":
            return Ne(field, value)
        elif op == "<":
            return Lt(field, value)
        elif op == "<=":
            return Le(field, value)
        elif op == ">":
            return Gt(field, value)
        elif op == ">=":
            return Ge(field, value)
        raise ValueError(f"Unknown comparison operator: {op}")

    def in_expr(self, ast: AST) -> In:
        """Parse an IN expression.

        :param ast: The AST node.
        :type ast: AST
        :return: An :class:`In` expression.
        :rtype: In
        """
        field = self._parse_field(ast.field if hasattr(ast, "field") else ast["field"])
        values_raw = ast.get("values") or ast.get("values_") or ast["values"]

        values = []
        if isinstance(values_raw, (list, tuple)) and len(values_raw) == 2:
            first_val, rest_vals = values_raw
            values.append(self._parse_value(first_val))
            if rest_vals:
                for sep_val_pair in rest_vals:
                    if isinstance(sep_val_pair, list) and len(sep_val_pair) == 2:
                        values.append(self._parse_value(sep_val_pair[1]))
        else:
            values.append(self._parse_value(values_raw))

        if self.sanitizer:
            values = [self.sanitizer(field.name, v, field.type_) for v in values]
        if field.type_:
            values = [field.cast(v) for v in values]

        return In(field, tuple(values))

    def regex_expr(self, ast: AST) -> Regex:
        """Parse a regex expression.

        :param ast: The AST node.
        :type ast: AST
        :return: A :class:`Regex` expression.
        :rtype: Regex
        """
        field = self._parse_field(ast.field if hasattr(ast, "field") else ast["field"])
        strings_raw = ast.get("strings") or ast["strings"]

        pattern, flags = None, None
        if isinstance(strings_raw, str):
            pattern = self._parse_value(strings_raw)
        elif isinstance(strings_raw, (list, tuple)):
            pattern = self._parse_value(strings_raw[0])
            if len(strings_raw) >= 3:
                try:
                    flags = int(self._parse_value(strings_raw[2]))
                except (ValueError, TypeError):
                    flags = None

        return Regex(field, pattern, flags)

    def not_expr(self, ast: Any) -> Expr:
        """Parse a NOT expression.

        :param ast: The AST node.
        :type ast: Any
        :return: A :class:`NotExpr` or the inner expression.
        :rtype: Expr
        """
        flat = self._flatten_primary(ast)
        if isinstance(flat, tuple) and len(flat) == 2 and flat[0] == "!":
            return NotExpr(flat[1])
        return ast  # type: ignore[no-any-return]

    def _flatten_primary(self, ast: Any) -> Any:
        """Flatten primary expression parentheses."""
        if isinstance(ast, (tuple, list)):
            if len(ast) == 3 and ast[0] == "(" and ast[2] == ")":
                return ast[1]
            if len(ast) == 2 and ast[0] == "!":
                return "!", self._flatten_primary(ast[1])
        return ast

    def primary(self, ast: Any) -> Expr:
        """Handle primary expressions.

        :param ast: The AST node.
        :type ast: Any
        :return: The expression.
        :rtype: Expr
        """
        return self._flatten_primary(ast)  # type: ignore[no-any-return]

    def and_expr(self, ast: Any) -> Expr:
        """Parse an AND expression.

        :param ast: The AST node.
        :type ast: Any
        :return: An :class:`AllExpr` or a single expression.
        :rtype: Expr
        """
        if isinstance(ast, (list, tuple)) and len(ast) == 2:
            first, rest_parts = ast
            if not rest_parts:
                return first  # type: ignore[no-any-return]
            exprs = [first] + [
                pair[1] for pair in rest_parts if isinstance(pair, list) and len(pair) == 2
            ]
            return AllExpr(*exprs)
        return ast  # type: ignore[no-any-return]

    def or_expr(self, ast: Any) -> Expr:
        """Parse an OR expression.

        :param ast: The AST node.
        :type ast: Any
        :return: An :class:`AnyExpr` or a single expression.
        :rtype: Expr
        """
        if isinstance(ast, (list, tuple)) and len(ast) == 2:
            first, rest_parts = ast
            if not rest_parts:
                return first  # type: ignore[no-any-return]
            exprs = [first] + [
                pair[1] for pair in rest_parts if isinstance(pair, list) and len(pair) == 2
            ]
            return AnyExpr(*exprs)
        return ast  # type: ignore[no-any-return]


class Serializer:
    """A serializer for converting expressions to and from string representations."""

    _DEFAULT_TYPE_REGISTRY: dict[type | Callable[..., Any], str] = {
        int: "int",
        float: "float",
        str: "str",
        bool: "bool",
    }

    def __init__(
        self,
        type_registry: dict[type | Callable[..., Any], str] | None = None,
        implicit_field_types: dict[str, type | Callable[..., Any]] | None = None,
        include_all_types: bool = False,
        url_encode: bool = False,
        sanitizer: Callable[[str, Any, type | Callable[..., Any] | None], Any] | None = None,
    ) -> None:
        """Initialize the serializer.

        :param type_registry: A custom registry for mapping types to names.
        :type type_registry: dict[type | Callable[..., Any], str] | None
        :param implicit_field_types: A mapping of field names to their implicit types.
        :type implicit_field_types: dict[str, type | Callable[..., Any]] | None
        :param include_all_types: If ``True``, include type annotations for all fields.
        :type include_all_types: bool
        :param url_encode: If ``True``, URL-encode serialized output and decode input.
        :type url_encode: bool
        :param sanitizer: An optional callback to sanitize values during deserialization.
        :type sanitizer: Callable[[str, Any, type | Callable[..., Any] | None], Any] | None
        """
        self.type_registry = {**self._DEFAULT_TYPE_REGISTRY, **(type_registry or {})}
        self.implicit_field_types = implicit_field_types or {}
        self.include_all_types = include_all_types
        self.url_encode = url_encode
        self.sanitizer = sanitizer
        self._reverse_type_registry = {v: k for k, v in self.type_registry.items()}

    def register_custom_type(self, type_obj: type | Callable[..., Any], type_name: str) -> None:
        """Register a custom type for serialization.

        :param type_obj: The Python type or callable.
        :type type_obj: type | Callable[..., Any]
        :param type_name: The string name to use in the serialized form.
        :type type_name: str
        """
        self.type_registry[type_obj] = type_name
        self._reverse_type_registry[type_name] = type_obj

    def register_field_type(self, field_name: str, field_type: type | Callable[..., Any]) -> None:
        """Register an implicit type for a field name.

        :param field_name: The name of the field.
        :type field_name: str
        :param field_type: The Python type or callable for this field.
        :type field_type: type | Callable[..., Any]
        """
        self.implicit_field_types[field_name] = field_type

    def serialize(self, expr: Expr) -> str:
        """Serialize an expression to a string.

        :param expr: The expression to be serialized.
        :type expr: Expr
        :return: The string representation of the expression.
        :rtype: str
        """
        visitor = _SerializerVisitor(self.type_registry, self.include_all_types, self.url_encode)
        result = expr.accept(visitor)
        return quote(result, safe="") if self.url_encode else result

    def deserialize(self, s: str) -> Expr:
        """Deserialize a string into an expression object.

        :param s: The string to be deserialized.
        :type s: str
        :return: The parsed expression object.
        :rtype: Expr
        :raises tatsu.exceptions.FailedParse: If the string is not a valid expression.
        """
        if self.url_encode:
            s = unquote(s)
        semantics = _DeserializerSemantics(
            self._reverse_type_registry, self.implicit_field_types, self.sanitizer
        )
        return _PARSER.parse(s, semantics=semantics)  # type: ignore[no-any-return]


__all__ = ["Serializer"]
