"""Optimization solvers for linear and mixed-integer programming.

Wraps OR-Tools behind the provider pattern for LP/MIP formulation and solving.
A PuLP stub is included for future expansion.

Usage via provider system:
    from platoon.learning.providers import get_provider
    opt = get_provider("optimization")
    result = opt.solve(problem)

Or use the convenience function:
    from platoon.learning.optimization import solve
    result = solve(problem)

Expression syntax (parsed by parse_linear_expression):
    "2*x + 3*y"          — coefficients with explicit multiply
    "x + y"              — implicit coefficient of 1
    "-3*z"               — negative coefficients
    "x - 2*y + 3*z"      — subtraction treated as negative term
    Limitation: var*coeff form (e.g. "x*2") is NOT supported.
"""

import re
from enum import Enum
from typing import Any, Dict, List, Optional, Tuple

from pydantic import BaseModel, Field

from platoon.learning.providers import ModelProvider, register_provider


# ---------------------------------------------------------------------------
# Models
# ---------------------------------------------------------------------------


class ProblemSense(str, Enum):
    """Optimization direction."""

    MINIMIZE = "minimize"
    MAXIMIZE = "maximize"


class VariableType(str, Enum):
    """Variable domain type."""

    CONTINUOUS = "continuous"
    INTEGER = "integer"
    BINARY = "binary"


class Variable(BaseModel):
    """Decision variable definition."""

    name: str
    lb: float = 0.0
    ub: float = Field(default=float("inf"))
    var_type: VariableType = VariableType.CONTINUOUS


class Constraint(BaseModel):
    """Linear constraint in string form.

    Expression uses coeff*var notation:
        "2*x + 3*y <= 10"
        "x + y >= 5"
        "x - 2*y == 0"
    """

    expression: str
    name: Optional[str] = None


class OptimizationProblem(BaseModel):
    """Full LP/MIP problem specification.

    Designed for LLM structured output — all fields are JSON-serializable
    and the expression strings are human-readable.
    """

    name: str = "problem"
    sense: ProblemSense = ProblemSense.MAXIMIZE
    variables: List[Variable]
    objective: str  # linear expression string, e.g. "5*x + 4*y"
    constraints: List[Constraint]


class VariableResult(BaseModel):
    """Solved value for a single variable."""

    name: str
    value: float


class OptimizationResult(BaseModel):
    """Solution returned by a solver backend."""

    status: str  # "optimal", "infeasible", "unbounded", "error"
    objective_value: Optional[float] = None
    variables: List[VariableResult] = Field(default_factory=list)
    solver: Optional[str] = None
    metadata: Dict[str, Any] = Field(default_factory=dict)


# ---------------------------------------------------------------------------
# Expression parser
# ---------------------------------------------------------------------------

# Matches terms like "2*x", "-3.5*var_name", "+x", "-y", "x"
_TERM_RE = re.compile(
    r"([+-]?)\s*(\d+\.?\d*)?\s*\*?\s*([A-Za-z_]\w*)"
)


def parse_linear_expression(expr: str) -> List[Tuple[float, str]]:
    """Parse a linear expression string into (coefficient, variable) pairs.

    Supports:
        "2*x + 3*y"      → [(2.0, "x"), (3.0, "y")]
        "x - 2*y"        → [(1.0, "x"), (-2.0, "y")]
        "-3*z"           → [(-3.0, "z")]
        "x"              → [(1.0, "x")]
        "2.5*x_1"        → [(2.5, "x_1")]

    Args:
        expr: Linear expression string using coeff*var notation.

    Returns:
        List of (coefficient, variable_name) tuples.

    Raises:
        ValueError: If no valid terms are found.
    """
    terms = []
    for match in _TERM_RE.finditer(expr):
        sign_str, coeff_str, var_name = match.groups()
        sign = -1.0 if sign_str == "-" else 1.0
        coeff = float(coeff_str) if coeff_str else 1.0
        terms.append((sign * coeff, var_name))

    if not terms:
        raise ValueError(f"Could not parse any terms from expression: '{expr}'")

    return terms


def parse_constraint(expr: str) -> Tuple[List[Tuple[float, str]], str, float]:
    """Parse a constraint string into LHS terms, operator, and RHS constant.

    Args:
        expr: Constraint string like "2*x + 3*y <= 10"

    Returns:
        (terms, operator, rhs) where terms is from parse_linear_expression,
        operator is "<=", ">=", or "==", and rhs is the float constant.

    Raises:
        ValueError: If the constraint cannot be parsed.
    """
    # Split on operator
    for op in ("<=", ">=", "=="):
        if op in expr:
            lhs_str, rhs_str = expr.split(op, 1)
            terms = parse_linear_expression(lhs_str.strip())
            rhs = float(rhs_str.strip())
            return terms, op, rhs

    raise ValueError(f"No operator (<=, >=, ==) found in constraint: '{expr}'")


# ---------------------------------------------------------------------------
# OR-Tools provider
# ---------------------------------------------------------------------------


def _ortools_available() -> bool:
    try:
        from ortools.linear_solver import pywraplp  # noqa: F401

        return True
    except ImportError:
        return False


