import {
  flattenStaticPipeline,
  flattenStaticSubsteps,
  type PlayStaticPipeline,
  type PlayStaticSubstep,
} from '../plays/static-pipeline';

export const EXECUTION_PLAN_DEFAULTS = {
  inlineRowsLimit: 1_000,
  largeMapChunkSize: 5_000,
  complexMapChunkSize: 1_000,
  workflowSoftStepBudget: 20_000,
  workflowHardStepBudget: 25_000,
  ingestStepCount: 1,
  finalizationStepCount: 2,
} as const;

export type DatasetHandleReference = {
  kind: 'dataset_handle';
  datasetId: string;
  datasetKind: 'csv' | 'map';
  tableNamespace: string;
  count: number;
  previewRows: Record<string, unknown>[];
  backing: 'inline' | 'neon_sheet' | 'r2_file';
};

export type ExecutionPlanMap = {
  mapName: string;
  tableNamespace: string;
  outputFields: string[];
  stepFields: string[];
  externalStepFields: string[];
  waterfallStages: Array<{
    waterfallId: string;
    stageIds: string[];
  }>;
  defaultChunkSize: number;
  stepsPerChunk: number;
};

export type ExecutionPlanChunkPlan = {
  inlineRowsLimit: number;
  defaultLargeMapChunkSize: number;
  softWorkflowStepBudget: number;
  hardWorkflowStepBudget: number;
  estimatedWorkflowSteps: number | null;
};

export type ExecutionPlan = {
  graphHash: string;
  artifactHash: string;
  entrypoint: string;
  assets: Array<{ playPath: string; storageKey: string }>;
  maps: ExecutionPlanMap[];
  toolDeclarations: Array<{ toolId: string; field?: string | null }>;
  chunkPlan: ExecutionPlanChunkPlan;
  persistencePlan: {
    datasetBacking: 'neon_sheet';
    compactWorkflowState: true;
  };
};

export type BuildExecutionPlanInput = {
  graphHash: string | null | undefined;
  artifactHash: string | null | undefined;
  entrypoint?: string | null;
  assets?: Array<{ playPath?: string | null; storageKey: string }>;
  staticPipeline?: PlayStaticPipeline | null;
  totalRows?: number | null;
};

