import {
  ArithmeticOperator,
  BaseExpressionFields,
  BooleanExpression,
  BooleanValueType,
  ComparisonOperator,
  ContinuousDimensionType,
  DiscreteDimensionType,
  EnumExpression,
  Expression,
  FloatExpression,
  FloatValueType,
  FunctionValueType,
  IntExpression,
  IntValueType,
  ListValueType,
  NoOpExpression,
  ObjectValueType,
  Parameter,
  RegexExpression,
  StringExpression,
  StringValueType,
  ValueType,
  VariableExpression,
  VoidValueType,
} from "../types";

type FoldPartialResult<TFoldResult> =
  | null
  | NoOpExpression
  | BooleanExpression
  | IntExpression
  | FloatExpression
  | StringExpression
  | RegexExpression
  | EnumExpression
  | (BaseExpressionFields & {
      type: "ObjectExpression";
      valueType: ObjectValueType;
      objectTypeName: string;
      fields: { [fieldName: string]: TFoldResult };
    })
  | (BaseExpressionFields & {
      type: "GetFieldExpression";
      valueType: ValueType;
      object: TFoldResult;
      fieldPath: string | null;
    })
  | (BaseExpressionFields & {
      type: "UpdateObjectExpression";
      valueType: ObjectValueType;
      object: TFoldResult;
      updates: { [fieldName: string]: TFoldResult };
    })
  | (BaseExpressionFields & {
      type: "ListExpression";
      valueType: ListValueType;
      items: TFoldResult[];
    })
  | (BaseExpressionFields & {
      type: "SwitchExpression";
      valueType: ValueType;
      control: TFoldResult;
      cases: {
        id: string;
        when: TFoldResult;
        then: TFoldResult;
      }[];
      default: TFoldResult;
    })
  | (BaseExpressionFields & {
      type: "EnumSwitchExpression";
      valueType: ValueType;
      control: TFoldResult;
      cases: { [enumValue: string]: TFoldResult };
    })
  | (BaseExpressionFields & {
      type: "ComparisonExpression";
      valueType: BooleanValueType;
      operator: ComparisonOperator | null;
      a: TFoldResult;
      b: TFoldResult;
    })
  | (BaseExpressionFields & {
      type: "ArithmeticExpression";
      valueType: IntValueType | FloatValueType;
      operator: ArithmeticOperator | null;
      a: TFoldResult;
      b: TFoldResult;
    })
  | (BaseExpressionFields & {
      type: "RoundNumberExpression";
      valueType: IntValueType;
      number: TFoldResult;
    })
  | (BaseExpressionFields & {
      type: "StringifyNumberExpression";
      valueType: StringValueType;
      number: TFoldResult;
    })
  | (BaseExpressionFields & {
      type: "StringConcatExpression";
      valueType: StringValueType;
      strings: TFoldResult;
    })
  | (BaseExpressionFields & {
      type: "GetUrlQueryParameterExpression";
      valueType: StringValueType;
      url: TFoldResult;
      queryParameterName: TFoldResult;
    })
  | (BaseExpressionFields & {
      type: "SplitExpression";
      valueType: ValueType;
      splitId: string | null;
      dimensionId: string | null;
      expose: TFoldResult;
      unitId: TFoldResult;
      dimensionMapping:
        | {
            type: typeof DiscreteDimensionType;
            cases: { [armId: string]: TFoldResult };
          }
        | {
            type: typeof ContinuousDimensionType;
            function: TFoldResult;
          };
      eventObjectTypeName: string | null;
      eventPayload: TFoldResult;
    })
  | (BaseExpressionFields & {
      type: "LogEventExpression";
      valueType: VoidValueType;
      eventObjectTypeName: string | null;
      eventPayload: TFoldResult;
      eventTypeId: string | null;
      unitId: TFoldResult;
    })
  | (BaseExpressionFields & {
      type: "FunctionExpression";
      valueType: FunctionValueType;
      parameters: Parameter[];
      body: TFoldResult;
    })
  | VariableExpression
  | (BaseExpressionFields & {
      type: "ApplicationExpression";
      valueType: ValueType;
      function: TFoldResult;
      arguments: TFoldResult[];
    });

