/**
 * 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 {CompilerError} from '../CompilerError';
import {DependencyPath, Identifier, ReactiveScopeDependency} from '../HIR';
import {printIdentifier} from '../HIR/PrintHIR';
import {assertExhaustive} from '../Utils/utils';

/*
 * We need to understand optional member expressions only when determining
 * dependencies of a ReactiveScope (i.e. in {@link PropagateScopeDependencies}),
 * hence why this type lives here (not in HIR.ts)
 */
export type ReactiveScopePropertyDependency = ReactiveScopeDependency;

/*
 * Finalizes a set of ReactiveScopeDependencies to produce a set of minimal unconditional
 * dependencies, preserving granular accesses when possible.
 *
 * Correctness properties:
 *   - All dependencies to a ReactiveBlock must be tracked.
 *     We can always truncate a dependency's path to a subpath, due to Forget assuming
 *     deep immutability. If the value produced by a subpath has not changed, then
 *     dependency must have not changed.
 *     i.e. props.a === $[..] implies props.a.b === $[..]
 *
 *     Note the inverse is not true, but this only means a false positive (we run the
 *     reactive block more than needed).
 *     i.e. props.a !== $[..] does not imply props.a.b !== $[..]
 *
 *   - The dependencies of a finalized ReactiveBlock must be all safe to access
 *     unconditionally (i.e. preserve program semantics with respect to nullthrows).
 *     If a dependency is only accessed within a conditional, we must track the nearest
 *     unconditionally accessed subpath instead.
 * @param initialDeps
 * @returns
 */
export class ReactiveScopeDependencyTree {
  #roots: Map<Identifier, DependencyNode> = new Map();

  #getOrCreateRoot(identifier: Identifier): DependencyNode {
    // roots can always be accessed unconditionally in JS
    let rootNode = this.#roots.get(identifier);

