"""Closed, typed rule expressions for governed DataOps rules.

The module deliberately owns both parsing and validation.  It produces JSON
only; expressions are never delegated to Python, SQL, JavaScript, or a model
runtime.  Compilers consume the resulting AST in a later release.
"""

from __future__ import annotations

from dataclasses import dataclass
import math
import re
from typing import Any


MAX_SOURCE_LENGTH = 20_000
MAX_AST_DEPTH = 40
MAX_AST_NODES = 500
MAX_REGEX_LENGTH = 500
ALLOWED_FUNCTIONS = frozenset(
    {
        "matches",
        "lower",
        "upper",
        "trim",
        "length",
        "coalesce",
        "date",
        "timestamp",
        "abs",
        "round",
    }
)
SUPPORTED_BACKENDS = frozenset({"postgresql", "mysql", "polars"})

_IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_]{0,127}$")
_FIELD_TYPE_ALIASES = {
    "bigint": "integer",
    "bool": "boolean",
    "character varying": "string",
    "datetime": "timestamp",
    "int": "integer",
    "numeric": "decimal",
    "text": "string",
    "varchar": "string",
}
_FIELD_TYPES = {
    "binary",
    "boolean",
    "date",
    "decimal",
    "double",
    "float",
    "integer",
    "json",
    "string",
    "timestamp",
    "timestamptz",
}
_NUMERIC_TYPES = {"integer", "decimal", "float", "double"}
_LITERAL_TYPES = {"boolean", "decimal", "integer", "null", "string"}
_BINARY_OPERATORS = {
    "||",
    "&&",
    "==",
    "!=",
    "<",
    "<=",
    ">",
    ">=",
    "+",
    "-",
    "*",
    "/",
    "%",
}
_PRECEDENCE = {
    "||": 10,
    "&&": 20,
    "==": 30,
    "!=": 30,
    "<": 30,
    "<=": 30,
    ">": 30,
    ">=": 30,
    "+": 40,
    "-": 40,
    "*": 50,
    "/": 50,
    "%": 50,
}


@dataclass(frozen=True)
class _Token:
    kind: str
    value: Any
    position: int


def _error(position: int, message: str) -> ValueError:
    return ValueError(f"expression {message} at position {position}")


class _Lexer:
    def __init__(self, source: str):
        self.source = source
        self.length = len(source)
        self.position = 0

    def tokens(self) -> list[_Token]:
        result: list[_Token] = []
        while self.position < self.length:
            char = self.source[self.position]
            if char.isspace():
                self.position += 1
            elif char in "'\"":
                result.append(self._string())
            elif char.isdigit():
                result.append(self._number())
            elif char.isascii() and (char.isalpha() or char == "_"):
                result.append(self._identifier())
            elif char in "(),":
                result.append(_Token(char, char, self.position))
                self.position += 1
            else:
                result.append(self._operator())
        result.append(_Token("EOF", None, self.position))
        return result

    def _string(self) -> _Token:
        start = self.position
        quote = self.source[self.position]
        self.position += 1
        pieces: list[str] = []
        escapes = {"n": "\n", "r": "\r", "t": "\t", "b": "\b", "f": "\f"}
        while self.position < self.length:
            char = self.source[self.position]
            self.position += 1
            if char == quote:
                value = "".join(pieces)
                if len(value) > 4_000:
                    raise _error(start, "string literal exceeds 4000 characters")
                return _Token("LITERAL", ("string", value), start)
            if char != "\\":
                pieces.append(char)
                continue
            if self.position >= self.length:
                raise _error(start, "contains an unfinished string escape")
            escaped = self.source[self.position]
            self.position += 1
            if escaped in escapes:
                pieces.append(escapes[escaped])
            elif escaped in {"\\", "'", '\"'}:
                pieces.append(escaped)
            elif escaped == "u":
                encoded = self.source[self.position : self.position + 4]
                if len(encoded) != 4 or not re.fullmatch(r"[0-9a-fA-F]{4}", encoded):
                    raise _error(start, "contains an invalid unicode escape")
                pieces.append(chr(int(encoded, 16)))
                self.position += 4
            else:
                raise _error(start, "contains an unsupported string escape")
        raise _error(start, "contains an unterminated string")

    def _number(self) -> _Token:
        start = self.position
        while self.position < self.length and self.source[self.position].isdigit():
            self.position += 1
        kind = "integer"
        if self.position < self.length and self.source[self.position] == ".":
            kind = "decimal"
            self.position += 1
            decimal_start = self.position
            while self.position < self.length and self.source[self.position].isdigit():
                self.position += 1
            if self.position == decimal_start:
                raise _error(start, "contains an invalid numeric literal")
        raw = self.source[start : self.position]
        value: int | float = int(raw) if kind == "integer" else float(raw)
        if isinstance(value, float) and not math.isfinite(value):
            raise _error(start, "contains a non-finite numeric literal")
        return _Token("LITERAL", (kind, value), start)

    def _identifier(self) -> _Token:
        start = self.position
        self.position += 1
        while self.position < self.length:
            char = self.source[self.position]
            if not (char.isascii() and (char.isalnum() or char == "_")):
                break
            self.position += 1
        value = self.source[start : self.position]
        if len(value) > 128:
            raise _error(start, "identifier exceeds 128 characters")
        lowered = value.lower()
        if lowered == "true":
            return _Token("LITERAL", ("boolean", True), start)
        if lowered == "false":
            return _Token("LITERAL", ("boolean", False), start)
        if lowered == "null":
            return _Token("LITERAL", ("null", None), start)
        return _Token("IDENT", value, start)

    def _operator(self) -> _Token:
        start = self.position
        for operator in ("&&", "||", "==", "!=", "<=", ">="):
            if self.source.startswith(operator, start):
                self.position += len(operator)
                return _Token("OP", operator, start)
        char = self.source[self.position]
        if char in "!<>+-*/%":
            self.position += 1
            return _Token("OP", char, start)
        raise _error(start, f"contains unexpected character {char!r}")