type FoldFunction<TFoldResult> = (
  partialResult: FoldPartialResult<TFoldResult>
) => TFoldResult;

// Expensive
export default function fold<TFoldResult>(
  f: FoldFunction<TFoldResult>,
  expression: Expression | null
): TFoldResult {
  if (!expression) {
    return f(expression);
  }

  switch (expression.type) {
    case "NoOpExpression":
    case "BooleanExpression":
    case "IntExpression":
    case "FloatExpression":
    case "StringExpression":
    case "RegexExpression":
    case "EnumExpression":
      return f(expression);

    case "ObjectExpression":
      return f({
        ...expression,
        fields: Object.fromEntries(
          Object.entries(expression.fields).map(([fieldName, field]) => [
            fieldName,
            fold(f, field),
          ])
        ),
      });

    case "GetFieldExpression":
      return f({
        ...expression,
        object: fold(f, expression.object),
      });

    case "UpdateObjectExpression":
      return f({
        ...expression,
        object: fold(f, expression.object),
        updates: Object.fromEntries(
          Object.entries(expression.updates).map(([fieldName, field]) => [
            fieldName,
            fold(f, field),
          ])
        ),
      });

    case "ListExpression":
      return f({
        ...expression,
        items: expression.items.map((item) => fold(f, item)),
      });

    case "SwitchExpression":
      return f({
        ...expression,
        control: fold(f, expression.control),
        cases: expression.cases.map((item) => ({
          id: item.id,
          when: fold(f, item.when),
          then: fold(f, item.then),
        })),
        default: fold(f, expression.default),
      });

    case "EnumSwitchExpression":
      return f({
        ...expression,
        control: fold(f, expression.control),
        cases: Object.fromEntries(
          Object.entries(expression.cases).map(
            ([enumValue, caseExpression]) => [
              enumValue,
              fold(f, caseExpression),
            ]
          )
        ),
      });

    case "ComparisonExpression":
    case "ArithmeticExpression":
      return f({
        ...expression,
        a: fold(f, expression.a),
        b: fold(f, expression.b),
      });

    case "RoundNumberExpression":
    case "StringifyNumberExpression":
      return f({
        ...expression,
        number: fold(f, expression.number),
      });

    case "StringConcatExpression":
      return f({
        ...expression,
        strings: fold(f, expression.strings),
      });

    case "GetUrlQueryParameterExpression":
      return f({
        ...expression,
        url: fold(f, expression.url),
        queryParameterName: fold(f, expression.queryParameterName),
      });

    case "SplitExpression":
      return f({
        ...expression,
        expose: fold(f, expression.expose),
        unitId: fold(f, expression.unitId),
        dimensionMapping:
          expression.dimensionMapping.type === "discrete"
            ? {
                type: "discrete",
                cases: Object.fromEntries(
                  Object.entries(expression.dimensionMapping.cases).map(
                    ([armId, caseExpression]) => [
                      armId,
                      fold(f, caseExpression),
                    ]
                  )
                ),
              }
            : {
                type: "continuous",
                function: fold(f, expression.dimensionMapping.function),
              },
        eventObjectTypeName: expression.eventObjectTypeName,
        eventPayload: fold(f, expression.eventPayload),
      });

    case "LogEventExpression":
      return f({
        ...expression,
        unitId: fold(f, expression.unitId),
        eventPayload: fold(f, expression.eventPayload),
      });

    case "FunctionExpression":
      return f({
        ...expression,
        body: fold(f, expression.body),
      });

    case "VariableExpression":
      return f(expression);

    case "ApplicationExpression":
      return f({
        ...expression,
        function: fold(f, expression.function),
        arguments: expression.arguments.map((argument) => fold(f, argument)),
      });

    default: {
      const neverExpression: never = expression;
      throw new Error(
        `Unexpected expression: ${JSON.stringify(neverExpression)}`
      );
    }
  }
}