    if (rootNode === undefined) {
      rootNode = {
        properties: new Map(),
        accessType: PropertyAccessType.UnconditionalAccess,
      };
      this.#roots.set(identifier, rootNode);
    }
    return rootNode;
  }

  add(dep: ReactiveScopePropertyDependency, inConditional: boolean): void {
    const {path} = dep;
    let currNode = this.#getOrCreateRoot(dep.identifier);

    for (const item of path) {
      // all properties read 'on the way' to a dependency are marked as 'access'
      let currChild = getOrMakeProperty(currNode, item.property);
      const accessType = inConditional
        ? PropertyAccessType.ConditionalAccess
        : item.optional
          ? PropertyAccessType.OptionalAccess
          : PropertyAccessType.UnconditionalAccess;
      currChild.accessType = merge(currChild.accessType, accessType);
      currNode = currChild;
    }

    /**
     * The final property node should be marked as an conditional/unconditional
     * `dependency` as based on control flow.
     */
    const depType = inConditional
      ? PropertyAccessType.ConditionalDependency
      : isOptional(currNode.accessType)
        ? PropertyAccessType.OptionalDependency
        : PropertyAccessType.UnconditionalDependency;

    currNode.accessType = merge(currNode.accessType, depType);
  }

  deriveMinimalDependencies(): Set<ReactiveScopeDependency> {
    const results = new Set<ReactiveScopeDependency>();
    for (const [rootId, rootNode] of this.#roots.entries()) {
      const deps = deriveMinimalDependenciesInSubtree(rootNode, null);
      CompilerError.invariant(
        deps.every(
          dep =>
            dep.accessType === PropertyAccessType.UnconditionalDependency ||
            dep.accessType == PropertyAccessType.OptionalDependency,
        ),
        {
          reason:
            '[PropagateScopeDependencies] All dependencies must be reduced to unconditional dependencies.',
          description: null,
          loc: null,
          suggestions: null,
        },
      );

      for (const dep of deps) {
        results.add({
          identifier: rootId,
          path: dep.relativePath,
        });
      }
    }

    return results;
  }

  addDepsFromInnerScope(
    depsFromInnerScope: ReactiveScopeDependencyTree,
    innerScopeInConditionalWithinParent: boolean,
    checkValidDepIdFn: (dep: ReactiveScopeDependency) => boolean,
  ): void {
    for (const [id, otherRoot] of depsFromInnerScope.#roots) {
      if (!checkValidDepIdFn({identifier: id, path: []})) {
        continue;
      }
      let currRoot = this.#getOrCreateRoot(id);
      addSubtree(currRoot, otherRoot, innerScopeInConditionalWithinParent);
      if (!isUnconditional(currRoot.accessType)) {
        currRoot.accessType = isDependency(currRoot.accessType)
          ? PropertyAccessType.UnconditionalDependency
          : PropertyAccessType.UnconditionalAccess;
      }
    }
  }

  promoteDepsFromExhaustiveConditionals(
    trees: Array<ReactiveScopeDependencyTree>,
  ): void {
    CompilerError.invariant(trees.length > 1, {
      reason: 'Expected trees to be at least 2 elements long.',
      description: null,
      loc: null,
      suggestions: null,
    });

    for (const [id, root] of this.#roots) {
      const nodesForRootId = mapNonNull(trees, tree => {
        const node = tree.#roots.get(id);
        if (node != null && isUnconditional(node.accessType)) {
          return node;
        } else {
          return null;
        }
      });
      if (nodesForRootId) {
        addSubtreeIntersection(
          root.properties,
          nodesForRootId.map(root => root.properties),
        );
      }
    }
  }

  /*
   * Prints dependency tree to string for debugging.
   * @param includeAccesses
   * @returns string representation of DependencyTree
   */
  printDeps(includeAccesses: boolean): string {
    let res = [];

    for (const [rootId, rootNode] of this.#roots.entries()) {
      const rootResults = printSubtree(rootNode, includeAccesses).map(
        result => `${printIdentifier(rootId)}.${result}`,
      );
      res.push(rootResults);
    }
    return res.flat().join('\n');
  }

  debug(): string {
    const buf: Array<string> = [`tree() [`];
    for (const [rootId, rootNode] of this.#roots) {
      buf.push(`${printIdentifier(rootId)} (${rootNode.accessType}):`);
      this.#debugImpl(buf, rootNode, 1);
    }
    buf.push(']');
    return buf.length > 2 ? buf.join('\n') : buf.join('');
  }

  #debugImpl(
    buf: Array<string>,
    node: DependencyNode,
    depth: number = 0,
  ): void {
    for (const [property, childNode] of node.properties) {
      buf.push(`${'  '.repeat(depth)}.${property} (${childNode.accessType}):`);
      this.#debugImpl(buf, childNode, depth + 1);
    }
  }
}

/*
 * Enum representing the access type of single property on a parent object.
 * We distinguish on two independent axes:
 * Conditional / Unconditional:
 *    - whether this property is accessed unconditionally (within the ReactiveBlock)
 * Access / Dependency:
 *    - Access: this property is read on the path of a dependency. We do not
 *      need to track change variables for accessed properties. Tracking accesses
 *      helps Forget do more granular dependency tracking.
 *    - Dependency: this property is read as a dependency and we must track changes
 *      to it for correctness.
 *
 *    ```javascript
 *    // props.a is a dependency here and must be tracked
 *    deps: {props.a, props.a.b} ---> minimalDeps: {props.a}
 *    // props.a is just an access here and does not need to be tracked
 *    deps: {props.a.b} ---> minimalDeps: {props.a.b}
 *    ```
 */
enum PropertyAccessType {
  ConditionalAccess = 'ConditionalAccess',
  OptionalAccess = 'OptionalAccess',
  UnconditionalAccess = 'UnconditionalAccess',
  ConditionalDependency = 'ConditionalDependency',
  OptionalDependency = 'OptionalDependency',
  UnconditionalDependency = 'UnconditionalDependency',
}