export function buildExecutionPlan(
  input: BuildExecutionPlanInput,
): ExecutionPlan {
  const maps = extractPlanMaps(input.staticPipeline ?? null);
  const toolDeclarations = extractToolDeclarations(
    input.staticPipeline ?? null,
  );
  const estimatedWorkflowSteps =
    typeof input.totalRows === 'number' && Number.isFinite(input.totalRows)
      ? estimateWorkflowSteps({
          totalRows: Math.max(0, Math.floor(input.totalRows)),
          maps,
        })
      : null;

  return {
    graphHash: input.graphHash?.trim() ?? '',
    artifactHash: input.artifactHash?.trim() ?? '',
    entrypoint: input.entrypoint?.trim() || 'default',
    assets: (input.assets ?? [])
      .map((asset) => ({
        playPath: String(asset.playPath ?? '').replace(/^\.\//, ''),
        storageKey: asset.storageKey,
      }))
      .filter(
        (asset) => asset.playPath.length > 0 && asset.storageKey.length > 0,
      ),
    maps,
    toolDeclarations,
    chunkPlan: {
      inlineRowsLimit: EXECUTION_PLAN_DEFAULTS.inlineRowsLimit,
      defaultLargeMapChunkSize: EXECUTION_PLAN_DEFAULTS.largeMapChunkSize,
      softWorkflowStepBudget: EXECUTION_PLAN_DEFAULTS.workflowSoftStepBudget,
      hardWorkflowStepBudget: EXECUTION_PLAN_DEFAULTS.workflowHardStepBudget,
      estimatedWorkflowSteps,
    },
    persistencePlan: {
      datasetBacking: 'neon_sheet',
      compactWorkflowState: true,
    },
  };
}

export function chooseMapChunkSize(input: {
  /**
   * Known input row count. Pass null for streaming datasets where counting
   * would require an extra full scan; unknown counts use the preferred chunk.
   */
  totalRows: number | null;
  mapCount: number;
  stepsPerChunk: number;
  preferredChunkSize?: number | null;
  softWorkflowStepBudget?: number | null;
}): number {
  const totalRows =
    typeof input.totalRows === 'number' && Number.isFinite(input.totalRows)
      ? Math.max(0, Math.floor(input.totalRows))
      : null;
  const preferred = Math.max(
    1,
    Math.floor(
      input.preferredChunkSize ?? EXECUTION_PLAN_DEFAULTS.largeMapChunkSize,
    ),
  );
  if (totalRows === null) {
    return preferred;
  }
  if (totalRows <= EXECUTION_PLAN_DEFAULTS.inlineRowsLimit) {
    return Math.max(1, totalRows || 1);
  }

  const mapCount = Math.max(1, Math.floor(input.mapCount));
  const stepsPerChunk = Math.max(1, Math.floor(input.stepsPerChunk));
  const softBudget =
    input.softWorkflowStepBudget ??
    EXECUTION_PLAN_DEFAULTS.workflowSoftStepBudget;
  const nonChunkSteps =
    EXECUTION_PLAN_DEFAULTS.ingestStepCount +
    EXECUTION_PLAN_DEFAULTS.finalizationStepCount;
  const maxChunksAcrossMaps = Math.max(
    mapCount,
    Math.floor((softBudget - nonChunkSteps) / stepsPerChunk),
  );
  const maxChunksPerMap = Math.max(
    1,
    Math.floor(maxChunksAcrossMaps / mapCount),
  );
  const minimumSurvivalChunkSize = Math.max(
    1,
    Math.ceil(totalRows / maxChunksPerMap),
  );
  return Math.max(preferred, minimumSurvivalChunkSize);
}

export function deterministicMapChunkStepName(input: {
  mapName: string;
  chunkIndex: number;
  phase?: 'prepare' | 'execute' | 'persist' | string;
}): string {
  const phase = input.phase?.trim() || 'execute';
  return `map:${input.mapName}:chunk:${String(input.chunkIndex).padStart(4, '0')}:${phase}`;
}

function extractPlanMaps(
  pipeline: PlayStaticPipeline | null,
): ExecutionPlanMap[] {
  if (!pipeline) return [];
  const substeps = flattenStaticPipeline(pipeline);
  const fallbackWaterfalls = substeps.filter(
    (substep): substep is Extract<PlayStaticSubstep, { type: 'waterfall' }> =>
      substep.type === 'waterfall',
  );
  const fallbackStepSuites = substeps.filter(
    (substep): substep is Extract<PlayStaticSubstep, { type: 'step_suite' }> =>
      substep.type === 'step_suite',
  );
  const datasetSubsteps = substeps.filter(
    (substep): substep is Extract<PlayStaticSubstep, { type: 'dataset' }> =>
      substep.type === 'dataset',
  );
  const hasSiblingMaps = datasetSubsteps.length > 1;
  return datasetSubsteps.map((mapSubstep) => {
    const waterfalls = fallbackWaterfalls.filter((waterfall) => {
      if (!mapSubstep.waterfallIds?.length) {
        return (
          !hasSiblingMaps ||
          substepFieldBelongsToMap(waterfall.field, mapSubstep)
        );
      }
      return (
        (waterfall.id && mapSubstep.waterfallIds.includes(waterfall.id)) ||
        mapSubstep.waterfallIds.includes(waterfall.field)
      );
    });
    const waterfallStages = waterfalls.map((waterfall) => ({
      waterfallId: waterfall.id ?? waterfall.field,
      stageIds: waterfall.steps?.map((step) => step.id) ?? [],
    }));
    const stepSuites = fallbackStepSuites.filter((stepSuite) => {
      if (!mapSubstep.waterfallIds?.length) {
        return (
          !hasSiblingMaps ||
          substepFieldBelongsToMap(stepSuite.field, mapSubstep)
        );
      }
      return mapSubstep.waterfallIds.includes(stepSuite.field);
    });
    const looseMapSubsteps = substeps.filter((substep) => {
      if (!isLooseMapExecutionSubstep(substep)) {
        return false;
      }
      return substepFieldBelongsToMap(substep.field, mapSubstep);
    });
    const stepSuiteStepsPerChunk = stepSuites.reduce(
      (max, stepSuite) => Math.max(max, stepSuite.steps.length),
      0,
    );
    const waterfallStepsPerChunk = waterfallStages.reduce(
      (max, waterfall) => Math.max(max, waterfall.stageIds.length),
      0,
    );
    const directMapStepsPerChunk = maxExecutionSubstepsPerChunk(
      mapSubstep.steps ?? [],
    );
    const stepsPerChunk = Math.max(
      1,
      directMapStepsPerChunk,
      looseMapSubsteps.length,
      waterfallStepsPerChunk,
      stepSuiteStepsPerChunk,
    );
    const mapExecutionSubsteps = [
      ...(mapSubstep.steps ?? []),
      ...looseMapSubsteps,
      ...stepSuites,
      ...waterfalls,
    ];
    return {
      mapName: mapSubstep.name ?? mapSubstep.field,
      tableNamespace: mapSubstep.tableNamespace ?? mapSubstep.field,
      outputFields: mapSubstep.outputFields ?? [],
      stepFields: collectPlanMapStepFields(mapExecutionSubsteps),
      externalStepFields:
        collectPlanMapExternalStepFields(mapExecutionSubsteps),
      waterfallStages,
      defaultChunkSize:
        stepsPerChunk > 1
          ? EXECUTION_PLAN_DEFAULTS.complexMapChunkSize
          : EXECUTION_PLAN_DEFAULTS.largeMapChunkSize,
      stepsPerChunk,
    };
  });
}

function maxExecutionSubstepsPerChunk(substeps: PlayStaticSubstep[]): number {
  return substeps.reduce((max, substep) => {
    if (substep.type === 'step_suite') {
      return Math.max(max, substep.steps.length);
    }
    if (substep.type === 'waterfall') {
      return Math.max(max, substep.steps?.length ?? 0);
    }
    if (substep.type === 'control_flow') {
      return Math.max(max, maxExecutionSubstepsPerChunk(substep.steps));
    }
    return Math.max(max, 1);
  }, 0);
}

function isLooseMapExecutionSubstep(
  substep: PlayStaticSubstep,
): substep is Extract<
  PlayStaticSubstep,
  { type: 'tool' | 'play_call' | 'control_flow' | 'code' }
> {
  return (
    substep.type === 'tool' ||
    substep.type === 'play_call' ||
    substep.type === 'control_flow' ||
    substep.type === 'code'
  );
}

function substepFieldBelongsToMap(
  substepField: string,
  mapSubstep: Extract<PlayStaticSubstep, { type: 'dataset' }>,
): boolean {
  const field = substepField.trim();
  if (!field) {
    return false;
  }
  const mapFields = [
    mapSubstep.field,
    mapSubstep.name,
    mapSubstep.tableNamespace,
    ...(mapSubstep.outputFields ?? []),
  ].filter((candidate): candidate is string => Boolean(candidate?.trim()));
  return mapFields.some((mapField) => {
    const normalized = mapField.trim();
    return field === normalized || field.startsWith(`${normalized}.`);
  });
}

function collectPlanMapStepFields(substeps: PlayStaticSubstep[]): string[] {
  const fields = new Set<string>();
  for (const substep of flattenStaticSubsteps(substeps)) {
    if ('field' in substep && typeof substep.field === 'string') {
      const field = substep.field.trim();
      if (field) {
        fields.add(field);
      }
    }
  }
  return [...fields];
}

function collectPlanMapExternalStepFields(
  substeps: PlayStaticSubstep[],
): string[] {
  const fields = new Set<string>();
  for (const substep of flattenStaticSubsteps(substeps)) {
    if (
      substep.type !== 'tool' &&
      substep.type !== 'play_call' &&
      substep.type !== 'waterfall'
    ) {
      continue;
    }
    if ('field' in substep && typeof substep.field === 'string') {
      const field = substep.field.trim();
      if (field) {
        fields.add(field);
      }
    }
  }
  return [...fields];
}

function extractToolDeclarations(
  pipeline: PlayStaticPipeline | null,
): ExecutionPlan['toolDeclarations'] {
  if (!pipeline) return [];
  const seen = new Set<string>();
  const declarations: ExecutionPlan['toolDeclarations'] = [];
  for (const substep of flattenStaticPipeline(pipeline)) {
    if (substep.type === 'tool') {
      const key = `${substep.toolId}:${substep.field}`;
      if (!seen.has(key)) {
        seen.add(key);
        declarations.push({ toolId: substep.toolId, field: substep.field });
      }
      continue;
    }
    if (substep.type === 'waterfall') {
      for (const step of substep.steps ?? []) {
        if (!step.toolId) continue;
        const key = `${step.toolId}:${substep.field}`;
        if (!seen.has(key)) {
          seen.add(key);
          declarations.push({ toolId: step.toolId, field: substep.field });
        }
      }
    }
  }
  return declarations;
}

function estimateWorkflowSteps(input: {
  totalRows: number;
  maps: ExecutionPlanMap[];
}): number {
  const maps = input.maps.length > 0 ? input.maps : [];
  const chunkSteps = maps.reduce((sum, map) => {
    const chunkSize = chooseMapChunkSize({
      totalRows: input.totalRows,
      mapCount: maps.length,
      stepsPerChunk: map.stepsPerChunk,
      preferredChunkSize: map.defaultChunkSize,
    });
    return sum + Math.ceil(input.totalRows / chunkSize) * map.stepsPerChunk;
  }, 0);
  return (
    EXECUTION_PLAN_DEFAULTS.ingestStepCount +
    chunkSteps +
    EXECUTION_PLAN_DEFAULTS.finalizationStepCount
  );
}