import type {
  PlaySheetContract,
  PlayStaticPipeline,
  PlayStaticSubstep,
} from './static-pipeline';
import { compileSheetContract } from './static-pipeline';

export type ResolvablePlayStaticShape = {
  staticPipeline?: PlayStaticPipeline | null;
};

export type ResolveStaticPipelineInput = {
  rootPlayId: string;
  root: ResolvablePlayStaticShape;
  resolveReferencedPlay: (
    playId: string,
  ) => Promise<ResolvablePlayStaticShape | null>;
  maxDepth?: number;
};

function clonePipeline(pipeline: PlayStaticPipeline): PlayStaticPipeline {
  return {
    ...pipeline,
    inputFields: pipeline.inputFields ? [...pipeline.inputFields] : undefined,
    rowKeyFields: pipeline.rowKeyFields
      ? [...pipeline.rowKeyFields]
      : undefined,
    fields: [...pipeline.fields],
    stages: (pipeline.stages ?? []).map((stage) => ({ ...stage })),
    substeps: pipeline.substeps.map((substep) => ({ ...substep })),
    sheetContract: cloneSheetContract(pipeline.sheetContract),
    sheetContractErrors: undefined,
  };
}

function cloneSheetContract(contract: PlaySheetContract | null | undefined) {
  if (!contract) return contract;
  return {
    ...contract,
    columns: contract.columns.map((column) => ({ ...column })),
  };
}

function materializePipeline(
  shape: ResolvablePlayStaticShape | null,
): PlayStaticPipeline | null {
  if (!shape) return null;
  if (shape.staticPipeline) {
    return clonePipeline(shape.staticPipeline);
  }
  return null;
}

function annotateSubstep(
  substep: PlayStaticSubstep,
  input: { callDepth: number; callPath: string[] },
): PlayStaticSubstep {
  return {
    ...substep,
    callDepth: input.callDepth,
    callPath: input.callPath,
  };
}

async function resolvePipelineSubsteps(input: {
  steps: PlayStaticSubstep[];
  resolveReferencedPlay: (
    playId: string,
  ) => Promise<ResolvablePlayStaticShape | null>;
  stack: string[];
  maxDepth: number;
}): Promise<PlayStaticSubstep[]> {
  const { steps, resolveReferencedPlay, stack, maxDepth } = input;
  const callDepth = Math.max(0, stack.length - 1);

  return await Promise.all(
    steps.map(async (substep) => {
      const annotated = annotateSubstep(substep, {
        callDepth,
        callPath: [...stack],
      });

      if (
        (annotated.type === 'dataset' ||
          annotated.type === 'step_suite' ||
          annotated.type === 'control_flow') &&
        annotated.steps?.length
      ) {
        return {
          ...annotated,
          steps: await resolvePipelineSubsteps({
            steps: annotated.steps,
            resolveReferencedPlay,
            stack,
            maxDepth,
          }),
        } satisfies PlayStaticSubstep;
      }

      if (annotated.type !== 'play_call') {
        return annotated;
      }

      if (stack.includes(annotated.playId)) {
        return {
          ...annotated,
          cycleDetected: true,
          resolutionError: `Static cycle detected: ${[...stack, annotated.playId].join(' -> ')}`,
          pipeline: null,
        } satisfies PlayStaticSubstep;
      }

      if (stack.length >= maxDepth) {
        return {
          ...annotated,
          resolutionError: `Static nesting limit reached at ${annotated.playId}`,
          pipeline: null,
        } satisfies PlayStaticSubstep;
      }

      const childShape = await resolveReferencedPlay(annotated.playId);
      const childPipeline = materializePipeline(childShape);
      if (!childPipeline) {
        return {
          ...annotated,
          resolutionError: `Unable to statically resolve play "${annotated.playId}"`,
          pipeline: null,
        } satisfies PlayStaticSubstep;
      }

      const resolvedChildPipeline = {
        ...childPipeline,
        stages: await resolvePipelineSubsteps({
          steps: childPipeline.stages ?? [],
          resolveReferencedPlay,
          stack: [...stack, annotated.playId],
          maxDepth,
        }),
        substeps: await resolvePipelineSubsteps({
          steps: childPipeline.substeps,
          resolveReferencedPlay,
          stack: [...stack, annotated.playId],
          maxDepth,
        }),
      };
      const childContract = compileSheetContract(resolvedChildPipeline);

      return {
        ...annotated,
        pipeline: {
          ...resolvedChildPipeline,
          sheetContract: childContract.contract,
          sheetContractErrors: childContract.errors,
        },
      } satisfies PlayStaticSubstep;
    }),
  );
}

export async function resolveStaticPipelineTree(
  input: ResolveStaticPipelineInput,
): Promise<PlayStaticPipeline | null> {
  const maxDepth = input.maxDepth ?? 6;
  const rootPipeline = materializePipeline(input.root);
  if (!rootPipeline) {
    return null;
  }

  const resolvedPipeline = {
    ...rootPipeline,
    stages: await resolvePipelineSubsteps({
      steps: rootPipeline.stages ?? [],
      resolveReferencedPlay: input.resolveReferencedPlay,
      stack: [input.rootPlayId],
      maxDepth,
    }),
    substeps: await resolvePipelineSubsteps({
      steps: rootPipeline.substeps,
      resolveReferencedPlay: input.resolveReferencedPlay,
      stack: [input.rootPlayId],
      maxDepth,
    }),
  };
  const contract = compileSheetContract(resolvedPipeline);
  return {
    ...resolvedPipeline,
    sheetContract: contract.contract,
    sheetContractErrors: contract.errors,
  };
}