const MIN_ACCESS_TYPE = PropertyAccessType.ConditionalAccess;
function isUnconditional(access: PropertyAccessType): boolean {
  return (
    access === PropertyAccessType.UnconditionalAccess ||
    access === PropertyAccessType.UnconditionalDependency
  );
}
function isDependency(access: PropertyAccessType): boolean {
  return (
    access === PropertyAccessType.ConditionalDependency ||
    access === PropertyAccessType.OptionalDependency ||
    access === PropertyAccessType.UnconditionalDependency
  );
}
function isOptional(access: PropertyAccessType): boolean {
  return (
    access === PropertyAccessType.OptionalAccess ||
    access === PropertyAccessType.OptionalDependency
  );
}

function merge(
  access1: PropertyAccessType,
  access2: PropertyAccessType,
): PropertyAccessType {
  const resultIsUnconditional =
    isUnconditional(access1) || isUnconditional(access2);
  const resultIsDependency = isDependency(access1) || isDependency(access2);
  const resultIsOptional = isOptional(access1) || isOptional(access2);

  /*
   * Straightforward merge.
   * This can be represented as bitwise OR, but is written out for readability
   *
   * Observe that `UnconditionalAccess | ConditionalDependency` produces an
   * unconditionally accessed conditional dependency. We currently use these
   * as we use unconditional dependencies. (i.e. to codegen change variables)
   */
  if (resultIsUnconditional) {
    if (resultIsDependency) {
      return PropertyAccessType.UnconditionalDependency;
    } else {
      return PropertyAccessType.UnconditionalAccess;
    }
  } else if (resultIsOptional) {
    if (resultIsDependency) {
      return PropertyAccessType.OptionalDependency;
    } else {
      return PropertyAccessType.OptionalAccess;
    }
  } else {
    if (resultIsDependency) {
      return PropertyAccessType.ConditionalDependency;
    } else {
      return PropertyAccessType.ConditionalAccess;
    }
  }
}

type DependencyNode = {
  properties: Map<string, DependencyNode>;
  accessType: PropertyAccessType;
};

type ReduceResultNode = {
  relativePath: DependencyPath;
  accessType: PropertyAccessType;
};

function promoteResult(
  accessType: PropertyAccessType,
  path: {property: string; optional: boolean} | null,
): Array<ReduceResultNode> {
  const result: ReduceResultNode = {
    relativePath: [],
    accessType,
  };
  if (path !== null) {
    result.relativePath.push(path);
  }
  return [result];
}

function prependPath(
  results: Array<ReduceResultNode>,
  path: {property: string; optional: boolean} | null,
): Array<ReduceResultNode> {
  if (path === null) {
    return results;
  }
  return results.map(result => {
    return {
      accessType: result.accessType,
      relativePath: [path, ...result.relativePath],
    };
  });
}

/*
 * Recursively calculates minimal dependencies in a subtree.
 * @param dep DependencyNode representing a dependency subtree.
 * @returns a minimal list of dependencies in this subtree.
 */