class _Parser:
    def __init__(self, source: str):
        self.tokens = _Lexer(source).tokens()
        self.index = 0

    @property
    def current(self) -> _Token:
        return self.tokens[self.index]

    def consume(self, kind: str, value: str | None = None) -> _Token:
        token = self.current
        if token.kind != kind or (value is not None and token.value != value):
            expected = value if value is not None else kind
            raise _error(token.position, f"expected {expected!r}")
        self.index += 1
        return token

    def parse(self) -> dict[str, Any]:
        expression = self._expression(0)
        if self.current.kind != "EOF":
            raise _error(self.current.position, "contains unexpected trailing input")
        return expression

    def _expression(self, minimum_precedence: int) -> dict[str, Any]:
        token = self.current
        self.index += 1
        if token.kind == "LITERAL":
            literal_type, value = token.value
            left = {"kind": "literal", "type": literal_type, "value": value}
        elif token.kind == "IDENT":
            if self.current.kind == "(":
                self.index += 1
                arguments: list[dict[str, Any]] = []
                if self.current.kind != ")":
                    while True:
                        arguments.append(self._expression(0))
                        if self.current.kind != ",":
                            break
                        self.index += 1
                self.consume(")")
                left = {
                    "kind": "call",
                    "function": token.value,
                    "arguments": arguments,
                }
            else:
                left = {"kind": "identifier", "name": token.value}
        elif token.kind == "OP" and token.value in {"!", "-"}:
            left = {
                "kind": "unary",
                "operator": token.value,
                "operand": self._expression(60),
            }
        elif token.kind == "(":
            left = self._expression(0)
            self.consume(")")
        else:
            raise _error(token.position, "contains an unexpected token")

        while self.current.kind == "OP":
            operator = self.current.value
            precedence = _PRECEDENCE.get(operator)
            if precedence is None or precedence < minimum_precedence:
                break
            self.index += 1
            left = {
                "kind": "binary",
                "operator": operator,
                "left": left,
                "right": self._expression(precedence + 1),
            }
        return left


def parse_expression(source: str) -> dict:
    """Parse a V1 expression into the closed, JSON-serializable AST."""

    if not isinstance(source, str) or not source.strip():
        raise ValueError("expression source is required")
    if len(source) > MAX_SOURCE_LENGTH:
        raise ValueError(f"expression source exceeds {MAX_SOURCE_LENGTH} characters")
    ast = _Parser(source).parse()
    _validate_ast(ast)
    return ast


