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

enum SearchMethod {
  start,
  bus,
  transfer,
  walk;

  factory SearchMethod.fromJson(String json) => values.byName(json);
}

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

  static Iterable<Stop> findStopsNear(Coordinates location) {
    const numStops = 5;
    log("Finding stops near $location...");
    final stopDistances = [
      for (final stop in stops.values)
        (stop, stop.coordinates.distanceTo(location)),
    ];
    stopDistances.sort((a, b) => a.$2.compareTo(b.$2));
    return stopDistances.take(numStops).map((pair) => pair.$1);
  }

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

  factory StopState.start({
    required StopID stopID,
    required StopID goalID,
    required RouteID routeID,
    required Coordinates startPoint,
    required Coordinates endPoint,
  }) {
    final stop = stops[stopID]!;
    final goal = stops[goalID]!;
    final distanceToStart = startPoint.distanceTo(stop.coordinates);
    final distanceToEnd = endPoint.distanceTo(goal.coordinates);
    return StopState(
      stopID: stopID,
      goalID: goalID,
      routeID: routeID,
      distanceWalked: distanceToStart + distanceToEnd,
      depth: 0,
      method: SearchMethod.start,
    );
  }

  factory StopState.fromJson(Json json) => StopState(
    stopID: StopID.fromJson(json["stop_id"]),
    goalID: StopID.fromJson(json["goal_id"]),
    routeID: RouteID.fromJson(json["route_id"]),
    method: SearchMethod.fromJson(json["method"]),
    depth: 0,
    distanceWalked: 0,
  );

  @override
  String hash() {
    final stop = stops[stopID];
    final goal = stops[goalID];
    final route = routes[routeID];
    return "$distanceWalked-${method.name}-$stop-$goal-$route";
  }

  @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];
    if (route == null) return [];
    final stop = stops[stopID]!;
    final stopIndex = route.stops.indexOf(stopID);
    if (stopIndex == -1) return [];
    final result = <StopState>[];

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

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

    // 3) Make any available transfers
    for (final otherRouteID in stop.routeIDs.difference({routeID})) {
      final state = StopState(
        stopID: stopID,
        goalID: goalID,
        routeID: otherRouteID,
        depth: depth + 5,
        distanceWalked: distanceWalked,
        method: SearchMethod.transfer,
      );
      result.add(state);
    }

    // 4) Walk to a new stop
    for (final otherStop in findStopsNear(stop.coordinates)) {
      final distanceToOther = stop.coordinates.distanceTo(otherStop.coordinates);
      for (final otherRouteID in otherStop.routeIDs) {
        final state = StopState(
          stopID: otherStop.id,
          goalID: goalID,
          routeID: otherRouteID,
          depth: depth + 5 + 10*distanceToOther,
          distanceWalked: distanceWalked + distanceToOther,
          method: SearchMethod.walk,
        );
        result.add(state);
      }
    }

    return result;
  }

  @override
  Json toJson() => {
    "stop_id": stopID,
    "goal_id": goalID,
    "route_id": routeID,
    "method": method.name,
  };
}