// Applies f to children, constructs new expression, applies f to it and returns
// it along with merged map results
// Expensive
export function mapExpressionWithResult<TMapResult>(
  fn: (expr: Expression | null) => {
    newExpression: Expression | null;
    mapResult: TMapResult;
  },
  combineResults: (...results: TMapResult[]) => TMapResult,
  expression: Expression | null
): {
  newExpression: Expression | null;
  mapResult: TMapResult;
} {
  type TFoldResult = {
    newExpression: Expression | null;
    mapResult: TMapResult;
  };
  // eslint-disable-next-line func-style
  const foldFunction: FoldFunction<TFoldResult> = (partialResult) => {
    if (!partialResult) {
      return fn(partialResult);
    }

    switch (partialResult.type) {
      case "NoOpExpression":
      case "BooleanExpression":
      case "IntExpression":
      case "FloatExpression":
      case "StringExpression":
      case "RegexExpression":
      case "EnumExpression":
        return fn(partialResult);

      case "ObjectExpression": {
        const thisResult = fn({
          ...partialResult,
          fields: Object.fromEntries(
            Object.entries(partialResult.fields).map(
              ([fieldName, fieldResult]) => [
                fieldName,
                fieldResult.newExpression,
              ]
            )
          ),
        });
        return {
          newExpression: thisResult.newExpression,
          mapResult: combineResults(
            thisResult.mapResult,
            ...Object.values(partialResult.fields).map(
              (fieldResult) => fieldResult.mapResult
            )
          ),
        };
      }

      case "GetFieldExpression": {
        const thisResult = fn({
          ...partialResult,
          object: partialResult.object.newExpression,
        });
        return {
          newExpression: thisResult.newExpression,
          mapResult: combineResults(
            thisResult.mapResult,
            partialResult.object.mapResult
          ),
        };
      }

      case "UpdateObjectExpression": {
        const thisResult = fn({
          ...partialResult,
          object: partialResult.object.newExpression,
          updates: Object.fromEntries(
            Object.entries(partialResult.updates).map(
              ([fieldName, updateResult]) => [
                fieldName,
                updateResult.newExpression,
              ]
            )
          ),
        });
        return {
          newExpression: thisResult.newExpression,
          mapResult: combineResults(
            thisResult.mapResult,
            partialResult.object.mapResult,
            ...Object.values(partialResult.updates).map(
              (updateResult) => updateResult.mapResult
            )
          ),
        };
      }

      case "ListExpression": {
        const thisResult = fn({
          ...partialResult,
          items: partialResult.items.map(
            (itemResult) => itemResult.newExpression
          ),
        });
        return {
          newExpression: thisResult.newExpression,
          mapResult: combineResults(
            thisResult.mapResult,
            ...partialResult.items.map((itemResult) => itemResult.mapResult)
          ),
        };
      }

      case "SwitchExpression": {
        const thisResult = fn({
          ...partialResult,
          control: partialResult.control.newExpression,
          cases: partialResult.cases.map((caseResult) => ({
            id: caseResult.id,
            when: caseResult.when.newExpression,
            then: caseResult.then.newExpression,
          })),
          default: partialResult.default.newExpression,
        });
        return {
          newExpression: thisResult.newExpression,
          mapResult: combineResults(
            thisResult.mapResult,
            partialResult.control.mapResult,
            ...partialResult.cases.map((caseResult) =>
              combineResults(
                caseResult.when.mapResult,
                caseResult.then.mapResult
              )
            ),
            partialResult.default.mapResult
          ),
        };
      }

      case "EnumSwitchExpression": {
        const thisResult = fn({
          ...partialResult,
          control: partialResult.control.newExpression,
          cases: Object.fromEntries(
            Object.entries(partialResult.cases).map(
              ([enumValue, caseResult]) => [enumValue, caseResult.newExpression]
            )
          ),
        });
        return {
          newExpression: thisResult.newExpression,
          mapResult: combineResults(
            thisResult.mapResult,
            partialResult.control.mapResult,
            ...Object.values(partialResult.cases).map(
              (caseResult) => caseResult.mapResult
            )
          ),
        };
      }

      case "ComparisonExpression":
      case "ArithmeticExpression": {
        const thisResult = fn({
          ...partialResult,
          a: partialResult.a.newExpression,
          b: partialResult.b.newExpression,
        });
        return {
          newExpression: thisResult.newExpression,
          mapResult: combineResults(
            thisResult.mapResult,
            partialResult.a.mapResult,
            partialResult.b.mapResult
          ),
        };
      }

      case "RoundNumberExpression":
      case "StringifyNumberExpression": {
        const thisResult = fn({
          ...partialResult,
          number: partialResult.number.newExpression,
        });
        return {
          newExpression: thisResult.newExpression,
          mapResult: combineResults(
            thisResult.mapResult,
            partialResult.number.mapResult
          ),
        };
      }

      case "StringConcatExpression": {
        const thisResult = fn({
          ...partialResult,
          strings: partialResult.strings.newExpression,
        });
        return {
          newExpression: thisResult.newExpression,
          mapResult: combineResults(
            thisResult.mapResult,
            partialResult.strings.mapResult
          ),
        };
      }

      case "GetUrlQueryParameterExpression": {
        const thisResult = fn({
          ...partialResult,
          url: partialResult.url.newExpression,
          queryParameterName: partialResult.queryParameterName.newExpression,
        });
        return {
          newExpression: thisResult.newExpression,
          mapResult: combineResults(
            thisResult.mapResult,
            partialResult.url.mapResult,
            partialResult.queryParameterName.mapResult
          ),
        };
      }

      case "SplitExpression": {
        const thisResult = fn({
          ...partialResult,
          expose: partialResult.expose.newExpression,
          unitId: partialResult.unitId.newExpression,
          dimensionMapping:
            partialResult.dimensionMapping.type === "discrete"
              ? {
                  type: "discrete",
                  cases: Object.fromEntries(
                    Object.entries(partialResult.dimensionMapping.cases).map(
                      ([armId, caseResult]) => [armId, caseResult.newExpression]
                    )
                  ),
                }
              : {
                  type: "continuous",
                  function:
                    partialResult.dimensionMapping.function.newExpression,
                },
          eventObjectTypeName: partialResult.eventObjectTypeName,
          eventPayload: partialResult.eventPayload.newExpression,
          featuresMapping: {},
        });
        return {
          newExpression: thisResult.newExpression,
          mapResult: combineResults(
            thisResult.mapResult,
            partialResult.expose.mapResult,
            partialResult.unitId.mapResult,
            partialResult.eventPayload.mapResult,
            ...(partialResult.dimensionMapping.type === "discrete"
              ? Object.values(partialResult.dimensionMapping.cases).map(
                  (caseMapResult) => caseMapResult.mapResult
                )
              : [partialResult.dimensionMapping.function.mapResult])
          ),
        };
      }

      case "LogEventExpression": {
        const thisResult = fn({
          ...partialResult,
          eventObjectTypeName: partialResult.eventObjectTypeName,
          eventPayload: partialResult.eventPayload.newExpression,
          unitId: partialResult.unitId.newExpression,
          featuresMapping: {},
        });
        return {
          newExpression: thisResult.newExpression,
          mapResult: combineResults(
            thisResult.mapResult,
            partialResult.unitId.mapResult,
            partialResult.eventPayload.mapResult
          ),
        };
      }

      case "FunctionExpression": {
        const thisResult = fn({
          ...partialResult,
          body: partialResult.body.newExpression,
        });
        return {
          newExpression: thisResult.newExpression,
          mapResult: combineResults(
            thisResult.mapResult,
            partialResult.body.mapResult
          ),
        };
      }

      case "VariableExpression":
        return fn(partialResult);

      case "ApplicationExpression": {
        const thisResult = fn({
          ...partialResult,
          function: partialResult.function.newExpression,
          arguments: partialResult.arguments.map(
            (argumentResult) => argumentResult.newExpression
          ),
        });
        return {
          newExpression: thisResult.newExpression,
          mapResult: combineResults(
            thisResult.mapResult,
            partialResult.function.mapResult,
            ...partialResult.arguments.map(
              (argumentResult) => argumentResult.mapResult
            )
          ),
        };
      }

      default: {
        const neverPartialResult: never = partialResult;
        throw new Error(
          `Unexpected partial result: ${JSON.stringify(neverPartialResult)}`
        );
      }
    }
  };
  return fold(foldFunction, expression);
}

// Expensive
export function mapExpression(
  mapper: (expr: Expression | null) => Expression | null,
  expression: Expression | null
): Expression | null {
  const result = mapExpressionWithResult<null>(
    (expr) => ({
      newExpression: mapper(expr),
      mapResult: null,
    }),
    () => null,
    expression
  );
  return result.newExpression;
}