function deriveMinimalDependenciesInSubtree(
  dep: DependencyNode,
  property: string | null,
): Array<ReduceResultNode> {
  const results: Array<ReduceResultNode> = [];
  for (const [childName, childNode] of dep.properties) {
    const childResult = deriveMinimalDependenciesInSubtree(
      childNode,
      childName,
    );
    results.push(...childResult);
  }

  switch (dep.accessType) {
    case PropertyAccessType.UnconditionalDependency: {
      return promoteResult(
        PropertyAccessType.UnconditionalDependency,
        property !== null ? {property, optional: false} : null,
      );
    }
    case PropertyAccessType.UnconditionalAccess: {
      if (
        results.every(
          ({accessType}) =>
            accessType === PropertyAccessType.UnconditionalDependency ||
            accessType === PropertyAccessType.OptionalDependency,
        )
      ) {
        // all children are unconditional dependencies, return them to preserve granularity
        return prependPath(
          results,
          property !== null ? {property, optional: false} : null,
        );
      } else {
        /*
         * at least one child is accessed conditionally, so this node needs to be promoted to
         * unconditional dependency
         */
        return promoteResult(
          PropertyAccessType.UnconditionalDependency,
          property !== null ? {property, optional: false} : null,
        );
      }
    }
    case PropertyAccessType.OptionalDependency: {
      return promoteResult(
        PropertyAccessType.OptionalDependency,
        property !== null ? {property, optional: true} : null,
      );
    }
    case PropertyAccessType.OptionalAccess: {
      if (
        results.every(
          ({accessType}) =>
            accessType === PropertyAccessType.UnconditionalDependency ||
            accessType === PropertyAccessType.OptionalDependency,
        )
      ) {
        // all children are unconditional dependencies, return them to preserve granularity
        return prependPath(
          results,
          property !== null ? {property, optional: true} : null,
        );
      } else {
        /*
         * at least one child is accessed conditionally, so this node needs to be promoted to
         * unconditional dependency
         */
        return promoteResult(
          PropertyAccessType.OptionalDependency,
          property !== null ? {property, optional: true} : null,
        );
      }
    }
    case PropertyAccessType.ConditionalAccess:
    case PropertyAccessType.ConditionalDependency: {
      if (
        results.every(
          ({accessType}) =>
            accessType === PropertyAccessType.ConditionalDependency,
        )
      ) {
        /*
         * No children are accessed unconditionally, so we cannot promote this node to
         * unconditional access.
         * Truncate results of child nodes here, since we shouldn't access them anyways
         */
        return promoteResult(
          PropertyAccessType.ConditionalDependency,
          property !== null ? {property, optional: true} : null,
        );
      } else {
        /*
         * at least one child is accessed unconditionally, so this node can be promoted to
         * unconditional dependency
         */
        return promoteResult(
          PropertyAccessType.UnconditionalDependency,
          property !== null ? {property, optional: true} : null,
        );
      }
    }
    default: {
      assertExhaustive(
        dep.accessType,
        '[PropgateScopeDependencies] Unhandled access type!',
      );
    }
  }
}

/*
 * Demote all unconditional accesses + dependencies in subtree to the
 * conditional equivalent, mutating subtree in place.
 * @param subtree unconditional node representing a subtree of dependencies
 */
function demoteSubtreeToConditional(subtree: DependencyNode): void {
  const stack: Array<DependencyNode> = [subtree];

  let node;
  while ((node = stack.pop()) !== undefined) {
    const {accessType, properties} = node;
    if (!isUnconditional(accessType)) {
      // A conditionally accessed node should not have unconditional children
      continue;
    }
    node.accessType = isDependency(accessType)
      ? PropertyAccessType.ConditionalDependency
      : PropertyAccessType.ConditionalAccess;

    for (const childNode of properties.values()) {
      if (isUnconditional(accessType)) {
        /*
         * No conditional node can have an unconditional node as a child, so
         * we only process childNode if it is unconditional
         */
        stack.push(childNode);
      }
    }
  }
}

/*
 * Calculates currNode = union(currNode, otherNode), mutating currNode in place
 * If demoteOtherNode is specified, we demote the subtree represented by
 * otherNode to conditional access/deps before taking the union.
 *
 * This is a helper function used to join an inner scope to its parent scope.
 * @param currNode (mutable) return by argument
 * @param otherNode (move) {@link addSubtree} takes ownership of the subtree
 * represented by otherNode, which may be mutated or moved to currNode. It is
 * invalid to use otherNode after this call.
 *
 * Note that @param otherNode may contain both conditional and unconditional nodes,
 * due to inner control flow and conditional member expressions
 *
 * @param demoteOtherNode
 */
function addSubtree(
  currNode: DependencyNode,
  otherNode: DependencyNode,
  demoteOtherNode: boolean,
): void {
  let otherType = otherNode.accessType;
  if (demoteOtherNode) {
    otherType = isDependency(otherType)
      ? PropertyAccessType.ConditionalDependency
      : PropertyAccessType.ConditionalAccess;
  }
  currNode.accessType = merge(currNode.accessType, otherType);

  for (const [propertyName, otherChild] of otherNode.properties) {
    const currChild = currNode.properties.get(propertyName);
    if (currChild) {
      // recursively calculate currChild = union(currChild, otherChild)
      addSubtree(currChild, otherChild, demoteOtherNode);
    } else {
      /*
       * if currChild doesn't exist, we can just move otherChild
       * currChild = otherChild.
       */
      if (demoteOtherNode) {
        demoteSubtreeToConditional(otherChild);
      }
      currNode.properties.set(propertyName, otherChild);
    }
  }
}

