// ---------------------------------------------------------------------------
// Shared types — keep Edge<N> in a shared types file if A* uses the same shape.
// ---------------------------------------------------------------------------

import type { Point } from "@common/geometry";
import { buildDijkstraOptionsFrom2D, MinHeap, type Dijkstra2DOptions, type DijkstraOptions, type Edge } from "@common/navigation/utils";

// ---------------------------------------------------------------------------
// DijkstraMap
// ---------------------------------------------------------------------------

export class DijkstraMap<N> {
    private readonly distances: Map<string, number> = new Map();
    private readonly nodeRefs: Map<string, N> = new Map();

    private readonly _key: (node: N) => string;
    private readonly _neighbours: (node: N) => Iterable<Edge<N>>;

    constructor(sources: Iterable<N>, opts: DijkstraOptions<N>) {
        this._key = opts.key ?? ((n) => String(n));
        this._neighbours = opts.neighbours;
        this.build(sources, opts.maxExpansions ?? Infinity);
    }

    private build(sources: Iterable<N>, maxExpansions: number): void {
        // HeapItem.f = HeapItem.g = distance — no heuristic, so f and g are
        // always equal. The existing _less comparator (lower f wins, ties broken
        // by higher g) reduces to a plain min-distance order, which is correct.
        const heap = new MinHeap<N>();

        for (const source of sources) {
            const nodeKey = this._key(source);
            if (!this.distances.has(nodeKey)) {
                this.distances.set(nodeKey, 0);
                this.nodeRefs.set(nodeKey, source);
                heap.push({ node: source, nodeKey, g: 0, f: 0 });
            }
        }

        let expansions = 0;

        while (heap.size > 0 && expansions < maxExpansions) {
            const { node, nodeKey, g: dist } = heap.pop()!;

            // Discard stale entries produced by relaxation.
            if (dist > this.distances.get(nodeKey)!) continue;

            expansions++;

            for (const edge of this._neighbours(node)) {
                const neighbourKey = this._key(edge.node);
                const candidate = dist + edge.cost;
                const existing = this.distances.get(neighbourKey);

                if (existing === undefined || candidate < existing) {
                    this.distances.set(neighbourKey, candidate);
                    this.nodeRefs.set(neighbourKey, edge.node);
                    // f = g = candidate: Dijkstra is A* with a zero heuristic.
                    heap.push({ node: edge.node, nodeKey: neighbourKey, g: candidate, f: candidate });
                }
            }
        }
    }

    // ---------------------------------------------------------------------------
    // Distance queries
    // ---------------------------------------------------------------------------

    /**
     * Shortest-path cost from `node` to the nearest source.
     * Returns `undefined` if the node was not reached (unreachable or beyond
     * `maxExpansions`).
     */
    distance(node: N): number | undefined {
        return this.distances.get(this._key(node));
    }

    /** `true` if `node` was reached during map construction. */
    reachable(node: N): boolean {
        return this.distances.has(this._key(node));
    }

    /** Total number of nodes reached. */
    get size(): number {
        return this.distances.size;
    }

    // ---------------------------------------------------------------------------
    // Gradient navigation
    // ---------------------------------------------------------------------------

    /**
     * The neighbour of `node` with the lowest map distance — the next step an
     * entity should take to move *toward* the nearest source.
     *
     * Only reachable neighbours are considered. Returns `undefined` when `node`
     * is isolated or all neighbours are unreachable.
     *
     * @example
     * const next = playerMap.approach(enemy.position);
     * if (next) enemy.moveTo(next);
     */
    approach(node: N): N | undefined {
        return this.gradient(node, 'min');
    }

    /**
     * The neighbour of `node` with the highest map distance — the next step an
     * entity should take to move *away from* all sources.
     *
     * @example
     * const next = playerMap.flee(enemy.position);
     * if (next) enemy.moveTo(next);
     */
    flee(node: N): N | undefined {
        return this.gradient(node, 'max');
    }

    private gradient(node: N, mode: 'min' | 'max'): N | undefined {
        let best: N | undefined;
        let bestDist = mode === 'min' ? Infinity : -Infinity;

        for (const edge of this._neighbours(node)) {
            const d = this.distances.get(this._key(edge.node));
            if (d === undefined) continue; // unreachable — never move there
            if (mode === 'min' ? d < bestDist : d > bestDist) {
                bestDist = d;
                best = edge.node;
            }
        }

        return best;
    }

