"""Serialization protocol for prescriptive problem encoding.

Defines the wire format used to serialize decision problems into Hash-based
AST nodes for transmission to the rai-solver-service. The solver service
deserializes these op codes to reconstruct the problem in MathOptInterface
(MOI) form before dispatching to a backend solver (HiGHS, Gurobi, Ipopt,
MiniZinc, etc.).
"""

from __future__ import annotations

from typing import TypeAlias

from relationalai.semantics.frontend import base as b
from relationalai.semantics.frontend.core import Hash
from relationalai.semantics import Any as CoreAny

# Semantic type alias for wire-format op codes. Documents intent without
# breaking type compatibility (NewType is too strict for existing dict lookups).
OpCode: TypeAlias = int

# Library for custom aggregation relations used in prescriptive expression handling.
_library = b.Library("aggregates")

# Relation for computing hash values from argument lists.
_hash = _library.Relation("hash", [b.Field.input("args", CoreAny, is_list=True), b.Field("hash", Hash)])


def _hash_with_res(args: list[b.Value], res: b.Value) -> b.Expression:
    """Compute a hash of the given arguments and bind to res."""
    return _hash(b.TupleVariable(args), res)


# =============================================================================
# Op Codes
# =============================================================================

# Wire format codes for problem serialization. Each entry maps an operator
# name to (numeric_code, category). Categories group operators for dispatch.
_OP_CODES: dict[str, tuple[int, str]] = {
    # Variable types (codes 0-9)
    "cont": (0, "var"),
    "int": (1, "var"),
    "bin": (2, "var"),
    # Expression types (codes 10-19)
    "min_objective": (10, "expr"),
    "max_objective": (11, "expr"),
    "constraint": (12, "expr"),
    # Arithmetic — infix display (codes 20-24)
    "+": (20, "arith"),
    "-": (21, "arith"),
    "*": (22, "arith"),
    "/": (23, "arith"),
    "^": (24, "arith"),
    # Math functions (codes 40-47)
    "abs": (40, "math"),
    "exp": (41, "math"),
    "natural_log": (42, "math"),
    "log": (43, "math"),
    "sqrt": (44, "math"),
    "square": (45, "math"),
    "cbrt": (46, "math"),
    "power": (47, "math"),
    # Comparisons (codes 70-75)
    "=": (70, "cmp"),
    "!=": (71, "cmp"),
    "<": (72, "cmp"),
    "<=": (73, "cmp"),
    ">": (74, "cmp"),
    ">=": (75, "cmp"),
    # Logical (codes 80-89)
    "implies": (80, "logic"),
    # Aggregates (codes 90-93)
    "sum": (90, "agg"),
    "min": (91, "agg"),
    "max": (92, "agg"),
    "count": (93, "agg"),
    # Global constraints (codes 110-112)
    "all_different": (110, "global"),
    "special_ordered_set_type_1": (111, "global"),
    "special_ordered_set_type_2": (112, "global"),
}

# =============================================================================
# Derived Lookups
# =============================================================================

# Reverse mapping: numeric code → operator name (used by display formatting).
_CODE_TO_NAME: dict[OpCode, str] = {code: name for name, (code, _) in _OP_CODES.items()}
assert len(_CODE_TO_NAME) == len(_OP_CODES), (
    f"Duplicate op codes detected: {len(_OP_CODES)} names but only {len(_CODE_TO_NAME)} unique codes"
)

# Category-filtered subsets (used by expression compiler and display).
_VAR_CODES: dict[str, OpCode] = {n: c for n, (c, cat) in _OP_CODES.items() if cat == "var"}
_EXPR_CODES: dict[str, OpCode] = {n: c for n, (c, cat) in _OP_CODES.items() if cat == "expr"}
_EXPR_CONCEPT_NAMES: dict[str, str] = {
    "constraint": "Constraint",
    "min_objective": "MinObjective",
    "max_objective": "MaxObjective",
}
_EXPR_TYPE_LABELS: list[tuple[str, str]] = [
    ("min_objective", "Minimization objectives"),
    ("max_objective", "Maximization objectives"),
    ("constraint", "Constraints"),
]
_FIRST_ORDER_CODES: dict[str, OpCode] = {n: c for n, (c, cat) in _OP_CODES.items() if cat in ("arith", "math")}
_COMPARISON_CODES: dict[str, OpCode] = {n: c for n, (c, cat) in _OP_CODES.items() if cat in ("cmp", "logic")}
_HIGHER_ORDER_CODES: dict[str, OpCode] = {n: c for n, (c, cat) in _OP_CODES.items() if cat in ("agg", "global")}
_INFIX: frozenset[str] = frozenset(n for n, (_, cat) in _OP_CODES.items() if cat == "arith")
_INFIX_AND_CMP: frozenset[str] = _INFIX | frozenset(n for n, (_, cat) in _OP_CODES.items() if cat == "cmp")

# Operators not supported by any solver backend (error messages reference these).
_UNSUPPORTED_OPS: frozenset[str] = frozenset({"%", "//", "floor", "ceil", "trunc", "round"})