def _validate_ast(ast: Any) -> None:
    count = 0

    def walk(node: Any, depth: int) -> None:
        nonlocal count
        if depth > MAX_AST_DEPTH:
            raise ValueError(f"expression AST depth exceeds {MAX_AST_DEPTH}")
        count += 1
        if count > MAX_AST_NODES:
            raise ValueError(f"expression AST node count exceeds {MAX_AST_NODES}")
        if not isinstance(node, dict):
            raise ValueError("expression AST node must be an object")
        kind = node.get("kind")
        if kind == "literal":
            if set(node) != {"kind", "type", "value"}:
                raise ValueError("literal AST node must have a closed shape")
            literal_type = node.get("type")
            value = node.get("value")
            if literal_type not in _LITERAL_TYPES:
                raise ValueError("literal AST node has an unsupported type")
            if literal_type == "null" and value is not None:
                raise ValueError("null literal must have a null value")
            if literal_type == "boolean" and not isinstance(value, bool):
                raise ValueError("boolean literal must have a boolean value")
            if literal_type == "string" and (
                not isinstance(value, str) or len(value) > 4_000
            ):
                raise ValueError("string literal is invalid")
            if literal_type == "integer" and (
                isinstance(value, bool) or not isinstance(value, int)
            ):
                raise ValueError("integer literal is invalid")
            if literal_type == "decimal" and (
                isinstance(value, bool)
                or not isinstance(value, (int, float))
                or not math.isfinite(float(value))
            ):
                raise ValueError("decimal literal is invalid")
            return
        if kind == "identifier":
            if set(node) != {"kind", "name"} or not isinstance(node.get("name"), str):
                raise ValueError("identifier AST node must have a closed shape")
            if _IDENTIFIER.fullmatch(node["name"]) is None:
                raise ValueError("identifier AST node has an invalid name")
            return
        if kind == "unary":
            if set(node) != {"kind", "operator", "operand"}:
                raise ValueError("unary AST node must have a closed shape")
            if node.get("operator") not in {"!", "-"}:
                raise ValueError("unary AST node has an unsupported operator")
            walk(node.get("operand"), depth + 1)
            return
        if kind == "binary":
            if set(node) != {"kind", "operator", "left", "right"}:
                raise ValueError("binary AST node must have a closed shape")
            if node.get("operator") not in _BINARY_OPERATORS:
                raise ValueError("binary AST node has an unsupported operator")
            walk(node.get("left"), depth + 1)
            walk(node.get("right"), depth + 1)
            return
        if kind == "call":
            if set(node) != {"kind", "function", "arguments"}:
                raise ValueError("call AST node must have a closed shape")
            function = node.get("function")
            arguments = node.get("arguments")
            if not isinstance(function, str) or _IDENTIFIER.fullmatch(function) is None:
                raise ValueError("call AST node has an invalid function")
            if not isinstance(arguments, list) or len(arguments) > 100:
                raise ValueError("call AST arguments must be a bounded array")
            for argument in arguments:
                walk(argument, depth + 1)
            # This resource bound is independent of schema typing, so RuleSpec
            # normalization can reject an unsafe regex before compilation.
            if (
                function == "matches"
                and len(arguments) >= 2
                and arguments[1].get("kind") == "literal"
                and arguments[1].get("type") == "string"
            ):
                if len(arguments[1]["value"]) > MAX_REGEX_LENGTH:
                    raise ValueError(
                        f"regex pattern exceeds {MAX_REGEX_LENGTH} characters"
                    )
            return
        raise ValueError("expression AST node kind is unsupported")

    walk(ast, 1)


def _normalized_fields(fields: Any) -> dict[str, str]:
    if not isinstance(fields, dict) or len(fields) > 1_000:
        raise ValueError("expression fields must be a bounded object")
    normalized: dict[str, str] = {}
    for name, field_type in fields.items():
        if not isinstance(name, str) or _IDENTIFIER.fullmatch(name) is None:
            raise ValueError("expression field names must be identifiers")
        if not isinstance(field_type, str):
            raise ValueError(f"expression field {name} type is invalid")
        result_type = _FIELD_TYPE_ALIASES.get(
            field_type.strip().lower(), field_type.strip().lower()
        )
        if result_type not in _FIELD_TYPES:
            raise ValueError(f"expression field {name} has an unsupported type")
        normalized[name] = result_type
    return normalized


def _is_numeric(value_type: str) -> bool:
    return value_type in _NUMERIC_TYPES


def _promote_numeric(left: str, right: str) -> str:
    order = {"integer": 0, "decimal": 1, "float": 2, "double": 3}
    return max((left, right), key=lambda value: order[value])


def _compatible(left: str, right: str) -> bool:
    return left == "null" or right == "null" or left == right or (
        _is_numeric(left) and _is_numeric(right)
    )


def _require_arguments(function: str, arguments: list[dict], expected: int | tuple[int, int]) -> None:
    if isinstance(expected, int):
        valid = len(arguments) == expected
    else:
        valid = expected[0] <= len(arguments) <= expected[1]
    if not valid:
        raise ValueError(f"function {function} received an unsupported argument count")