@register_provider
class ORToolsProvider(ModelProvider):
    """Optimization solver backed by Google OR-Tools.

    Uses GLOP for pure LP problems and CBC for MIP (integer/binary variables).
    Auto-detects solver type from variable definitions.
    """

    name = "ortools-optimization"
    domain = "optimization"
    backend = "ortools"

    @classmethod
    def is_available(cls) -> bool:
        return _ortools_available()

    def solve(self, problem: OptimizationProblem) -> OptimizationResult:
        """Solve an LP or MIP problem.

        Args:
            problem: Fully specified optimization problem.

        Returns:
            OptimizationResult with status, objective value, and variable values.
        """
        from ortools.linear_solver import pywraplp

        # Pick solver: CBC for MIP, GLOP for pure LP
        has_integers = any(
            v.var_type != VariableType.CONTINUOUS for v in problem.variables
        )
        solver_id = "CBC" if has_integers else "GLOP"
        solver = pywraplp.Solver.CreateSolver(solver_id)

        if solver is None:
            return OptimizationResult(
                status="error",
                metadata={"error": f"Could not create {solver_id} solver"},
            )

        # Create variables
        var_map: Dict[str, Any] = {}
        for v in problem.variables:
            ub = v.ub if v.ub != float("inf") else solver.infinity()
            if v.var_type == VariableType.BINARY:
                var_map[v.name] = solver.BoolVar(v.name)
            elif v.var_type == VariableType.INTEGER:
                var_map[v.name] = solver.IntVar(int(v.lb), int(ub), v.name)
            else:
                var_map[v.name] = solver.NumVar(v.lb, ub, v.name)

        # Set objective
        objective = solver.Objective()
        for coeff, var_name in parse_linear_expression(problem.objective):
            if var_name not in var_map:
                return OptimizationResult(
                    status="error",
                    metadata={"error": f"Unknown variable in objective: '{var_name}'"},
                )
            objective.SetCoefficient(var_map[var_name], coeff)

        if problem.sense == ProblemSense.MAXIMIZE:
            objective.SetMaximization()
        else:
            objective.SetMinimization()

        # Add constraints
        for c in problem.constraints:
            terms, op, rhs = parse_constraint(c.expression)
            ct_name = c.name or ""

            if op == "<=":
                ct = solver.Constraint(-solver.infinity(), rhs, ct_name)
            elif op == ">=":
                ct = solver.Constraint(rhs, solver.infinity(), ct_name)
            else:  # ==
                ct = solver.Constraint(rhs, rhs, ct_name)

            for coeff, var_name in terms:
                if var_name not in var_map:
                    return OptimizationResult(
                        status="error",
                        metadata={
                            "error": f"Unknown variable in constraint: '{var_name}'"
                        },
                    )
                ct.SetCoefficient(var_map[var_name], coeff)

        # Solve
        raw_status = solver.Solve()

        status_map = {
            pywraplp.Solver.OPTIMAL: "optimal",
            pywraplp.Solver.FEASIBLE: "feasible",
            pywraplp.Solver.INFEASIBLE: "infeasible",
            pywraplp.Solver.UNBOUNDED: "unbounded",
            pywraplp.Solver.ABNORMAL: "error",
            pywraplp.Solver.NOT_SOLVED: "error",
        }
        status = status_map.get(raw_status, "error")

        if status in ("optimal", "feasible"):
            return OptimizationResult(
                status=status,
                objective_value=objective.Value(),
                variables=[
                    VariableResult(name=name, value=var.solution_value())
                    for name, var in var_map.items()
                ],
                solver=solver_id,
                metadata={
                    "iterations": solver.iterations(),
                    "wall_time_ms": solver.wall_time(),
                },
            )

        return OptimizationResult(status=status, solver=solver_id)


# ---------------------------------------------------------------------------
# PuLP provider (stub)
# ---------------------------------------------------------------------------


@register_provider
class PuLPProvider(ModelProvider):
    """Placeholder for a PuLP-backed optimization solver.

    PuLP provides a similar LP/MIP interface. This stub is here so the
    provider can be discovered; actual solving is not yet implemented.
    """

    name = "pulp-optimization"
    domain = "optimization"
    backend = "pulp"

    @classmethod
    def is_available(cls) -> bool:
        try:
            import pulp  # noqa: F401

            return True
        except ImportError:
            return False

    def solve(self, problem: OptimizationProblem) -> OptimizationResult:
        raise NotImplementedError(
            "PuLP provider is a placeholder. Use the OR-Tools backend."
        )


# ---------------------------------------------------------------------------
# Convenience function
# ---------------------------------------------------------------------------


def solve(
    problem: OptimizationProblem, backend: Optional[str] = None
) -> OptimizationResult:
    """Solve an optimization problem using the best available backend.

    Args:
        problem: Fully specified optimization problem.
        backend: Optional backend name ("ortools" or "pulp").

    Returns:
        OptimizationResult with status, objective value, and variable values.
    """
    from platoon.learning.providers import get_provider

    provider = get_provider("optimization", backend=backend)
    return provider.solve(problem)
