package propertymap;

import com.conveyal.r5.OneOriginResult;
import com.conveyal.r5.analyst.FreeFormPointSet;
import com.conveyal.r5.analyst.PointSet;
import com.conveyal.r5.analyst.StreetTimesAndModes;
import com.conveyal.r5.analyst.TravelTimeComputer;
import com.conveyal.r5.analyst.WebMercatorExtents;
import com.conveyal.r5.analyst.cluster.PathResult;
import com.conveyal.r5.analyst.cluster.RegionalTask;
import com.conveyal.r5.analyst.cluster.TravelTimeResult;
import com.conveyal.r5.api.util.LegMode;
import com.conveyal.r5.api.util.TransitModes;
import com.conveyal.r5.kryo.KryoNetworkSerializer;
import com.conveyal.r5.transit.TransitLayer;
import com.conveyal.r5.transit.TransportNetwork;
import com.conveyal.r5.transit.path.RouteSequence;
import com.google.common.collect.Multimap;
import org.locationtech.jts.geom.Coordinate;

import java.io.File;
import java.time.LocalDate;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.EnumSet;
import java.util.List;

/** R5 routing: network loading, spatial filtering, travel time computation. */
public class Router {

    private static final int ZOOM = 9; // R5 enforces range 9-12
    private static final int MAX_GRID_CELLS = 4_900_000; // under R5's 5M limit
    private static final int DEPARTURE_FROM_TIME = 7 * 3600 + 30 * 60; // 07:30
    private static final int DEPARTURE_TO_TIME = 8 * 3600 + 30 * 60;   // 08:30
    private static final int MAX_TRIP_DURATION_MINUTES = 90;

    /** R5 PathResult throws if destinations > 5000. Chunks must be smaller when recording paths. */
    private static final int PATH_MAX_DESTINATIONS = 5000;

    // Percentile indices in R5 result arrays (order must match task.percentiles in buildTask)
    private static final int PERCENTILE_BEST = 0;   // 5th percentile (transit only)
    private static final int PERCENTILE_MEDIAN = 1;  // 50th percentile (transit: index 1, others: index 0)

    /** Result of computing travel times for a single origin with spatial pre-filtering. */
    record FilteredResult(int[] originalIndices, short[] times, short[] bestTimes, String[] journeys) {}

    /** Max plausible travel radius in km for {@link #MAX_TRIP_DURATION_MINUTES}-minute trips. */
    static double maxRadiusKm(String mode) {
        return switch (mode) {
            case "car" -> 150;
            case "transit" -> 150;
            case "bicycle" -> 60;
            case "walking" -> 12;
            default -> throw new IllegalArgumentException("Unknown mode: " + mode);
        };
    }

    /**
     * Load or build the transport network with Kryo caching.
     * The returned network is read-only after buildDistanceTables — safe for concurrent use.
     */
    static TransportNetwork loadNetwork(String dataDir, String cacheDir) throws Exception {
        System.err.println("Loading transport network...");
        File cacheFile = new File(cacheDir, "network.dat");
        TransportNetwork network;

        if (cacheFile.exists()) {
            System.err.println("  Loading cached network from " + cacheFile);
            network = KryoNetworkSerializer.read(cacheFile);
        } else {
            System.err.println("  Building network (first run, takes a few minutes)...");
            network = TransportNetwork.fromDirectory(new File(dataDir));
            new File(cacheDir).mkdirs();
            KryoNetworkSerializer.write(network, cacheFile);
            System.err.println("  Cached to " + cacheFile);
        }

        System.err.println("  Building distance tables...");
        network.transitLayer.buildDistanceTables(null);
        System.err.println("  Network ready");
        return network;
    }