def type_check_expression(ast: dict, fields: dict[str, str]) -> str:
    """Validate schema-bound identifiers and return the expression result type."""

    _validate_ast(ast)
    normalized_fields = _normalized_fields(fields)

    def check(node: dict) -> str:
        kind = node["kind"]
        if kind == "literal":
            return node["type"]
        if kind == "identifier":
            name = node["name"]
            if name not in normalized_fields:
                raise ValueError(f"unknown expression field: {name}")
            return normalized_fields[name]
        if kind == "unary":
            operand_type = check(node["operand"])
            operator = node["operator"]
            if operator == "!":
                if operand_type != "boolean":
                    raise ValueError("logical not requires a boolean operand")
                return "boolean"
            if not _is_numeric(operand_type):
                raise ValueError("numeric negation requires a numeric operand")
            return operand_type
        if kind == "binary":
            left_type = check(node["left"])
            right_type = check(node["right"])
            operator = node["operator"]
            if operator in {"&&", "||"}:
                if left_type != "boolean" or right_type != "boolean":
                    raise ValueError("logical operators require boolean operands")
                return "boolean"
            if operator in {"+", "-", "*", "/", "%"}:
                if not _is_numeric(left_type) or not _is_numeric(right_type):
                    raise ValueError("arithmetic operators require numeric operands")
                return _promote_numeric(left_type, right_type)
            if not _compatible(left_type, right_type):
                raise ValueError(
                    "comparison operands must have compatible types; use an explicit cast"
                )
            return "boolean"
        function = node["function"]
        arguments = node["arguments"]
        if function not in ALLOWED_FUNCTIONS:
            raise ValueError(f"unsupported expression function: {function}")
        argument_types = [check(argument) for argument in arguments]
        if function == "matches":
            _require_arguments(function, arguments, 2)
            if argument_types != ["string", "string"]:
                raise ValueError("matches requires string arguments")
            pattern = arguments[1]
            if pattern["kind"] != "literal" or pattern["type"] != "string":
                raise ValueError("matches pattern must be a string literal")
            if len(pattern["value"]) > MAX_REGEX_LENGTH:
                raise ValueError(f"regex pattern exceeds {MAX_REGEX_LENGTH} characters")
            try:
                re.compile(pattern["value"])
            except re.error as exc:
                raise ValueError("regex pattern is invalid") from exc
            return "boolean"
        if function in {"lower", "upper", "trim"}:
            _require_arguments(function, arguments, 1)
            if argument_types != ["string"]:
                raise ValueError(f"{function} requires a string argument")
            return "string"
        if function == "length":
            _require_arguments(function, arguments, 1)
            if argument_types != ["string"]:
                raise ValueError("length requires a string argument")
            return "integer"
        if function == "coalesce":
            _require_arguments(function, arguments, (2, 100))
            concrete = [value for value in argument_types if value != "null"]
            if not concrete:
                return "null"
            if not all(_compatible(concrete[0], value) for value in concrete[1:]):
                raise ValueError("coalesce arguments must have compatible types")
            if all(_is_numeric(value) for value in concrete):
                return max(concrete, key=lambda value: {"integer": 0, "decimal": 1, "float": 2, "double": 3}[value])
            return concrete[0]
        if function == "date":
            _require_arguments(function, arguments, 1)
            if argument_types[0] not in {"string", "date", "timestamp", "timestamptz"}:
                raise ValueError("date requires a string or temporal argument")
            return "date"
        if function == "timestamp":
            _require_arguments(function, arguments, 1)
            if argument_types[0] not in {"string", "date", "timestamp", "timestamptz"}:
                raise ValueError("timestamp requires a string or temporal argument")
            return "timestamp"
        if function == "abs":
            _require_arguments(function, arguments, 1)
            if not _is_numeric(argument_types[0]):
                raise ValueError("abs requires a numeric argument")
            return argument_types[0]
        _require_arguments(function, arguments, (1, 2))
        if not _is_numeric(argument_types[0]):
            raise ValueError("round requires a numeric value")
        if len(argument_types) == 2 and argument_types[1] != "integer":
            raise ValueError("round precision requires an integer")
        return argument_types[0]

    return check(ast)


def backend_support(ast: dict) -> frozenset[str]:
    """Return backends that can preserve this grammar's V1 semantics.

    The result is intentionally a capability analysis, not a compiler.  It
    validates the closed AST again so hand-authored JSON cannot bypass the
    function allowlist before a compiler sees it.
    """

    _validate_ast(ast)

    def validate_functions(node: dict) -> None:
        if node["kind"] == "call":
            if node["function"] not in ALLOWED_FUNCTIONS:
                raise ValueError(
                    f"unsupported expression function: {node['function']}"
                )
            for argument in node["arguments"]:
                validate_functions(argument)
        elif node["kind"] == "unary":
            validate_functions(node["operand"])
        elif node["kind"] == "binary":
            validate_functions(node["left"])
            validate_functions(node["right"])

    validate_functions(ast)
    return SUPPORTED_BACKENDS