/*
 * Adds intersection(otherProperties) to currProperties, mutating
 * currProperties in place. i.e.
 *    currProperties = union(currProperties, intersection(otherProperties))
 *
 * Used to merge unconditional accesses from exhaustive conditional branches
 * into the parent ReactiveDeps Tree.
 * intersection(currProperties) is determined as such:
 *   - a node is present in the intersection iff it is present in all every
 *     branch
 *   - the type of an added node is `UnconditionalDependency` if it is a
 *     dependency in at least one branch (otherwise `UnconditionalAccess`)
 *
 * @param otherProperties (read-only) an array of node properties containing
 *         conditionally and unconditionally accessed nodes. Each element
 *         represents asubtree of reactive dependencies from a single CFG
 *         branch.
 *        otherProperties must represent all reachable branches.
 * @param currProperties (mutable) return by argument properties of a node
 *
 * otherProperties and currProperties must be properties of disjoint nodes
 * that represent the same dependency (identifier + path).
 */
function addSubtreeIntersection(
  currProperties: Map<string, DependencyNode>,
  otherProperties: Array<Map<string, DependencyNode>>,
): void {
  CompilerError.invariant(otherProperties.length > 1, {
    reason:
      '[DeriveMinimalDependencies] Expected otherProperties to be at least 2 elements long.',
    description: null,
    loc: null,
    suggestions: null,
  });

  /*
   * otherProperties here may contain unconditional nodes as the result of
   * recursively merging exhaustively conditional children with unconditionally
   * accessed nodes (e.g. in the test condition itself)
   * See `reduce-reactive-cond-deps-cfg-nested-testifelse` fixture for example
   */

  for (const [propertyName, currNode] of currProperties) {
    const otherNodes = mapNonNull(otherProperties, properties => {
      const node = properties.get(propertyName);
      if (node != null && isUnconditional(node.accessType)) {
        return node;
      } else {
        return null;
      }
    });

    /*
     * intersection(otherNodes[propertyName]) only exists if each element in
     * otherProperties accesses propertyName.
     */
    if (otherNodes) {
      addSubtreeIntersection(
        currNode.properties,
        otherNodes.map(node => node.properties),
      );

      const isDep = otherNodes.some(tree => isDependency(tree.accessType));
      const externalAccessType = isDep
        ? PropertyAccessType.UnconditionalDependency
        : PropertyAccessType.UnconditionalAccess;
      currNode.accessType = merge(externalAccessType, currNode.accessType);
    }
  }
}

function printSubtree(
  node: DependencyNode,
  includeAccesses: boolean,
): Array<string> {
  const results: Array<string> = [];
  for (const [propertyName, propertyNode] of node.properties) {
    if (includeAccesses || isDependency(propertyNode.accessType)) {
      results.push(`${propertyName} (${propertyNode.accessType})`);
    }
    const propertyResults = printSubtree(propertyNode, includeAccesses);
    results.push(...propertyResults.map(result => `${propertyName}.${result}`));
  }
  return results;
}

function getOrMakeProperty(
  node: DependencyNode,
  property: string,
): DependencyNode {
  let child = node.properties.get(property);
  if (child == null) {
    child = {
      properties: new Map(),
      accessType: MIN_ACCESS_TYPE,
    };
    node.properties.set(property, child);
  }
  return child;
}

function mapNonNull<T extends NonNullable<V>, V, U>(
  arr: Array<U>,
  fn: (arg0: U) => T | undefined | null,
): Array<T> | null {
  const result = [];
  for (let i = 0; i < arr.length; i++) {
    const element = fn(arr[i]);
    if (element) {
      result.push(element);
    } else {
      return null;
    }
  }
  return result;
}