import { Value } from "@convex-dev/common";
import { GenericAPI, NamedQuery, QueryNames } from "../api.js";
import { createError } from "../logging.js";
import { MutationId } from "./protocol.js";
import { QueryResult } from "./remote_query_set.js";
import { QueryToken, serializePathAndArgs } from "./udf_path_utils.js";

/**
 * A view of the query results currently in the Convex client for use within
 * optimistic updates.
 *
 * @public
 */
export interface OptimisticLocalStore<API extends GenericAPI = GenericAPI> {
  /**
   * Retrieve the result of a query from the client.
   *
   * Important: Query results should be treated as immutable!
   * Always make new copies of structures within query results to avoid
   * corrupting data within the client.
   *
   * @param name - The name of the query.
   * @param args - An array of the arguments for this query.
   * @returns The query result or `undefined` if the query is not currently
   * in the client.
   */
  getQuery<Name extends QueryNames<API>>(
    name: Name,
    args: Parameters<NamedQuery<API, Name>>
  ): undefined | ReturnType<NamedQuery<API, Name>>;

  /**
   * Optimistically update the result of a query.
   *
   * This can either be a new value (perhaps derived from the old value from
   * {@link OptimisticLocalStore.getQuery}) or `undefined` to remove the query.
   * Removing a query is useful to create loading states while Convex recomputes
   * the query results.
   *
   * @param name - The name of the query.
   * @param args - An array of the arguments for this query.
   * @param value - The new value to set the query to or `undefined` to remove
   * it from the client.
   */
  setQuery<Name extends QueryNames<API>>(
    name: Name,
    args: Parameters<NamedQuery<API, Name>>,
    value: undefined | ReturnType<NamedQuery<API, Name>>
  ): void;
}
/**
 * A temporary, local update to query results within this client.
 *
 * This update will always be executed when a mutation is synced to the Convex
 * server and rolled back when the mutation completes.
 *
 * Note that optimistic updates can be called multiple times! If the client
 * loads new data while the mutation is in progress, the update will be replayed
 * again.
 *
 * @param localQueryStore - An interface to read and edit local query results.
 * @param args - The arguments to the mutation.
 *
 * @public
 */
export type OptimisticUpdate<
  API extends GenericAPI,
  Arguments extends any[]
> = (localQueryStore: OptimisticLocalStore<API>, ...args: Arguments) => void;

/**
 * An optimistic update function that has been curried over its arguments.
 */
type WrappedOptimisticUpdate = (locaQueryStore: OptimisticLocalStore) => void;

/**
 * The implementation of `OptimisticLocalStore`.
 *
 * This class provides the interface for optimistic updates to modify query results.
 */
class OptimisticLocalStoreImpl implements OptimisticLocalStore {
  // A references of the query results in OptimisticQueryResults
  private readonly queryResults: QueryResultsMap;

  // All of the queries modified by this class
  readonly modifiedQueries: QueryToken[];

  constructor(queryResults: QueryResultsMap) {
    this.queryResults = queryResults;
    this.modifiedQueries = [];
  }

  getQuery(name: string, args: Value[]): Value | undefined {
    const query = this.queryResults.get(serializePathAndArgs(name, args));
    if (query === undefined) {
      return undefined;
    }
    const result = query.result;
    if (result === undefined) {
      return undefined;
    } else if (result.success) {
      return result.value;
    } else {
      // If the query is an error state, just return `undefined` as though
      // it's loading. Optimistic updates should already handle `undefined` well
      // and there isn't a need to break the whole update because it tried
      // to load a single query that errored.
      return undefined;
    }
  }

  setQuery(name: string, args: Value[], value: Value | undefined): void {
    const queryToken = serializePathAndArgs(name, args);

    let result: QueryResult | undefined;
    if (value === undefined) {
      result = undefined;
    } else {
      result = {
        success: true,
        value,
      };
    }
    const query: Query = {
      udfPath: name,
      args,
      result,
    };
    this.queryResults.set(queryToken, query);
    this.modifiedQueries.push(queryToken);
  }
}

type OptimisticUpdateAndId = {
  update: WrappedOptimisticUpdate;
  mutationId: MutationId;
};

type Query = {
  // undefined means the query was set to be loading (undefined) in an optimistic update.
  // Note that we can also have queries not present in the QueryResultMap
  // at all because they are still loading from the server.
  result: QueryResult | undefined;
  udfPath: string;
  args: Value[];
};
export type QueryResultsMap = Map<QueryToken, Query>;

type ChangedQueries = QueryToken[];

/**
 * A view of all of our query results with optimistic updates applied on top.
 */
export class OptimisticQueryResults {
  private queryResults: QueryResultsMap;
  private optimisticUpdates: OptimisticUpdateAndId[];

  constructor() {
    this.queryResults = new Map();
    this.optimisticUpdates = [];
  }

  ingestQueryResultsFromServer(
    serverQueryResults: QueryResultsMap,
    optimisticUpdatesToDrop: Set<MutationId>
  ): ChangedQueries {
    this.optimisticUpdates = this.optimisticUpdates.filter(updateAndId => {
      return !optimisticUpdatesToDrop.has(updateAndId.mutationId);
    });

    const oldQueryResults = this.queryResults;
    this.queryResults = new Map(serverQueryResults);
    const localStore = new OptimisticLocalStoreImpl(this.queryResults);
    for (const updateAndId of this.optimisticUpdates) {
      updateAndId.update(localStore);
    }

    // To find the changed queries, just do a shallow comparison
    // TODO(CX-733): Change this so we avoid unnecessary rerenders
    const changedQueries: ChangedQueries = [];
    for (const [queryToken, query] of this.queryResults) {
      const oldQuery = oldQueryResults.get(queryToken);
      if (oldQuery === undefined || oldQuery.result !== query.result) {
        changedQueries.push(queryToken);
      }
    }

    return changedQueries;
  }

  applyOptimisticUpdate(
    update: WrappedOptimisticUpdate,
    mutationId: MutationId
  ): ChangedQueries {
    // Apply the update to our store
    this.optimisticUpdates.push({
      update,
      mutationId,
    });
    const localStore = new OptimisticLocalStoreImpl(this.queryResults);
    update(localStore);

    // Notify about any query results that changed
    // TODO(CX-733): Change this so we avoid unnecessary rerenders
    return localStore.modifiedQueries;
  }

  queryResult(queryToken: QueryToken): Value | undefined {
    const query = this.queryResults.get(queryToken);
    if (query === undefined) {
      return undefined;
    }
    const result = query.result;
    if (result === undefined) {
      return undefined;
    } else if (result.success) {
      return result.value;
    } else {
      throw createError("query", query.udfPath, result.errorMessage);
    }
  }
}