import {
  CONCURRENCY_SCHEDULE_KINDS,
  type ConcurrencySchedule,
  type ConcurrencyScheduleResult,
  type InstrumentedSystem,
  type ScheduleStepRef,
  type StepContext,
  type TraceEntry,
} from "./types.js";

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

function assertScheduleKind(value: string): void {
  if (!(CONCURRENCY_SCHEDULE_KINDS as readonly string[]).includes(value)) {
    throw new Error(`unknown schedule kind: ${value}`);
  }
}

function safeExportState(state: unknown): unknown {
  try {
    if (typeof structuredClone === "function") {
      return structuredClone(state);
    }
  } catch {
    // Fall through to manual export for function-bearing fixture state.
  }

  if (state === null || typeof state !== "object") {
    return state;
  }

  const exported: Record<string, unknown> = {};
  for (const [key, value] of Object.entries(state as Record<string, unknown>)) {
    if (typeof value === "function") {
      continue;
    }
    if (value instanceof Set) {
      exported[key] = [...value];
      continue;
    }
    if (value instanceof Map) {
      exported[key] = Object.fromEntries(value.entries());
      continue;
    }
    if (Array.isArray(value)) {
      exported[key] = value.map((entry: unknown) => {
        if (entry !== null && typeof entry === "object" && !Array.isArray(entry)) {
          try {
            return JSON.parse(JSON.stringify(entry)) as unknown;
          } catch {
            return entry;
          }
        }
        return entry;
      });
      continue;
    }
    if (value !== null && typeof value === "object") {
      try {
        exported[key] = JSON.parse(JSON.stringify(value)) as unknown;
      } catch {
        exported[key] = "[non-serializable]";
      }
      continue;
    }
    exported[key] = value;
  }
  return exported;
}

function isScheduleStepRef(value: unknown): value is ScheduleStepRef {
  if (typeof value !== "object" || value === null) {
    return false;
  }
  const candidate = value as { taskId?: unknown; stepId?: unknown };
  return typeof candidate.taskId === "string" && typeof candidate.stepId === "string";
}

function validateSchedule(schedule: ConcurrencySchedule): void {
  assertNonEmptyString(schedule.id, "id");
  assertScheduleKind(schedule.kind);
  const interleaving: readonly ScheduleStepRef[] = schedule.interleaving;
  if (!Array.isArray(interleaving) || interleaving.length === 0) {
    throw new Error("interleaving must contain at least one scheduled step");
  }
  interleaving.forEach((scheduledStep, index) => {
    if (!isScheduleStepRef(scheduledStep)) {
      throw new Error(`interleaving[${String(index)}] must include taskId and stepId`);
    }
    if (scheduledStep.taskId.trim().length === 0) {
      throw new Error(`interleaving[${String(index)}].taskId must be a non-empty string`);
    }
    if (scheduledStep.stepId.trim().length === 0) {
      throw new Error(`interleaving[${String(index)}].stepId must be a non-empty string`);
    }
  });
}

function validateSystem<TState>(system: InstrumentedSystem<TState>): void {
  assertNonEmptyString(system.id, "system.id");
  if (typeof system.createState !== "function") {
    throw new Error("createState must be a function");
  }
  if (typeof system.assertOutcome !== "function") {
    throw new Error("assertOutcome must be a function");
  }
  if (typeof system.tasks !== "object") {
    throw new Error("tasks must be an object");
  }
}

function validateScheduledStep<TState>(
  system: InstrumentedSystem<TState>,
  taskId: string,
  stepId: string,
): void {
  const task = system.tasks[taskId];
  if (task === undefined) {
    throw new Error(`unknown task: ${taskId}`);
  }
  if (!task.steps.includes(stepId)) {
    throw new Error(`unknown step ${stepId} on task ${taskId}`);
  }
  if (typeof task.runStep !== "function") {
    throw new Error(`task ${taskId} is missing runStep`);
  }
}

function createStepContext<TState>(
  state: TState,
  yieldPoints: string[],
): StepContext<TState> {
  return {
    state,
    yield: (point: string) => {
      assertNonEmptyString(point, "yield point");
      yieldPoints.push(point);
      return Promise.resolve();
    },
  };
}

export async function runSchedule<TState>(
  schedule: ConcurrencySchedule,
  system: InstrumentedSystem<TState>,
): Promise<ConcurrencyScheduleResult> {
  validateSchedule(schedule);
  validateSystem(system);

  for (const step of schedule.interleaving) {
    validateScheduledStep(system, step.taskId, step.stepId);
  }

  const state = await system.createState();
  const trace: TraceEntry[] = [];

  try {
    for (const [index, step] of schedule.interleaving.entries()) {
      const task = system.tasks[step.taskId];
      if (task === undefined) {
        throw new Error(`unknown task: ${step.taskId}`);
      }

      const yieldPoints: string[] = [];
      const context = createStepContext(state, yieldPoints);
      await task.runStep(step.stepId, context);
      trace.push({
        index,
        taskId: step.taskId,
        stepId: step.stepId,
        yieldPoints: [...yieldPoints],
      });
    }

    system.assertOutcome(state);

    return {
      id: schedule.id,
      kind: schedule.kind,
      holds: true,
      trace,
      finalState: safeExportState(state),
    };
  } catch (error) {
    const message = error instanceof Error ? error.message : String(error);
    return {
      id: schedule.id,
      kind: schedule.kind,
      holds: false,
      trace,
      finalState: safeExportState(state),
      error: message,
    };
  }
}