    /**
     * Filter destinations by distance, build chunks, compute travel times for one origin.
     * Returns only the filtered subset indices and their travel times.
     */
    static FilteredResult computeForOrigin(
            TransportNetwork network,
            double[] allLats, double[] allLons,
            double originLat, double originLon,
            String mode, LocalDate date, boolean enablePaths) {

        double maxRadius = maxRadiusKm(mode);

        // 1. Filter destinations by bounding box
        int[] filtered = filterByDistance(allLats, allLons, originLat, originLon, maxRadius);
        if (filtered.length == 0) {
            return new FilteredResult(new int[0], new short[0], null, null);
        }

        // 2. Extract filtered coordinate arrays
        double[] fLats = new double[filtered.length];
        double[] fLons = new double[filtered.length];
        for (int i = 0; i < filtered.length; i++) {
            fLats[i] = allLats[filtered[i]];
            fLons[i] = allLons[filtered[i]];
        }

        // 3. Build chunks — smaller when path recording is active (R5 PathResult limit: 5000)
        boolean isTransit = mode.equals("transit");
        boolean recordPaths = isTransit && enablePaths;
        int maxDestsPerChunk = recordPaths ? PATH_MAX_DESTINATIONS : Integer.MAX_VALUE;
        List<DestinationChunk> chunks = buildDestinationChunks(fLats, fLons, maxDestsPerChunk);

        // 4. Compute travel times (and optionally paths)
        String[] journeys = recordPaths ? new String[fLats.length] : null;
        short[][] allTimes = computeTravelTimes(
            network, chunks, originLat, originLon, mode, fLats.length, date,
            recordPaths, journeys);

        // Transit requests [5th, 50th] percentiles; others request [50th] only
        short[] medianTimes = isTransit ? allTimes[PERCENTILE_MEDIAN] : allTimes[0];
        short[] bestTimes = isTransit ? allTimes[PERCENTILE_BEST] : null;
        return new FilteredResult(filtered, medianTimes, bestTimes, journeys);
    }

    /**
     * Filter destination indices to those within a bounding box of maxRadiusKm from origin.
     * Uses degree-based approximation — slightly overestimates at corners, which is fine.
     */
    private static int[] filterByDistance(
            double[] lats, double[] lons,
            double originLat, double originLon,
            double maxRadiusKm) {

        double degLat = maxRadiusKm / 111.0;
        double degLon = maxRadiusKm / (111.0 * Math.cos(Math.toRadians(originLat)));

        double minLat = originLat - degLat;
        double maxLat = originLat + degLat;
        double minLon = originLon - degLon;
        double maxLon = originLon + degLon;

        // Two-pass: count then fill (avoids ArrayList/boxing overhead)
        int count = 0;
        for (int i = 0; i < lats.length; i++) {
            if (lats[i] >= minLat && lats[i] <= maxLat && lons[i] >= minLon && lons[i] <= maxLon) {
                count++;
            }
        }

        int[] result = new int[count];
        int j = 0;
        for (int i = 0; i < lats.length; i++) {
            if (lats[i] >= minLat && lats[i] <= maxLat && lons[i] >= minLon && lons[i] <= maxLon) {
                result[j++] = i;
            }
        }
        return result;
    }

    /**
     * Split destinations into geographic chunks that each fit within R5's grid cell limit
     * and optionally a maximum destination count (required for path recording).
     * Sorts by latitude and splits into bands.
     */
    private static List<DestinationChunk> buildDestinationChunks(
            double[] lats, double[] lons, int maxDestsPerChunk) {
        int n = lats.length;

        // Sort indices by latitude for geographic chunking
        Integer[] sorted = new Integer[n];
        for (int i = 0; i < n; i++) sorted[i] = i;
        Arrays.sort(sorted, (a, b) -> Double.compare(lats[a], lats[b]));

        // Determine grid width (longitude span is the same for all chunks)
        double minLon = Double.MAX_VALUE, maxLon = -Double.MAX_VALUE;
        for (double lon : lons) {
            minLon = Math.min(minLon, lon);
            maxLon = Math.max(maxLon, lon);
        }
        int totalPixels = 256 << ZOOM;
        int gridWidth = lonToPixel(maxLon, totalPixels) - lonToPixel(minLon, totalPixels) + 1;
        int maxHeight = MAX_GRID_CELLS / gridWidth;

        // Greedily build chunks: extend each band until it would exceed maxHeight or maxDestsPerChunk
        List<DestinationChunk> chunks = new ArrayList<>();
        int start = 0;
        while (start < n) {
            int end = start + 1;
            int topPixel = latToPixel(lats[sorted[start]], totalPixels);

            while (end < n) {
                if (end - start >= maxDestsPerChunk) break;
                int bottomPixel = latToPixel(lats[sorted[end]], totalPixels);
                if (Math.abs(bottomPixel - topPixel) + 1 > maxHeight) break;
                end++;
            }

            chunks.add(buildChunk(lats, lons, sorted, start, end));
            start = end;
        }

        return chunks;
    }

