from typing import Sequence, TypeVar

from google.protobuf.internal.containers import RepeatedCompositeFieldContainer
from ortools.sat.cp_model_pb2 import ConstraintProto
from ortools.sat.python.cp_model import CpModel, CpSolver, IntVar, LiteralT

from .callbacks import (
    ObjectiveBoundRecorder,
    ObjectiveValueRecorder,
)


class CustomCpModel(CpModel):
    """A custom CpModel class that extends the OR-Tools CpModel class."""

    solver: CpSolver
    """CpSolver object for solving the model."""
    solution_callback: ObjectiveValueRecorder
    """Recorder for objective values during search."""
    obj_bound_recorder: ObjectiveBoundRecorder
    """Recorder for objective bounds during search."""

    def __init__(self) -> None:
        super().__init__()
        self.num_base_constraints: int = 0
        """Number of base constraints in the model."""

    # variable functions

    def change_domain(self, var: IntVar, domain: list[int]) -> None:
        """Changes the domain of a variable.

        Args:
            var (IntVar)
            domain (list[int]): A list of two integers representing the new domain.
        """
        assert len(domain) == 2, (
            f"Domain must be a list of two integers; {domain} given."
        )

        var.Proto().domain[:] = domain

    # objective functions

    def is_maximize(self) -> bool:
        """
        Returns:
            bool: True if the objective is maximize, False if minimize.
        """
        proto = self.Proto()
        # If the objective is not set to maximize, it defaults to minimize
        # in OR-Tools, so we return False if maximize is not set.
        return getattr(proto.objective, "maximize", False)

    # constraint functions

    def _get_constraints(self) -> RepeatedCompositeFieldContainer[ConstraintProto]:
        proto = self.Proto()
        if not hasattr(proto, "constraints") or proto.constraints is None:
            raise RuntimeError("No constraints defined in the model.")
        return proto.constraints

    def get_next_constr_idx(self) -> int:
        """Returns the index of the next constraint.

        Returns:
            int: The index of the next constraint.
        """
        return len(self._get_constraints())

    def set_num_base_constraints(self) -> None:
        """Sets the base number of constraints to the current number of constraints."""
        self.num_base_constraints = self.get_next_constr_idx()

    def add_linear_constraint_fast(
        self,
        var_list: Sequence[IntVar],
        coeff_list: Sequence[int],
        domain: tuple[int, int],
    ) -> None:
        """Adds a linear constraint of the form:

        domain[0] <= sum(coeff_list[i] * var_list[i]) <= domain[1]

        Args:
            var_list (Sequence[IntVar]): The list of integer variables.
            coeff_list (Sequence[int]): Matching coefficients for each variable.
            domain (tuple[int, int]): Lower and upper bound (inclusive) for the sum.
        """
        # assert len(var_list) == len(coeff_list), (
        #     "Length of var_list and coeff_list must be the same."
        # )
        # assert len(domain) == 2, "Domain must be a tuple of two integers."

        ct = self.proto.constraints.add()
        ct.linear.vars.extend([v.Index() for v in var_list])
        ct.linear.coeffs.extend(coeff_list)
        ct.linear.domain.extend(domain)

    def add_linear_constraint_enforced_fast(
        self,
        var_list: Sequence[IntVar],
        coeff_list: Sequence[int],
        domain: tuple[int, int],
        enforcers: Sequence[LiteralT],
    ) -> None:
        """
        Adds a linear constraint of the form:

        domain[0] <= sum(coeff_list[i] * var_list[i]) <= domain[1]

        that is *conditionally enforced* by the given literals.

        Args:
            var_list (Sequence[IntVar]): The list of integer variables.
            coeff_list (Sequence[int]): Matching coefficients for each variable.
            domain (Tuple[int, int]): Lower and upper bound (inclusive) for the sum.
            enforcers (Sequence[LiteralT]): One or more boolean literals whose
                conjunction enforces the constraint. If all literals are true,
                the linear inequality must hold; otherwise it is ignored.

        Example:
            To enforce `end_j - start_jp <= 0` only if arc(j, jp) is selected:

                self.add_linear_constraint_enforced_fast(
                    [end_j, start_jp],
                    [1, -1],
                    (-self.horizon, 0),
                    [arc_j_jp]
                )
        """
        # assert len(var_list) == len(coeff_list), (
        #     "Length of var_list and coeff_list must be the same."
        # )
        # assert len(domain) == 2, "Domain must be a tuple of two integers."
        # lb, ub = int(domain[0]), int(domain[1])
        # assert lb <= ub, f"Invalid domain: {domain}"

        ct = self.proto.constraints.add()
        ct.linear.vars.extend([v.Index() for v in var_list])
        ct.linear.coeffs.extend(coeff_list)
        ct.linear.domain.extend(domain)
        ct.enforcement_literal.extend(
            [self.get_or_make_boolean_index(lit) for lit in enforcers]
        )

    # methods to delete constraints

    def delete_constraints(self, idx_start: int, idx_end: int) -> None:
        del self._get_constraints()[idx_start:idx_end]

    def delete_added_constraints(self):
        """Deletes all constraints added after base model was built.

        Raises:
            ValueError: If no constraints were added after the base model was built.
        """

        if self.num_base_constraints == 0:
            raise ValueError("No base model constraints defined.")
        current_num_constraints = self.get_next_constr_idx()
        if current_num_constraints > self.num_base_constraints:
            self.delete_constraints(self.num_base_constraints, current_num_constraints)


CustomCpModelT = TypeVar("CustomCpModelT", bound=CustomCpModel)
"""
Type variable for CustomCpModel, allowing methods to specify
that they return or accept an instance of CustomCpModel or its subclasses.
"""
