from __future__ import annotations
from typing import List

# Try OR-Tools; fallback to a simple nearest-neighbor if unavailable
try:
    from ortools.constraint_solver import routing_enums_pb2
    from ortools.constraint_solver import pywrapcp
except Exception:  # pragma: no cover
    routing_enums_pb2 = None
    pywrapcp = None


def solve_tsp(durations: List[List[float]]) -> List[int]:
    """
    Solve a single-vehicle TSP with start/end at node 0 on the given duration matrix (seconds).
    Returns a list of node indices in visiting order, including the final 0 (start=end).
    Falls back to nearest-neighbor if OR-Tools is not available.
    """
    n = len(durations)
    if n == 0:
        return []
    if n == 1:
        return [0]

    # Fallback: nearest-neighbor tour + return to start
    if pywrapcp is None:
        order = [0]
        remaining = set(range(1, n))
        cur = 0
        while remaining:
            nxt = min(remaining, key=lambda j: durations[cur][j])
            order.append(nxt)
            remaining.remove(nxt)
            cur = nxt
        order.append(0)
        return order

    # OR-Tools exact/heuristic solution
    manager = pywrapcp.RoutingIndexManager(n, 1, 0)  # 1 vehicle, depot=0
    routing = pywrapcp.RoutingModel(manager)

    def time_cb(from_index, to_index):
        i = manager.IndexToNode(from_index)
        j = manager.IndexToNode(to_index)
        return int(max(0, durations[i][j]))

    transit_cb_index = routing.RegisterTransitCallback(time_cb)
    routing.SetArcCostEvaluatorOfAllVehicles(transit_cb_index)

    search_params = pywrapcp.DefaultRoutingSearchParameters()
    search_params.first_solution_strategy = routing_enums_pb2.FirstSolutionStrategy.PATH_CHEAPEST_ARC
    search_params.local_search_metaheuristic = routing_enums_pb2.LocalSearchMetaheuristic.GUIDED_LOCAL_SEARCH
    search_params.time_limit.FromSeconds(5)

    solution = routing.SolveWithParameters(search_params)
    if solution is None:
        # Fallback if solver fails
        order = [0]
        remaining = set(range(1, n))
        cur = 0
        while remaining:
            nxt = min(remaining, key=lambda j: durations[cur][j])
            order.append(nxt)
            remaining.remove(nxt)
            cur = nxt
        order.append(0)
        return order

    # Extract route
    index = routing.Start(0)
    order = []
    while not routing.IsEnd(index):
        order.append(manager.IndexToNode(index))
        index = solution.Value(routing.NextVar(index))
    order.append(manager.IndexToNode(index))  # return to depot
    return order