import fc, { type Arbitrary, type Command } from "fast-check";
import {
  PROPERTY_INVARIANT_CLASSES,
  type PropertyExecutionRecord,
  type PropertyInvariantClass,
  type StateMachineInvariantResult,
  type StateMachineTransitionMode,
} from "./types.js";

const DEFAULT_NUM_RUNS = 100;
const DEFAULT_SEED = 0x1a2b_3c4d;
const DEFAULT_MAX_COMMANDS = 12;

export type StateMachineInvariantSpec<
  Model extends object,
  Real,
  Cmd extends Command<Model, Real>,
> = {
  id: string;
  invariantClass: PropertyInvariantClass;
  setup: () => { model: Model; real: Real };
  commandArbs: Arbitrary<Cmd>[];
  transitionMode: StateMachineTransitionMode;
  postCondition?: (model: Model, real: Real) => void;
  expectRejection?: boolean;
  maxCommands?: number;
  numRuns?: number;
  seed?: number;
  path?: string;
};

function assertNonEmptyString(value: string, field: string): void {
  if (value.trim().length === 0) {
    throw new Error(`${field} must be a non-empty string`);
  }
}

function assertInvariantClass(value: string): asserts value is PropertyInvariantClass {
  if (!(PROPERTY_INVARIANT_CLASSES as readonly string[]).includes(value)) {
    throw new Error(`unknown invariant class: ${value}`);
  }
}

function validateStateMachineSpec<Model extends object, Real, Cmd extends Command<Model, Real>>(
  spec: StateMachineInvariantSpec<Model, Real, Cmd>,
): void {
  assertNonEmptyString(spec.id, "id");
  assertInvariantClass(spec.invariantClass);
  if (spec.commandArbs.length === 0) {
    throw new Error("commandArbs must contain at least one command arbitrary");
  }
}

function cloneValue<T>(value: T): T {
  if (typeof structuredClone === "function") {
    return structuredClone(value);
  }
  return JSON.parse(JSON.stringify(value)) as T;
}

function serializeCommand<Model extends object, Real>(
  command: Command<Model, Real>,
): string {
  return command.toString();
}

function snapshotModel<Model extends object>(model: Model): Model {
  return cloneValue(model);
}

function runLegalCommands<Model extends object, Real, Cmd extends Command<Model, Real>>(
  setup: () => { model: Model; real: Real },
  cmds: Iterable<Cmd>,
  postCondition?: (model: Model, real: Real) => void,
): void {
  const { model, real } = setup();
  for (const command of cmds) {
    if (!command.check(model)) {
      throw new Error(`illegal command precondition violated: ${command.toString()}`);
    }
    command.run(model, real);
  }
  postCondition?.(model, real);
}

function runIllegalCommands<Model extends object, Real, Cmd extends Command<Model, Real>>(
  setup: () => { model: Model; real: Real },
  cmds: Iterable<Cmd>,
  expectRejection: boolean,
): number {
  let rejected = 0;
  const { model, real } = setup();
  for (const command of cmds) {
    try {
      if (!command.check(model)) {
        rejected += 1;
        continue;
      }
      command.run(model, real);
      if (expectRejection) {
        throw new Error(`expected rejection for illegal command: ${command.toString()}`);
      }
    } catch {
      rejected += 1;
    }
  }
  return rejected;
}

function materializeCommands<Cmd>(commands: Iterable<Cmd>): Cmd[] {
  const result: Cmd[] = [];
  for (const command of commands) {
    result.push(command);
  }
  return result;
}

function buildStateMachineRecord<Model extends object, Real, Cmd extends Command<Model, Real>>(
  setup: () => { model: Model; real: Real },
  commands: Cmd[],
  seed: number,
  shrinkPath: string | null,
  counterexample: Cmd[] | null,
): PropertyExecutionRecord {
  const initial = setup();
  const final = setup();
  for (const command of commands) {
    try {
      if (typeof command.check === "function" && command.check(final.model)) {
        command.run(final.model, final.real);
      }
    } catch {
      // Illegal or rejected commands leave state unchanged for record capture.
    }
  }

  return {
    generatedInput: null,
    initialState: snapshotModel(initial.model),
    actionSequence: commands.map((command) => serializeCommand(command)),
    finalState: snapshotModel(final.model),
    seed,
    shrinkPath,
    minimalCounterexample:
      counterexample === null
        ? null
        : counterexample.map((command) => serializeCommand(command)),
  };
}

export function runStateMachineInvariant<
  Model extends object,
  Real,
  Cmd extends Command<Model, Real>,
>(spec: StateMachineInvariantSpec<Model, Real, Cmd>): StateMachineInvariantResult {
  validateStateMachineSpec(spec);

  const seed = spec.seed ?? DEFAULT_SEED;
  const numRuns = spec.numRuns ?? DEFAULT_NUM_RUNS;
  const maxCommands = spec.maxCommands ?? DEFAULT_MAX_COMMANDS;
  const expectRejection = spec.expectRejection ?? spec.transitionMode === "illegal";

  const commandArb = fc.commands(spec.commandArbs as Arbitrary<Command<Model, Real>>[], {
    maxCommands,
    ...(spec.path !== undefined ? { replayPath: spec.path } : {}),
  });

  const property =
    spec.transitionMode === "legal"
      ? fc.property(commandArb, (commands) => {
          runLegalCommands(spec.setup, materializeCommands(commands), spec.postCondition);
        })
      : fc.property(commandArb, (commands) => {
          const materialized = materializeCommands(commands);
          if (materialized.length === 0) {
            return;
          }
          const rejected = runIllegalCommands(spec.setup, materialized, expectRejection);
          if (rejected === 0) {
            throw new Error("illegal exploration produced no rejected commands");
          }
        });

  const details = fc.check(property, {
    numRuns,
    seed,
    ...(spec.path !== undefined ? { path: spec.path } : {}),
  });

  if (details.failed) {
    const counterexample = materializeCommands(details.counterexample as Iterable<Cmd>);
    const executionRecord = buildStateMachineRecord(
      spec.setup,
      counterexample,
      details.seed,
      details.counterexamplePath,
      counterexample,
    );
    const error =
      details.errorInstance instanceof Error
        ? details.errorInstance.message
        : String(details.errorInstance);
    return {
      id: spec.id,
      invariantClass: spec.invariantClass,
      holds: false,
      transitionMode: spec.transitionMode,
      executionRecord,
      error,
    };
  }

  const sampleCommands = materializeCommands(
    fc.sample(commandArb, { numRuns: 1, seed: details.seed })[0] as Iterable<Cmd>,
  );

  let rejectedCommandCount: number | undefined;
  if (spec.transitionMode === "illegal") {
    rejectedCommandCount =
      sampleCommands.length === 0
        ? 0
        : runIllegalCommands(spec.setup, sampleCommands, expectRejection);
  }

  const successResult: StateMachineInvariantResult = {
    id: spec.id,
    invariantClass: spec.invariantClass,
    holds: true,
    transitionMode: spec.transitionMode,
    executionRecord: buildStateMachineRecord(
      spec.setup,
      sampleCommands,
      details.seed,
      null,
      null,
    ),
  };
  if (rejectedCommandCount !== undefined) {
    successResult.rejectedCommandCount = rejectedCommandCount;
  }
  return successResult;
}