    /**
     * Compute travel times from one origin to all destinations across all chunks.
     * Returns one short[] per requested percentile (transit gets 2: best + median, others get 1: median).
     * When recordPaths is true, also extracts journey instructions into journeysOut.
     */
    private static short[][] computeTravelTimes(
            TransportNetwork network, List<DestinationChunk> chunks,
            double originLat, double originLon, String mode, int nDest, LocalDate date,
            boolean recordPaths, String[] journeysOut) {

        boolean isTransit = mode.equals("transit");
        int nPercentiles = isTransit ? 2 : 1;
        short[][] allTimes = new short[nPercentiles][nDest];
        for (short[] arr : allTimes) Arrays.fill(arr, (short) -1);

        for (DestinationChunk chunk : chunks) {
            RegionalTask task = buildTask(chunk, originLat, originLon, mode, date, recordPaths);
            TravelTimeComputer computer = new TravelTimeComputer(task, network);
            OneOriginResult result = computer.computeTravelTimes();

            TravelTimeResult tt = result.travelTimes;
            if (tt == null) {
                throw new RuntimeException("R5 returned null travelTimes for chunk with "
                    + chunk.originalIndices.length + " destinations");
            }
            int[][] values = tt.getValues();
            if (values.length < nPercentiles) {
                throw new RuntimeException("R5 returned " + values.length + " percentiles, expected "
                    + nPercentiles);
            }
            for (int p = 0; p < nPercentiles; p++) {
                if (values[p].length < chunk.originalIndices.length) {
                    throw new RuntimeException("R5 returned " + values[p].length
                        + " travel times for percentile " + p + ", expected "
                        + chunk.originalIndices.length);
                }
                for (int i = 0; i < chunk.originalIndices.length; i++) {
                    if (values[p][i] != Integer.MAX_VALUE) {
                        allTimes[p][chunk.originalIndices[i]] = (short) values[p][i];
                    }
                }
            }

            // Extract path data for transit
            if (recordPaths && journeysOut != null && result.paths != null) {
                extractPaths(result.paths, chunk.originalIndices, network.transitLayer, journeysOut);
            }
        }
        return allTimes;
    }

    /**
     * Extract the most common journey pattern for each destination in a chunk.
     * Produces a JSON array of legs: [{mode, from?, to?, minutes}, ...].
     */
    @SuppressWarnings("unchecked")
    private static void extractPaths(
            PathResult paths, int[] originalIndices, TransitLayer transitLayer,
            String[] journeysOut) {
        Multimap<RouteSequence, PathResult.Iteration>[] allPaths = paths.iterationsForPathTemplates;
        for (int i = 0; i < originalIndices.length && i < allPaths.length; i++) {
            Multimap<RouteSequence, PathResult.Iteration> destPaths = allPaths[i];
            if (destPaths == null || destPaths.isEmpty()) continue;

            // Find the RouteSequence used by the most departure-time iterations
            RouteSequence bestRoute = null;
            int maxCount = 0;
            for (RouteSequence rs : destPaths.keySet()) {
                int count = destPaths.get(rs).size();
                if (count > maxCount) {
                    maxCount = count;
                    bestRoute = rs;
                }
            }
            if (bestRoute == null) continue;

            journeysOut[originalIndices[i]] = buildJourneyJson(bestRoute, transitLayer);
        }
    }

    /** Build a JSON array of journey legs from a RouteSequence. */
    private static String buildJourneyJson(RouteSequence routeSequence, TransitLayer transitLayer) {
        StringBuilder sb = new StringBuilder("[");
        boolean first = true;

        // Access leg (walk/bike to first transit stop)
        StreetTimesAndModes.StreetTimeAndMode access = routeSequence.stopSequence.access;
        if (access != null && access.time > 0) {
            int mins = (access.time + 30) / 60; // round to nearest minute
            sb.append("{\"mode\":\"").append(access.mode.name().toLowerCase())
              .append("\",\"minutes\":").append(mins).append("}");
            first = false;
        }

        // Transit legs
        for (RouteSequence.TransitLeg leg : routeSequence.transitLegs(transitLayer)) {
            if (!first) sb.append(",");
            sb.append("{\"mode\":\"").append(escapeJson(leg.route))
              .append("\",\"from\":\"").append(escapeJson(leg.board))
              .append("\",\"to\":\"").append(escapeJson(leg.alight))
              .append("\",\"minutes\":").append(Math.round(leg.inVehicleTime))
              .append("}");
            first = false;
        }

        // Egress leg (walk/bike from last transit stop)
        StreetTimesAndModes.StreetTimeAndMode egress = routeSequence.stopSequence.egress;
        if (egress != null && egress.time > 0) {
            if (!first) sb.append(",");
            int mins = (egress.time + 30) / 60;
            sb.append("{\"mode\":\"").append(egress.mode.name().toLowerCase())
              .append("\",\"minutes\":").append(mins).append("}");
        }

        sb.append("]");
        return sb.toString();
    }

