"""Expression compiler: DSL expressions → Hash-based AST nodes.

Compiles Python DSL expressions (arithmetic, comparisons, aggregates,
constraints) into Hash-based AST nodes stored as model defines. The solver
service deserializes these AST nodes to reconstruct the problem in MOI form.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Optional

from relationalai.semantics.frontend import base as b
from relationalai.semantics.frontend.core import Hash, cast as core_cast
from relationalai.semantics import std

from .wire_format import (
    _hash_with_res,
    _FIRST_ORDER_CODES,
    _COMPARISON_CODES,
    _HIGHER_ORDER_CODES,
    _UNSUPPORTED_OPS,
)

if TYPE_CHECKING:
    from .problem import Problem


class SymbolicNode:
    """Wrapper marking an expression as containing decision variables."""

    def __init__(self, expr: Any) -> None:
        if isinstance(expr, SymbolicNode):
            raise TypeError("Cannot double-wrap SymbolicNode.")
        self.expr = expr


def _expr_dict_key(expr: Any) -> int:
    """Get unique integer key for an expression."""
    key = expr.to_dict_key()
    if not isinstance(key, int):
        raise TypeError(f"Expression key must be int, got {type(key).__name__}.")
    return key


class ExpressionCompiler:
    """Compiles DSL expressions into Hash-based AST nodes for the solver service."""

    # Expression types not supported in problem constraints (checked by rewrite()).
    _UNSUPPORTED_EXPRS: list[tuple[type | tuple[type, ...], str]] = [
        (b.New, "Cannot use 'New' expressions in problem constraints."),
        (
            b.FieldRef,
            "Cannot use field indexing in problem constraints. "
            "Use an explicit ref binding: qty = Float.ref(); sum(qty * X.var).where(X.rel(Y, qty)).per(Y)",
        ),
        (b.MetaRef, "Cannot use meta-references in problem constraints."),
        (b.AsBool, "Cannot use 'as_bool()' in problem constraints. Use comparison operators directly."),
        ((b.Not, b.Distinct, b.Group), "Not/Distinct/Group not supported in problem constraints."),
        ((b.DerivedTable, b.DerivedColumn), "Cannot use derived tables/columns in problem constraints."),
    ]

    _REWRITE_DISPATCH: dict[type, str] = {
        b.Concept: "_rewrite_variable_ref",
        b.Ref: "_rewrite_variable_ref",
        b.Chain: "_rewrite_chain",
        b.Relationship: "_rewrite_relationship",
        b.Property: "_rewrite_relationship",  # Property is subclass of Relationship
        b.Reading: "_rewrite_reading",  # Resolves to parent variable if value field unchanged
        b.Match: "_rewrite_match",
        b.Union: "_rewrite_match",
        b.FilterBy: "_rewrite_filter_by",
        b.Expression: "_compile_arithmetic",
        b.Aggregate: "_rewrite_aggregate",
        b.Fragment: "_rewrite_fragment",
    }

    _MAX_REWRITE_DEPTH = 400

    def __init__(self, problem: Problem) -> None:
        self.model = problem._model
        self.problem = problem
        self.bound_variables: dict[int, Any] = {}
        self.defines: list[Any] = []
        self.wheres: list[Any] = []
        self.children: list[ExpressionCompiler] = []
        self._rewrite_depth = 0

    def _create_scope(self) -> ExpressionCompiler:
        """Create child scope for nested expressions (e.g., aggregates)."""
        child = ExpressionCompiler(self.problem)
        child.bound_variables = self.bound_variables.copy()
        child.defines = self.defines.copy()
        child.wheres = self.wheres.copy()
        child._rewrite_depth = self._rewrite_depth
        self.children.append(child)
        return child

    # =========================================================================
    # Public Rewriting Methods
    # =========================================================================

    def rewrite_where(self, *exprs: Any) -> list[Any]:
        """Rewrite where clauses in two passes: variable relationships first, then rest."""
        results: list[Any] = []

        # Pass 1: Variable relationships populate bound_variables.
        for expr in exprs:
            if isinstance(expr, b.Expression) and isinstance(expr._op, b.Relationship):
                self._check_reading_value_field(expr._op)
                if self._variable_lookup_key(expr._op) in self.problem._variable_relationships:
                    results.append(self._bind_variable_entity(expr))
                    continue
            results.append(None)

        # Pass 2: Rewrite remaining expressions (may reference pass 1 variables).
        for i, expr in enumerate(exprs):
            if results[i] is None:
                # Scalars/literals don't have to_dict_key(); skip bound_variables check.
                if hasattr(expr, "to_dict_key"):
                    key = _expr_dict_key(expr)
                    if key in self.bound_variables:
                        results[i] = expr
                        continue
                results[i] = self._passthrough(expr)

        return results

    def rewrite(self, expr: Any) -> Optional[SymbolicNode | Any]:
        """Rewrite expression using dispatch table.

        Returns
        -------
        SymbolicNode or Any or None
            ``SymbolicNode`` if expression contains decision variables,
            a rewritten expression if modified but not symbolic,
            or ``None`` if no rewriting was needed.
        """
        # Base cases: literals, None, and data sources (no rewriting needed).
        if expr is None or isinstance(expr, (int, float, str, b.Literal, b.Table, b.Data)):
            return None

        self._rewrite_depth += 1
        if self._rewrite_depth > self._MAX_REWRITE_DEPTH:
            raise RecursionError(
                f"Expression tree exceeds maximum depth ({self._MAX_REWRITE_DEPTH}). "
                f"Simplify the expression or break it into smaller sub-expressions."
            )
        try:
            return self._rewrite_inner(expr)
        finally:
            self._rewrite_depth -= 1

    def _rewrite_inner(self, expr: Any) -> Optional[SymbolicNode | Any]:
        # Alias: delegate to source item.
        if isinstance(expr, b.Alias):
            return self.rewrite(expr._source_item)

        # Dispatch to type-specific handler.
        handler_name = self._REWRITE_DISPATCH.get(type(expr))
        if handler_name is not None:
            return getattr(self, handler_name)(expr)

        # Unsupported expression types.
        for types, msg in self._UNSUPPORTED_EXPRS:
            if isinstance(expr, types):  # type: ignore[arg-type]
                raise ValueError(msg)
        raise NotImplementedError(f"Expression type '{type(expr).__name__}' not supported in problem constraints.")

    # -------------------------------------------------------------------------
    # Type-specific rewrite handlers
    # -------------------------------------------------------------------------

    def _rewrite_variable_ref(self, expr: b.Concept | b.Ref) -> Optional[SymbolicNode]:
        """Rewrite Concept or Ref that may be a decision variable."""
        key = _expr_dict_key(expr)
        return SymbolicNode(self.bound_variables[key]) if key in self.bound_variables else None

    def _rewrite_chain(self, expr: b.Chain) -> Optional[SymbolicNode]:
        """Rewrite Chain (entity.relationship). Returns SymbolicNode if relationship is a decision variable."""
        start, next_expr = expr._start, expr._next

        # Handle nested chains (e.g., Factory.product.amount):
        # resolve intermediate chain's target concept and use it as parent for variable lookup.
        if isinstance(start, b.Chain):
            start_result = self.rewrite(start)
            if isinstance(start_result, SymbolicNode):
                raise ValueError("Chain start must not be symbolic.")
            if isinstance(next_expr, b.Relationship):
                self._check_reading_value_field(next_expr)
                result = self._get_variable_ref(next_expr, start)
                if result is not None:
                    key, var_ref = result
                    self.bound_variables[key] = var_ref
                    return SymbolicNode(var_ref)
            return start_result

        if not isinstance(start, (b.Concept, b.Ref)):
            raise TypeError(f"Chain start must be Concept or Ref, got {type(start).__name__}.")
        if self.rewrite(start) is not None:
            raise ValueError("Chain start must not be symbolic.")

        # Check if the relationship part is a decision variable.
        if isinstance(next_expr, b.Relationship):
            self._check_reading_value_field(next_expr)
            result = self._get_variable_ref(next_expr, start)
            if result is not None:
                key, var_ref = result
                self.bound_variables[key] = var_ref
                return SymbolicNode(var_ref)
        return None

    @staticmethod
    def _variable_lookup_key(rel: b.Relationship) -> int:
        """Get the variable lookup key, resolving Readings to their parent.

        Readings reference the same underlying relationship as their parent,
        so variable lookups should use the parent's key.
        """
        parent = rel._relationship if isinstance(rel, b.Reading) else rel
        return _expr_dict_key(parent)

    def _check_reading_value_field(self, expr: b.Relationship) -> None:
        """Raise if a Reading moves the value field out of last position.

        Entity-field reordering is fine — the solver binds by field name,
        not position. But the last field is the decision variable value;
        moving it breaks variable binding.
        """
        if not isinstance(expr, b.Reading):
            return
        parent = expr._relationship
        key = self._variable_lookup_key(expr)
        if key in self.problem._variable_relationships or key in self.bound_variables:
            if expr._fields[-1] is not parent._fields[-1]:
                parent_name = parent._short_name if parent._short_name else parent._name
                raise ValueError(
                    f"Cannot use alternative reading (`.alt()`) that reorders the value field "
                    f"of relationship '{parent_name}'. The last field of a decision-variable "
                    f"relationship must remain the value field. Use the original relationship "
                    f"or a reading that only reorders entity fields."
                )

    def _rewrite_reading(self, expr: b.Reading) -> Optional[SymbolicNode]:
        """Rewrite Reading, resolving to parent variable if value field is unchanged."""
        self._check_reading_value_field(expr)
        key = self._variable_lookup_key(expr)
        if key in self.bound_variables:
            return SymbolicNode(self.bound_variables[key])
        result = self._get_variable_ref(expr, None)
        if result is not None:
            ret_key, var_ref = result
            self.bound_variables[ret_key] = var_ref
            return SymbolicNode(var_ref)
        return None

    def _rewrite_relationship(self, expr: b.Relationship) -> Optional[SymbolicNode]:
        """Rewrite Relationship that may be a decision variable."""
        key = _expr_dict_key(expr)
        if key in self.bound_variables:
            return SymbolicNode(self.bound_variables[key])
        result = self._get_variable_ref(expr, None)
        if result is not None:
            ret_key, var_ref = result
            self.bound_variables[ret_key] = var_ref
            return SymbolicNode(var_ref)
        return None

    def _rewrite_match(self, expr: b.Match) -> Optional[SymbolicNode | Any]:
        """Rewrite Match/Union expression."""
        items = expr._items
        rewritten: list[Any] = []
        has_symbolic = False

        for item in items:
            result = self.rewrite(item)
            if isinstance(result, SymbolicNode):
                has_symbolic = True
                rewritten.append(result.expr)
            else:
                rewritten.append(result if result is not None else item)

        if not has_symbolic and rewritten == list(items):
            return None

        # Build the appropriate Match or Union.
        if isinstance(expr, b.Union):
            combined = self.model.union(*rewritten)
        else:
            combined = b.Match(self.model, *rewritten)

        return SymbolicNode(combined) if has_symbolic else combined

    def _rewrite_filter_by(self, expr: b.FilterBy) -> Optional[Any]:
        """Rewrite FilterBy expression (symbolic values not allowed)."""
        new_kwargs, modified = {}, False
        for k, v in expr._kwargs.items():
            r = self.rewrite(v)
            if isinstance(r, SymbolicNode):
                raise ValueError(f"Cannot use symbolic value in FilterBy argument '{k}'.")
            new_kwargs[k] = r if r is not None else v
            modified = modified or r is not None

        return expr._params[0].filter_by(**new_kwargs) if modified else None

    def _compile_arithmetic(self, expr: b.Expression) -> Optional[SymbolicNode | Any]:
        """Compile arithmetic/comparison expression. Builds AST node if any arg is symbolic."""
        op = self.rewrite(expr._op)
        if isinstance(op, SymbolicNode):
            raise ValueError("Expression operator cannot be symbolic.")
        op = op if op is not None else expr._op

        args = expr._args
        rewritten: list[Any] = []
        has_symbolic = False
        for arg in args:
            result = self.rewrite(arg)
            if isinstance(result, SymbolicNode):
                has_symbolic = True
            rewritten.append(result if result is not None else arg)

        if not has_symbolic and rewritten == list(args):
            return None
        if not has_symbolic:
            return b.Expression(op, rewritten)

        # Build symbolic AST node.
        op_name = op._short_name
        if not isinstance(op_name, str):
            raise TypeError(f"Operator name must be str, got {type(op_name).__name__}.")

        if op_name in _FIRST_ORDER_CODES:
            op_code = _FIRST_ORDER_CODES[op_name]
            rewritten = rewritten[:-1]  # Drop auto-filled output arg.
        elif op_name in _COMPARISON_CODES:
            op_code = _COMPARISON_CODES[op_name]
        else:
            # Provide helpful error for unsupported operators
            if op_name in _UNSUPPORTED_OPS:
                msg = (
                    f"Operator '{op_name}' is not supported by any solver backend "
                    f"(not available in MOI/MiniZinc/Ipopt/HiGHS)."
                )
            else:
                supported = sorted(list(_FIRST_ORDER_CODES.keys()) + list(_COMPARISON_CODES.keys()))
                msg = (
                    f"Operator '{op_name}' is not supported by the solver engine. "
                    f"Supported operators: {', '.join(supported)}"
                )
            raise NotImplementedError(msg)

        # Create Hash node unique per (expr_id, args).
        hash_args = [a.expr if isinstance(a, SymbolicNode) else a for a in rewritten]
        node = Hash.ref()
        self.wheres.append(_hash_with_res([expr._id, *hash_args], node))
        self.defines.append(self.problem._operator(node, op_code))

        for i, arg in enumerate(rewritten):
            if isinstance(arg, SymbolicNode):
                self.defines.append(self.problem._ordered_args_hash(node, i, arg.expr))
            else:
                self.defines.append(self.problem._ordered_args_data(node, i, arg))

        return SymbolicNode(node)

    def _extract_aggregate_args(self, expr: b.Aggregate, op_name: str) -> tuple[Any, list[Any], list[Any]]:
        """Extract symbolic arg and other args based on operator type."""
        args, proj_args = list(expr._args), list(expr._projection_args)
        if op_name == "count":
            return proj_args.pop(), args, proj_args
        elif op_name == "rank":
            if not isinstance(args[0], b.TupleVariable):
                raise TypeError("Rank aggregate expects TupleVariable as first arg.")
            return None, args[0]._items, []
        elif op_name in ("special_ordered_set_type_1", "special_ordered_set_type_2"):
            return args[1], args[:1], proj_args
        else:
            return args[0], args[1:], proj_args

    def _rewrite_aggregate(self, expr: b.Aggregate) -> Optional[SymbolicNode | Any]:
        """Rewrite aggregate (sum/min/max/count/etc)."""
        if expr._distinct:
            raise NotImplementedError("Distinct aggregates not yet supported.")

        op = self.rewrite(expr._op)
        if isinstance(op, SymbolicNode):
            raise ValueError("Aggregate operator cannot be symbolic.")
        op = op if op is not None else expr._op
        op_name = op._short_name

        # Extract arguments based on operator type
        sym_arg, args, proj_args = self._extract_aggregate_args(expr, op_name)

        # Rewrite components in subcontext
        child = self._create_scope()
        rw_where = child.rewrite_where(*expr._where._where)
        rw_proj = [child._passthrough(a) for a in proj_args]
        rw_args = [child._passthrough(a) for a in args]
        rw_group = [child._passthrough(a) for a in expr._group._args]
        rw_sym = None if op_name == "rank" else child.rewrite(sym_arg)

        # Early return if nothing changed
        if (
            rw_sym is None
            and rw_proj == list(proj_args)
            and rw_args == list(args)
            and rw_where == list(expr._where._where)
        ):
            return None

        # Non-symbolic case: rebuild aggregate
        if not isinstance(rw_sym, SymbolicNode):
            sym_arg_final = rw_sym if rw_sym is not None else sym_arg
            if op_name == "rank":
                new_args = [b.TupleVariable(rw_args), expr._args[1], expr._args[2]]
            elif op_name in (
                "special_ordered_set_type_1",
                "special_ordered_set_type_2",
            ):
                new_args = [*rw_proj, *rw_args, sym_arg_final]
            else:
                new_args = [*rw_proj, sym_arg_final, *rw_args]
            return b.Aggregate(op, *new_args, distinct=expr._distinct).per(*rw_group).where(*rw_where, *child.wheres)

        # SymbolicNode case: build Hash-based AST node
        if op_name not in _HIGHER_ORDER_CODES:
            supported_aggs = list(_HIGHER_ORDER_CODES.keys())
            msg = f"Aggregate '{op_name}' is not supported by the solver engine."
            if op_name in ("avg", "string_join", "rank", "limit", "top", "bottom"):
                msg += f" Only these aggregates are supported: {', '.join(sorted(supported_aggs))}"
            else:
                msg += f" Supported aggregates: {', '.join(sorted(supported_aggs))}"
            raise NotImplementedError(msg)

        sym_expr = rw_sym.expr
        node = Hash.ref()
        child.wheres.append(_hash_with_res([expr._id, *rw_group], node))
        child.wheres.extend(rw_where)
        child.defines.append(self.problem._operator(node, _HIGHER_ORDER_CODES[op_name]))

        if op_name == "special_ordered_set_type_2":
            if len(rw_proj) != 0 or len(rw_args) != 1:
                raise ValueError("SOS2 expects exactly 2 arguments (rank, variable).")
            child.defines.append(self.problem._ordered_args_hash(node, rw_args[0], sym_expr))
        else:
            arg_hash = std.common.hash(*rw_proj, sym_expr)
            child.defines.append(self.problem._unordered_args_hash(node, arg_hash, sym_expr))

        return SymbolicNode(self.model.select(node).where(*child.wheres))

    def _rewrite_fragment(self, expr: b.Fragment) -> Optional[SymbolicNode | Any]:
        """Rewrite inline select/where fragment."""
        if expr._define or expr._require:
            raise ValueError("Fragments with define/require not supported.")
        if len(expr._select) != 1:
            raise ValueError(f"Fragment must have exactly 1 select, got {len(expr._select)}.")

        child = self._create_scope()
        rw_where = child.rewrite_where(*expr._where)
        rw_select = child.rewrite(expr._select[0])
        wheres = self.model.where(*rw_where, *child.wheres)

        if isinstance(rw_select, SymbolicNode):
            return SymbolicNode(wheres.select(rw_select.expr))
        if rw_select is not None:
            return wheres.select(rw_select)
        return None

    # =========================================================================
    # Private Helper Methods
    # =========================================================================

    def _bind_variable_entity(self, expr: b.Expression) -> Any:
        """Bind entity ref to a variable: store ref in bound_variables, return field bindings."""
        relationship = expr._op
        if not isinstance(relationship, b.Relationship):
            raise TypeError(f"Expected Relationship operator, got {type(relationship).__name__}.")

        params = expr._args
        if len(params) != len(relationship._fields):
            raise ValueError(f"Parameter count {len(params)} != field count {len(relationship._fields)}.")

        # The last parameter is the output variable (the symbolic value).
        last_param = params[-1]
        if isinstance(last_param, b.Alias):
            last_param = last_param._source_item
        if not isinstance(last_param, (b.Concept, b.Ref)):
            raise TypeError(f"Last parameter must be Concept or Ref, got {type(last_param).__name__}.")

        # Build field bindings (grounding conditions).
        # Field names come from relationship (possibly a Reading) — they match
        # the user's argument order regardless of entity-field reordering.
        fields: dict[str, Any] = {}
        for i, param in enumerate(params[:-1]):
            rewritten = self.rewrite(param)
            if isinstance(rewritten, SymbolicNode):
                raise ValueError(f"Variable relationship parameter {i} cannot be symbolic.")
            fields[relationship._fields[i].name] = rewritten if rewritten is not None else param

        lookup_key = self._variable_lookup_key(relationship)
        var_ref = self.problem._variable_relationships[lookup_key].ref()
        self.bound_variables[_expr_dict_key(last_param)] = core_cast(b.MetaRef(Hash), var_ref)

        return self.problem._model.where(*[getattr(var_ref, name) == val for name, val in fields.items()])

    def _passthrough(self, expr: Any) -> Any:
        """Rewrite expression, raising if result is unexpectedly symbolic."""
        result = self.rewrite(expr)
        if isinstance(result, SymbolicNode):
            raise ValueError(f"Unexpected SymbolicNode result for {expr}.")
        return result if result is not None else expr

    def _get_variable_ref(
        self, relationship: b.Relationship, parent: b.Concept | b.Ref | b.Chain | None
    ) -> Optional[tuple[int, Any]]:
        """Get (lookup_key, Hash-cast var ref) if relationship is a variable, else None.

        Uses ``_variable_lookup_key`` to resolve Readings to their parent for
        the dict lookup, but ``relationship`` for field names (so entity binding
        matches the user's argument order, even for Readings).
        """
        key = self._variable_lookup_key(relationship)
        VarConcept = self.problem._variable_relationships.get(key)
        if VarConcept is None:
            return None

        var_ref = VarConcept.ref()
        if parent is not None:
            self.wheres.append(getattr(var_ref, relationship._fields[0].name)(parent))
        return (key, core_cast(b.MetaRef(Hash), var_ref))
