import { z } from "zod";

import type { RunConfig } from "./runConfig.ts";

export const DEFAULT_MAX_COMBINATIONS = 64;
export const HARD_MAX_COMBINATIONS = 4096;
export const SWEEP_POINT_ANNOTATION = "sweep.point";
export const SWEEP_ENABLED_TARGET = "$enabled";

export const sweepTargetSchema = z.strictObject({
  instanceId: z.string().min(1),
  paramPath: z.string().min(1),
  valueKey: z.string().min(1).optional(),
});

export const sweepAxisSchema = z.strictObject({
  id: z.string().min(1),
  targets: z.array(sweepTargetSchema).min(1),
  values: z.array(z.json()).min(1),
});

export const sweepSchema = z.strictObject({
  axes: z.array(sweepAxisSchema).min(1),
  maxCombinations: z
    .number()
    .int()
    .positive()
    .max(HARD_MAX_COMBINATIONS)
    .default(DEFAULT_MAX_COMBINATIONS),
});

export type SweepTarget = z.infer<typeof sweepTargetSchema>;
export type SweepAxis = z.infer<typeof sweepAxisSchema>;
export type SweepConfig = z.infer<typeof sweepSchema>;
export type SweepPoint = Record<string, z.infer<ReturnType<typeof z.json>>>;

export interface SweepCell {
  readonly config: RunConfig;
  readonly point: SweepPoint;
}

export interface SweepSkippedCell {
  readonly point: SweepPoint;
  readonly reason: string;
}

export interface SweepStats {
  readonly combinations: number;
  feasibleCells: number;
  readonly skipped: SweepSkippedCell[];
}

export class SweepError extends Error {
  constructor(message: string) {
    super(message);
    this.name = "SweepError";
  }
}

export function sweepCombinationCount(sweep: SweepConfig): number {
  return sweep.axes.reduce((count, axis) => count * axis.values.length, 1);
}

/** Expand a filled config in declaration order, with the last axis varying fastest. */
export function expandSweepCells(filled: RunConfig): SweepCell[] {
  if (!filled.sweep) return [{ config: filled, point: {} }];

  const axes = filled.sweep.axes;
  const indexes = axes.map(() => 0);
  const count = sweepCombinationCount(filled.sweep);
  const cells: SweepCell[] = [];

  for (let combination = 0; combination < count; combination++) {
    const point: SweepPoint = {};
    const blocks = structuredClone(filled.pipeline.blocks);
    const byId = new Map(blocks.map((block) => [block.instanceId, block]));

    for (let axisIndex = 0; axisIndex < axes.length; axisIndex++) {
      const axis = axes[axisIndex]!;
      const axisPoint = structuredClone(axis.values[indexes[axisIndex]!]!);
      point[axis.id] = axisPoint;
      for (const target of axis.targets) {
        const block = byId.get(target.instanceId)!;
        const value =
          target.valueKey === undefined
            ? structuredClone(axisPoint)
            : structuredClone(
                (
                  axisPoint as Record<
                    string,
                    z.infer<ReturnType<typeof z.json>>
                  >
                )[target.valueKey],
              );
        if (target.paramPath === SWEEP_ENABLED_TARGET) {
          block.enabled = value as boolean;
        } else {
          setDotPath(block.params, target.paramPath, value);
        }
      }
    }

    const withoutSweep = { ...filled };
    delete withoutSweep.sweep;
    cells.push({
      point,
      config: {
        ...withoutSweep,
        pipeline: { ...filled.pipeline, blocks },
      },
    });

    for (let axisIndex = axes.length - 1; axisIndex >= 0; axisIndex--) {
      indexes[axisIndex] = indexes[axisIndex]! + 1;
      if (indexes[axisIndex]! < axes[axisIndex]!.values.length) break;
      indexes[axisIndex] = 0;
    }
  }
  return cells;
}

export function expandSweep(filled: RunConfig): RunConfig[] {
  return expandSweepCells(filled).map((cell) => cell.config);
}

function setDotPath(
  target: Record<string, unknown>,
  path: string,
  value: unknown,
): void {
  const segments = path.split(".");
  let current = target;
  for (let i = 0; i < segments.length - 1; i++) {
    current = current[segments[i]!] as Record<string, unknown>;
  }
  current[segments.at(-1)!] = value;
}