    private static String escapeJson(String s) {
        if (s == null) return "";
        return s.replace("\\", "\\\\").replace("\"", "\\\"");
    }

    // --- Private helpers ---

    private record DestinationChunk(FreeFormPointSet pointSet, WebMercatorExtents extents, int[] originalIndices) {}

    private static DestinationChunk buildChunk(
            double[] lats, double[] lons, Integer[] sorted, int start, int end) {
        int size = end - start;
        int[] originalIndices = new int[size];
        Coordinate[] coords = new Coordinate[size];
        double minLat = Double.MAX_VALUE, maxLat = -Double.MAX_VALUE;
        double minLon = Double.MAX_VALUE, maxLon = -Double.MAX_VALUE;

        for (int i = 0; i < size; i++) {
            int idx = sorted[start + i];
            originalIndices[i] = idx;
            double lat = lats[idx], lon = lons[idx];
            coords[i] = new Coordinate(lon, lat); // x=lon, y=lat
            minLat = Math.min(minLat, lat);
            maxLat = Math.max(maxLat, lat);
            minLon = Math.min(minLon, lon);
            maxLon = Math.max(maxLon, lon);
        }

        FreeFormPointSet pointSet = new FreeFormPointSet(coords);
        int totalPixels = 256 << ZOOM;
        int west = lonToPixel(minLon, totalPixels);
        int north = latToPixel(maxLat, totalPixels);
        int width = lonToPixel(maxLon, totalPixels) - west + 1;
        int height = latToPixel(minLat, totalPixels) - north + 1;
        WebMercatorExtents extents = new WebMercatorExtents(west, north, width, height, ZOOM);

        return new DestinationChunk(pointSet, extents, originalIndices);
    }

    private static RegionalTask buildTask(
            DestinationChunk chunk, double originLat, double originLon, String mode, LocalDate date,
            boolean recordPaths) {
        RegionalTask task = new RegionalTask();
        task.fromLat = originLat;
        task.fromLon = originLon;
        task.date = date;
        task.percentiles = mode.equals("transit") ? new int[]{5, 50} : new int[]{50};
        task.recordTimes = true;
        task.destinationPointSets = new PointSet[]{chunk.pointSet};
        task.zoom = chunk.extents.zoom;
        task.west = chunk.extents.west;
        task.north = chunk.extents.north;
        task.width = chunk.extents.width;
        task.height = chunk.extents.height;
        task.fromTime = DEPARTURE_FROM_TIME;
        task.toTime = DEPARTURE_TO_TIME;
        task.maxTripDurationMinutes = MAX_TRIP_DURATION_MINUTES;

        if (recordPaths) {
            task.includePathResults = true;
            task.nPathsPerTarget = 3;
        }

        configureMode(task, mode);
        return task;
    }

    private static void configureMode(RegionalTask task, String mode) {
        switch (mode) {
            case "car" -> setDirectMode(task, LegMode.CAR);
            case "bicycle" -> setDirectMode(task, LegMode.BICYCLE);
            case "walking" -> setDirectMode(task, LegMode.WALK);
            case "transit" -> {
                task.maxRides = 4;
                // R5 requires directModes ⊆ accessModes. BICYCLE egress is too expensive
                // (builds cost tables from 59k stops × N destinations), so keep WALK only
                // for egress and match access/direct to avoid the R5 validation error.
                task.accessModes = EnumSet.of(LegMode.WALK);
                task.egressModes = EnumSet.of(LegMode.WALK);
                task.directModes = EnumSet.of(LegMode.WALK);
                task.transitModes = EnumSet.allOf(TransitModes.class);
            }
            default -> throw new IllegalArgumentException("Unknown mode: " + mode);
        }
    }

    private static void setDirectMode(RegionalTask task, LegMode legMode) {
        task.maxRides = 0;
        task.accessModes = EnumSet.of(legMode);
        task.egressModes = EnumSet.of(legMode);
        task.directModes = EnumSet.of(legMode);
        task.transitModes = EnumSet.noneOf(TransitModes.class);
    }

    private static int lonToPixel(double lon, int totalPixels) {
        return (int) Math.floor(totalPixels * (lon + 180.0) / 360.0);
    }

    private static int latToPixel(double lat, int totalPixels) {
        double latRad = Math.toRadians(lat);
        return (int) Math.floor(totalPixels * (1.0 - Math.log(Math.tan(latRad) + 1.0 / Math.cos(latRad)) / Math.PI) / 2.0);
    }
}