// Lazy schedule enumeration.
// Wraps an AsyncIterable<Schedule> so that filter/map/take/branch stay lazy
// and only the terminal ops (bestBy, bestByLex, kBestByLex, paretoFrontier,
// collect) actually drain.

import type { Schedule, ScheduleStream, Scorer } from "./types.ts";

export class ScheduleStreamImpl implements ScheduleStream {
  private readonly source: AsyncIterable<Schedule>;

  constructor(source: AsyncIterable<Schedule>) {
    this.source = source;
  }

  [Symbol.asyncIterator](): AsyncIterator<Schedule> {
    return this.source[Symbol.asyncIterator]();
  }

  filter(pred: (s: Schedule) => boolean): ScheduleStream {
    const source = this.source;
    return new ScheduleStreamImpl({
      async *[Symbol.asyncIterator]() {
        for await (const s of source) {
          if (pred(s)) yield s;
        }
      },
    });
  }

  map(fn: (s: Schedule) => Schedule): ScheduleStream {
    const source = this.source;
    return new ScheduleStreamImpl({
      async *[Symbol.asyncIterator]() {
        for await (const s of source) yield fn(s);
      },
    });
  }

  take(k: number): ScheduleStream {
    const source = this.source;
    return new ScheduleStreamImpl({
      async *[Symbol.asyncIterator]() {
        if (k <= 0) return;
        let n = 0;
        for await (const s of source) {
          yield s;
          if (++n >= k) return;
        }
      },
    });
  }

  branch(fork: (s: Schedule) => ScheduleStream): ScheduleStream {
    const source = this.source;
    return new ScheduleStreamImpl({
      async *[Symbol.asyncIterator]() {
        for await (const s of source) {
          for await (const child of fork(s)) yield child;
        }
      },
    });
  }

  async bestBy(scorer: Scorer): Promise<Schedule | null> {
    let bestSchedule: Schedule | null = null;
    let bestScore: number | null = null;
    for await (const s of this.source) {
      const score = scorer.score(s);
      if (
        bestScore === null ||
        (scorer.direction === "min" ? score < bestScore : score > bestScore)
      ) {
        bestScore = score;
        bestSchedule = s;
      }
    }
    return bestSchedule;
  }

  async bestByLex(scorers: readonly Scorer[]): Promise<Schedule | null> {
    let bestSchedule: Schedule | null = null;
    let bestScores: number[] | null = null;
    for await (const s of this.source) {
      const scores = scorers.map((sc) => sc.score(s));
      if (bestScores === null || lexBetter(scores, bestScores, scorers)) {
        bestScores = scores;
        bestSchedule = s;
      }
    }
    return bestSchedule;
  }

  async kBestByLex(
    scorers: readonly Scorer[],
    k: number,
  ): Promise<readonly Schedule[]> {
    if (k <= 0) return [];
    const scored: { schedule: Schedule; scores: number[] }[] = [];
    for await (const s of this.source) {
      scored.push({ schedule: s, scores: scorers.map((sc) => sc.score(s)) });
    }
    // Stable sort (ES2019 guarantees it): equal-on-all-tiers keeps stream
    // order, matching bestByLex where first seen wins.
    scored.sort((a, b) =>
      lexBetter(a.scores, b.scores, scorers)
        ? -1
        : lexBetter(b.scores, a.scores, scorers)
          ? 1
          : 0,
    );
    return scored.slice(0, k).map((x) => x.schedule);
  }

  async paretoFrontier(
    scorers: readonly Scorer[],
  ): Promise<readonly Schedule[]> {
    const collected: { schedule: Schedule; scores: number[] }[] = [];
    for await (const s of this.source) {
      collected.push({
        schedule: s,
        scores: scorers.map((sc) => sc.score(s)),
      });
    }
    const frontier: Schedule[] = [];
    for (const candidate of collected) {
      let dominated = false;
      for (const other of collected) {
        if (other === candidate) continue;
        if (dominates(other.scores, candidate.scores, scorers)) {
          dominated = true;
          break;
        }
      }
      if (!dominated) frontier.push(candidate.schedule);
    }
    return frontier;
  }

  async collect(limit?: number): Promise<readonly Schedule[]> {
    const out: Schedule[] = [];
    for await (const s of this.source) {
      out.push(s);
      if (limit !== undefined && out.length >= limit) break;
    }
    return out;
  }
}

// `a` beats `b` on the first tier where they differ; equal on all tiers is
// not better (first seen wins, keeping selection deterministic).
function lexBetter(
  a: readonly number[],
  b: readonly number[],
  scorers: readonly Scorer[],
): boolean {
  for (let i = 0; i < scorers.length; i++) {
    const ai = a[i]!;
    const bi = b[i]!;
    if (ai === bi) continue;
    return scorers[i]!.direction === "min" ? ai < bi : ai > bi;
  }
  return false;
}

// `a` dominates `b` iff a is at least as good on every scorer and strictly
// better on at least one.
function dominates(
  a: readonly number[],
  b: readonly number[],
  scorers: readonly Scorer[],
): boolean {
  let strictlyBetter = false;
  for (let i = 0; i < scorers.length; i++) {
    const dir = scorers[i]!.direction;
    const ai = a[i]!;
    const bi = b[i]!;
    const aBetter = dir === "min" ? ai < bi : ai > bi;
    const aWorse = dir === "min" ? ai > bi : ai < bi;
    if (aWorse) return false;
    if (aBetter) strictlyBetter = true;
  }
  return strictlyBetter;
}

export function streamFromFactory(
  factory: () => AsyncGenerator<Schedule>,
): ScheduleStream {
  return new ScheduleStreamImpl({
    [Symbol.asyncIterator]: factory,
  });
}

/** Lazily concatenate streams in declaration order. A later factory is not
 * invoked until every schedule from the preceding stream has been consumed. */
export function concatStreams(
  factories: readonly (() => ScheduleStream)[],
): ScheduleStream {
  return new ScheduleStreamImpl({
    async *[Symbol.asyncIterator]() {
      for (const factory of factories) {
        for await (const schedule of factory()) yield schedule;
      }
    },
  });
}