    // ---------------------------------------------------------------------------
    // Iteration
    // ---------------------------------------------------------------------------

    /**
     * Iterates all `[node, distance]` pairs in the map.
     *
     * @example
     * for (const [node, dist] of playerMap.entries()) {
     *   renderer.tint(node, distToColor(dist));
     * }
     */
    *entries(): IterableIterator<[N, number]> {
        for (const [k, dist] of this.distances) {
            yield [this.nodeRefs.get(k)!, dist];
        }
    }

    [Symbol.iterator]() {
        return this.entries();
    }

    // ---------------------------------------------------------------------------
    // Static factory: influence map
    // ---------------------------------------------------------------------------

    /**
     * Combines multiple `DijkstraMap` layers into a single **influence map**.
     *
     * Each layer contributes `weight × distance(node)` to a total score.
     * Use **negative weights** to attract toward a source, **positive** to repel.
     *
     * `approach()` moves toward the lowest combined score (most desirable cell).
     * `flee()` moves toward the highest combined score (emergency escape).
     *
     * @example
     * const goblinAI = DijkstraMap.combine(
     *   [
     *     { map: playerMap, weight: -1.5 }, // chase
     *     { map: archerMap, weight:  2.0 }, // dodge archers
     *     { map: lootMap,   weight: -0.5 }, // weakly drawn to loot
     *   ],
     *   { neighbours: grid.walkable, key: grid.key },
     * );
     *
     * const next = goblinAI.approach(goblin.position);
     */
    static combine<N>(
        layers: ReadonlyArray<{ map: DijkstraMap<N>; weight: number }>,
        opts: DijkstraOptions<N>,
    ): InfluenceMap<N> {
        return new InfluenceMap(layers, opts);
    }
}

// ---------------------------------------------------------------------------
// InfluenceMap — weighted sum of Dijkstra layers
// ---------------------------------------------------------------------------

type Layers<N> = ReadonlyArray<{ map: DijkstraMap<N>; weight: number }>;

export class InfluenceMap<N> {
    private readonly layers: Layers<N>;
    private readonly _neighbours: (node: N) => Iterable<Edge<N>>;

    /** @internal — use `DijkstraMap.combine()` instead. */
    constructor(
        layers: Layers<N>,
        opts: DijkstraOptions<N>,
    ) {
        this.layers = layers;
        this._neighbours = opts.neighbours;
    }

    /**
     * Combined influence score for `node`: `Σ (weight_i × distance_i(node))`.
     *
     * Returns `undefined` if the node is unreachable in every layer.
     * Layers that don't reach the node are omitted from the sum rather than
     * treated as 0, so a node reachable by only some layers still gets a
     * meaningful score.
     */
    score(node: N): number | undefined {
        let total = 0;
        let hasAny = false;

        for (const { map, weight } of this.layers) {
            const d = map.distance(node);
            if (d !== undefined) {
                total += weight * d;
                hasAny = true;
            }
        }

        return hasAny ? total : undefined;
    }

    /** Step toward the lowest-scoring neighbour (most desirable position). */
    approach(node: N): N | undefined {
        return this.gradient(node, 'min');
    }

    /** Step toward the highest-scoring neighbour (least desirable position). */
    flee(node: N): N | undefined {
        return this.gradient(node, 'max');
    }

    private gradient(node: N, mode: 'min' | 'max'): N | undefined {
        let best: N | undefined;
        let bestScore = mode === 'min' ? Infinity : -Infinity;

        for (const edge of this._neighbours(node)) {
            const s = this.score(edge.node);
            if (s === undefined) continue;
            if (mode === 'min' ? s < bestScore : s > bestScore) {
                bestScore = s;
                best = edge.node;
            }
        }

        return best;
    }
}
// ---------------------------------------------------------------------------
// DijkstraMap2D — convenience subclass for 2-D grids
// ---------------------------------------------------------------------------

export class DijkstraMap2D extends DijkstraMap<Point> {
    constructor(sources: Iterable<Point>, options: Dijkstra2DOptions) {
        super(sources, buildDijkstraOptionsFrom2D(options));
    }
}

// ---------------------------------------------------------------------------
// InfluenceMap2D — convenience subclass for 2-D grids
// ---------------------------------------------------------------------------

export class InfluenceMap2D extends InfluenceMap<Point> {
    constructor(layers: Layers<Point>, options: Dijkstra2DOptions) {
        super(layers, buildDijkstraOptionsFrom2D(options));
    }
}
