/**
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

import {
  HIRFunction,
  InstructionId,
  Place,
  PrunedReactiveScopeBlock,
  ReactiveBlock,
  ReactiveFunction,
  ReactiveInstruction,
  ReactiveScopeBlock,
  ReactiveStatement,
  ReactiveTerminal,
  ReactiveTerminalStatement,
  ReactiveValue,
} from "../HIR/HIR";
import {
  eachInstructionLValue,
  eachInstructionValueOperand,
  eachTerminalOperand,
} from "../HIR/visitors";
import { assertExhaustive } from "../Utils/utils";

export function visitReactiveFunction<TState>(
  fn: ReactiveFunction,
  visitor: ReactiveFunctionVisitor<TState>,
  state: TState
): void {
  visitor.visitBlock(fn.body, state);
}

export class ReactiveFunctionVisitor<TState = void> {
  visitID(_id: InstructionId, _state: TState): void {}
  visitParam(_place: Place, _state: TState): void {}
  visitLValue(_id: InstructionId, _lvalue: Place, _state: TState): void {}
  visitPlace(_id: InstructionId, _place: Place, _state: TState): void {}
  visitReactiveFunctionValue(
    _id: InstructionId,
    _dependencies: Array<Place>,
    _fn: ReactiveFunction,
    _state: TState
  ): void {}

  visitValue(id: InstructionId, value: ReactiveValue, state: TState): void {
    this.traverseValue(id, value, state);
  }
  traverseValue(id: InstructionId, value: ReactiveValue, state: TState): void {
    switch (value.kind) {
      case "OptionalExpression": {
        this.visitValue(id, value.value, state);
        break;
      }
      case "LogicalExpression": {
        this.visitValue(id, value.left, state);
        this.visitValue(id, value.right, state);
        break;
      }
      case "ConditionalExpression": {
        this.visitValue(id, value.test, state);
        this.visitValue(id, value.consequent, state);
        this.visitValue(id, value.alternate, state);
        break;
      }
      case "SequenceExpression": {
        for (const instr of value.instructions) {
          this.visitInstruction(instr, state);
        }
        this.visitValue(value.id, value.value, state);
        break;
      }
      case "ReactiveFunctionValue": {
        this.visitReactiveFunctionValue(
          id,
          value.dependencies,
          value.fn,
          state
        );
        break;
      }
      default: {
        for (const place of eachInstructionValueOperand(value)) {
          this.visitPlace(id, place, state);
        }
      }
    }
  }

  visitInstruction(instruction: ReactiveInstruction, state: TState): void {
    this.traverseInstruction(instruction, state);
  }
  traverseInstruction(instruction: ReactiveInstruction, state: TState): void {
    this.visitID(instruction.id, state);
    for (const operand of eachInstructionLValue(instruction)) {
      this.visitLValue(instruction.id, operand, state);
    }
    this.visitValue(instruction.id, instruction.value, state);
  }

  visitTerminal(stmt: ReactiveTerminalStatement, state: TState): void {
    this.traverseTerminal(stmt, state);
  }
  traverseTerminal(stmt: ReactiveTerminalStatement, state: TState): void {
    const { terminal } = stmt;
    if (terminal.id !== null) {
      this.visitID(terminal.id, state);
    }
    switch (terminal.kind) {
      case "break":
      case "continue": {
        break;
      }
      case "return": {
        this.visitPlace(terminal.id, terminal.value, state);
        break;
      }
      case "throw": {
        this.visitPlace(terminal.id, terminal.value, state);
        break;
      }
      case "for": {
        this.visitValue(terminal.id, terminal.init, state);
        this.visitValue(terminal.id, terminal.test, state);
        this.visitBlock(terminal.loop, state);
        if (terminal.update !== null) {
          this.visitValue(terminal.id, terminal.update, state);
        }
        break;
      }
      case "for-of": {
        this.visitValue(terminal.id, terminal.init, state);
        this.visitValue(terminal.id, terminal.test, state);
        this.visitBlock(terminal.loop, state);
        break;
      }
      case "for-in": {
        this.visitValue(terminal.id, terminal.init, state);
        this.visitBlock(terminal.loop, state);
        break;
      }
      case "do-while": {
        this.visitBlock(terminal.loop, state);
        this.visitValue(terminal.id, terminal.test, state);
        break;
      }
      case "while": {
        this.visitValue(terminal.id, terminal.test, state);
        this.visitBlock(terminal.loop, state);
        break;
      }
      case "if": {
        this.visitPlace(terminal.id, terminal.test, state);
        this.visitBlock(terminal.consequent, state);
        if (terminal.alternate !== null) {
          this.visitBlock(terminal.alternate, state);
        }
        break;
      }
      case "switch": {
        this.visitPlace(terminal.id, terminal.test, state);
        for (const case_ of terminal.cases) {
          if (case_.test !== null) {
            this.visitPlace(terminal.id, case_.test, state);
          }
          if (case_.block !== undefined) {
            this.visitBlock(case_.block, state);
          }
        }
        break;
      }
      case "label": {
        this.visitBlock(terminal.block, state);
        break;
      }
      case "try": {
        this.visitBlock(terminal.block, state);
        this.visitBlock(terminal.handler, state);
        break;
      }
      default: {
        assertExhaustive(
          terminal,
          `Unexpected terminal kind \`${(terminal as any).kind}\``
        );
      }
    }
  }

  visitScope(scope: ReactiveScopeBlock, state: TState): void {
    this.traverseScope(scope, state);
  }
  traverseScope(scope: ReactiveScopeBlock, state: TState): void {
    this.visitBlock(scope.instructions, state);
  }

  visitPrunedScope(scopeBlock: PrunedReactiveScopeBlock, state: TState): void {
    this.traversePrunedScope(scopeBlock, state);
  }
  traversePrunedScope(
    scopeBlock: PrunedReactiveScopeBlock,
    state: TState
  ): void {
    this.visitBlock(scopeBlock.instructions, state);
  }

  visitBlock(block: ReactiveBlock, state: TState): void {
    this.traverseBlock(block, state);
  }
  traverseBlock(block: ReactiveBlock, state: TState): void {
    for (const instr of block) {
      switch (instr.kind) {
        case "instruction": {
          this.visitInstruction(instr.instruction, state);
          break;
        }
        case "scope": {
          this.visitScope(instr, state);
          break;
        }
        case "pruned-scope": {
          this.visitPrunedScope(instr, state);
          break;
        }
        case "terminal": {
          this.visitTerminal(instr, state);
          break;
        }
        default: {
          assertExhaustive(
            instr,
            `Unexpected instruction kind \`${(instr as any).kind}\``
          );
        }
      }
    }
  }

  visitHirFunction(fn: HIRFunction, state: TState): void {
    for (const param of fn.params) {
      const place = param.kind === "Identifier" ? param : param.place;
      this.visitParam(place, state);
    }
    for (const [, block] of fn.body.blocks) {
      for (const instr of block.instructions) {
        this.visitInstruction(instr, state);
        if (
          instr.value.kind === "FunctionExpression" ||
          instr.value.kind === "ObjectMethod"
        ) {
          this.visitHirFunction(instr.value.loweredFunc.func, state);
        }
      }
      for (const operand of eachTerminalOperand(block.terminal)) {
        this.visitPlace(block.terminal.id, operand, state);
      }
    }
  }
}

export type TransformedValue =
  | { kind: "keep" }
  | { kind: "replace"; value: ReactiveValue };

export type Transformed<T> =
  | { kind: "remove" }
  | { kind: "keep" }
  | { kind: "replace"; value: T }
  | { kind: "replace-many"; value: Array<T> };

export class ReactiveFunctionTransform<
  TState = void,
> extends ReactiveFunctionVisitor<TState> {
  override traverseBlock(block: ReactiveBlock, state: TState): void {
    let nextBlock: ReactiveBlock | null = null;
    for (let i = 0; i < block.length; i++) {
      const instr = block[i]!;
      let transformed: Transformed<ReactiveStatement>;
      switch (instr.kind) {
        case "instruction": {
          transformed = this.transformInstruction(instr.instruction, state);
          break;
        }
        case "scope": {
          transformed = this.transformScope(instr, state);
          break;
        }
        case "pruned-scope": {
          transformed = this.transformPrunedScope(instr, state);
          break;
        }
        case "terminal": {
          transformed = this.transformTerminal(instr, state);
          break;
        }
        default: {
          assertExhaustive(
            instr,
            `Unexpected instruction kind \`${(instr as any).kind}\``
          );
        }
      }
      switch (transformed.kind) {
        case "keep": {
          if (nextBlock !== null) {
            nextBlock.push(instr);
          }
          break;
        }
        case "remove": {
          if (nextBlock === null) {
            nextBlock = block.slice(0, i);
          }
          break;
        }
        case "replace": {
          nextBlock ??= block.slice(0, i);
          nextBlock.push(transformed.value);
          break;
        }
        case "replace-many": {
          nextBlock ??= block.slice(0, i);
          nextBlock.push(...transformed.value);
          break;
        }
      }
    }
    if (nextBlock !== null) {
      block.length = 0;
      block.push(...nextBlock);
    }
  }

  transformInstruction(
    instruction: ReactiveInstruction,
    state: TState
  ): Transformed<ReactiveStatement> {
    this.visitInstruction(instruction, state);
    return { kind: "keep" };
  }

  transformTerminal(
    stmt: ReactiveTerminalStatement,
    state: TState
  ): Transformed<ReactiveStatement> {
    this.visitTerminal(stmt, state);
    return { kind: "keep" };
  }

  transformScope(
    scope: ReactiveScopeBlock,
    state: TState
  ): Transformed<ReactiveStatement> {
    this.visitScope(scope, state);
    return { kind: "keep" };
  }

  transformPrunedScope(
    scope: PrunedReactiveScopeBlock,
    state: TState
  ): Transformed<ReactiveStatement> {
    this.visitPrunedScope(scope, state);
    return { kind: "keep" };
  }

  transformValue(
    id: InstructionId,
    value: ReactiveValue,
    state: TState
  ): TransformedValue {
    this.visitValue(id, value, state);
    return { kind: "keep" };
  }

  transformReactiveFunctionValue(
    id: InstructionId,
    dependencies: Array<Place>,
    fn: ReactiveFunction,
    state: TState
  ): { kind: "keep" } | { kind: "replace"; value: ReactiveFunction } {
    this.visitReactiveFunctionValue(id, dependencies, fn, state);
    return { kind: "keep" };
  }

  override traverseValue(
    id: InstructionId,
    value: ReactiveValue,
    state: TState
  ): void {
    switch (value.kind) {
      case "OptionalExpression": {
        const nextValue = this.transformValue(id, value.value, state);
        if (nextValue.kind === "replace") {
          value.value = nextValue.value;
        }
        break;
      }
      case "LogicalExpression": {
        const left = this.transformValue(id, value.left, state);
        if (left.kind === "replace") {
          value.left = left.value;
        }
        const right = this.transformValue(id, value.right, state);
        if (right.kind === "replace") {
          value.right = right.value;
        }
        break;
      }
      case "ConditionalExpression": {
        const test = this.transformValue(id, value.test, state);
        if (test.kind === "replace") {
          value.test = test.value;
        }
        const consequent = this.transformValue(id, value.consequent, state);
        if (consequent.kind === "replace") {
          value.consequent = consequent.value;
        }
        const alternate = this.transformValue(id, value.alternate, state);
        if (alternate.kind === "replace") {
          value.alternate = alternate.value;
        }
        break;
      }
      case "SequenceExpression": {
        for (const instr of value.instructions) {
          this.visitInstruction(instr, state);
        }
        const nextValue = this.transformValue(value.id, value.value, state);
        if (nextValue.kind === "replace") {
          value.value = nextValue.value;
        }
        break;
      }
      case "ReactiveFunctionValue": {
        const nextValue = this.transformReactiveFunctionValue(
          id,
          value.dependencies,
          value.fn,
          state
        );
        if (nextValue.kind === "replace") {
          value.fn = nextValue.value;
        }
        break;
      }
      default: {
        for (const place of eachInstructionValueOperand(value)) {
          this.visitPlace(id, place, state);
        }
      }
    }
  }

  override traverseInstruction(
    instruction: ReactiveInstruction,
    state: TState
  ): void {
    this.visitID(instruction.id, state);
    for (const operand of eachInstructionLValue(instruction)) {
      this.visitLValue(instruction.id, operand, state);
    }
    const nextValue = this.transformValue(
      instruction.id,
      instruction.value,
      state
    );
    if (nextValue.kind === "replace") {
      instruction.value = nextValue.value;
    }
  }

  override traverseTerminal(
    stmt: ReactiveTerminalStatement,
    state: TState
  ): void {
    const { terminal } = stmt;
    if (terminal.id !== null) {
      this.visitID(terminal.id, state);
    }
    switch (terminal.kind) {
      case "break":
      case "continue": {
        break;
      }
      case "return": {
        this.visitPlace(terminal.id, terminal.value, state);
        break;
      }
      case "throw": {
        this.visitPlace(terminal.id, terminal.value, state);
        break;
      }
      case "for": {
        const init = this.transformValue(terminal.id, terminal.init, state);
        if (init.kind === "replace") {
          terminal.init = init.value;
        }
        const test = this.transformValue(terminal.id, terminal.test, state);
        if (test.kind === "replace") {
          terminal.test = test.value;
        }
        if (terminal.update !== null) {
          const update = this.transformValue(
            terminal.id,
            terminal.update,
            state
          );
          if (update.kind === "replace") {
            terminal.update = update.value;
          }
        }
        this.visitBlock(terminal.loop, state);
        break;
      }
      case "for-of": {
        const init = this.transformValue(terminal.id, terminal.init, state);
        if (init.kind === "replace") {
          terminal.init = init.value;
        }
        const test = this.transformValue(terminal.id, terminal.test, state);
        if (test.kind === "replace") {
          terminal.test = test.value;
        }
        this.visitBlock(terminal.loop, state);
        break;
      }
      case "for-in": {
        const init = this.transformValue(terminal.id, terminal.init, state);
        if (init.kind === "replace") {
          terminal.init = init.value;
        }
        this.visitBlock(terminal.loop, state);
        break;
      }
      case "do-while": {
        this.visitBlock(terminal.loop, state);
        const test = this.transformValue(terminal.id, terminal.test, state);
        if (test.kind === "replace") {
          terminal.test = test.value;
        }
        break;
      }
      case "while": {
        const test = this.transformValue(terminal.id, terminal.test, state);
        if (test.kind === "replace") {
          terminal.test = test.value;
        }
        this.visitBlock(terminal.loop, state);
        break;
      }
      case "if": {
        this.visitPlace(terminal.id, terminal.test, state);
        this.visitBlock(terminal.consequent, state);
        if (terminal.alternate !== null) {
          this.visitBlock(terminal.alternate, state);
        }
        break;
      }
      case "switch": {
        this.visitPlace(terminal.id, terminal.test, state);
        for (const case_ of terminal.cases) {
          if (case_.test !== null) {
            this.visitPlace(terminal.id, case_.test, state);
          }
          if (case_.block !== undefined) {
            this.visitBlock(case_.block, state);
          }
        }
        break;
      }
      case "label": {
        this.visitBlock(terminal.block, state);
        break;
      }
      case "try": {
        this.visitBlock(terminal.block, state);
        if (terminal.handlerBinding !== null) {
          this.visitPlace(terminal.id, terminal.handlerBinding, state);
        }
        this.visitBlock(terminal.handler, state);
        break;
      }
      default: {
        assertExhaustive(
          terminal,
          `Unexpected terminal kind \`${(terminal as any).kind}\``
        );
      }
    }
  }
}

export function* eachReactiveValueOperand(
  instrValue: ReactiveValue
): Iterable<Place> {
  switch (instrValue.kind) {
    case "OptionalExpression": {
      yield* eachReactiveValueOperand(instrValue.value);
      break;
    }
    case "LogicalExpression": {
      yield* eachReactiveValueOperand(instrValue.left);
      yield* eachReactiveValueOperand(instrValue.right);
      break;
    }
    case "SequenceExpression": {
      for (const instr of instrValue.instructions) {
        yield* eachReactiveValueOperand(instr.value);
      }
      yield* eachReactiveValueOperand(instrValue.value);
      break;
    }
    case "ConditionalExpression": {
      yield* eachReactiveValueOperand(instrValue.test);
      yield* eachReactiveValueOperand(instrValue.consequent);
      yield* eachReactiveValueOperand(instrValue.alternate);
      break;
    }
    case "ReactiveFunctionValue": {
      yield* instrValue.dependencies;
      break;
    }
    default: {
      yield* eachInstructionValueOperand(instrValue);
    }
  }
}

export function mapTerminalBlocks(
  terminal: ReactiveTerminal,
  fn: (block: ReactiveBlock) => ReactiveBlock
): void {
  switch (terminal.kind) {
    case "break":
    case "continue":
    case "return":
    case "throw": {
      break;
    }
    case "for": {
      terminal.loop = fn(terminal.loop);
      break;
    }
    case "for-of": {
      terminal.loop = fn(terminal.loop);
      break;
    }
    case "for-in": {
      terminal.loop = fn(terminal.loop);
      break;
    }
    case "do-while":
    case "while": {
      terminal.loop = fn(terminal.loop);
      break;
    }
    case "if": {
      terminal.consequent = fn(terminal.consequent);
      if (terminal.alternate !== null) {
        terminal.alternate = fn(terminal.alternate);
      }
      break;
    }
    case "switch": {
      for (const case_ of terminal.cases) {
        if (case_.block !== undefined) {
          case_.block = fn(case_.block);
        }
      }
      break;
    }
    case "label": {
      terminal.block = fn(terminal.block);
      break;
    }
    case "try": {
      terminal.block = fn(terminal.block);
      terminal.handler = fn(terminal.handler);
      break;
    }
    default: {
      assertExhaustive(
        terminal,
        `Unexpected terminal kind \`${(terminal as any).kind}\``
      );
    }
  }
}