import "package:a_star/a_star.dart";
import "package:shared/data.dart";

class StopState extends AStarState<StopState> {
  static late final Map<RouteID, Route> routes;
  static late final Map<StopID, Stop> stops;

  static void init(Iterable<Route> allRoutes, Iterable<Stop> allStops) {
    routes = {
      for (final route in allRoutes)
        route.id: route,
    };
    stops = {
      for (final stop in allStops)
        stop.id: stop,
    };
  }

  final StopID stopID;
  final StopID goalID;
  final RouteID routeID;
  StopState({
    required this.stopID,
    required this.goalID,
    required this.routeID,
    required super.depth,
  });

  @override
  String hash() => "$stopID-$goalID-$routeID";

  @override
  bool isGoal() => stopID == goalID;

  @override
  double heuristic() {
    final stop = stops[stopID]!;
    final goal = stops[goalID]!;
    return stop.coordinates.distanceTo(goal.coordinates);
  }

  @override
  Iterable<StopState> expand() {
    final route = routes[routeID]!;
    final stop = stops[stopID]!;
    final stopIndex = route.stops.indexOf(stopID);
    final result = <StopState>[];

    // 1): Go forward one stop on the same route
    final nextIndex = stopIndex + 1;
    if (nextIndex < route.stops.length) {
      final nextStop = route.stops[nextIndex];
      final state = StopState(stopID: nextStop, goalID: goalID, routeID: routeID, depth: depth + 1);
      result.add(state);
    }

    // 2) Go back one stop on the same route
    final prevIndex = stopIndex - 1;
    if (prevIndex > 0) {
      final prevStop = route.stops[prevIndex];
      final state = StopState(stopID: prevStop, goalID: goalID, routeID: routeID, depth: depth + 1);
      result.add(state);
    }

    // 3) Make any available transfers
    for (final otherRoute in stop.routeIDs.difference({routeID})) {
      final state = StopState(stopID: stopID, goalID: goalID, routeID: otherRoute, depth: depth + 1);
      result.add(state);
    }

    return result;
  }
}
