// Copyright 2024 Signal Messenger, LLC
// SPDX-License-Identifier: AGPL-3.0-only
import { Aci } from '@signalapp/libsignal-client';
import {
  groupSendEndorsementsDataSchema,
  toGroupSendToken,
} from '../types/GroupSendEndorsements.std.ts';
import type {
  GroupSendCombinedEndorsementRecord,
  GroupSendMemberEndorsementRecord,
  GroupSendToken,
  GroupSendEndorsementsData,
} from '../types/GroupSendEndorsements.std.ts';
import { strictAssert } from './assert.std.ts';
import {
  GroupSecretParams,
  GroupSendEndorsement,
  GroupSendEndorsementsResponse,
  ServerPublicParams,
} from './zkgroup.node.ts';
import type { ServiceIdString } from '../types/ServiceId.std.ts';
import { fromAciObject } from '../types/ServiceId.std.ts';
import { createLogger } from '../logging/log.std.ts';
import type { GroupV2MemberType } from '../model-types.d.ts';
import { DurationInSeconds } from './durations/index.std.ts';
import { parseStrict } from './schemas.std.ts';
import { DataReader } from '../sql/Client.preload.ts';
import { maybeUpdateGroup } from '../groups.preload.ts';
import * as Bytes from '../Bytes.std.ts';
import { isGroupV2 } from './whatTypeOfConversation.dom.ts';
import { itemStorage } from '../textsecure/Storage.preload.ts';

const log = createLogger('groupSendEndorsements');

export function decodeGroupSendEndorsementsResponse({
  groupId,
  groupSendEndorsementsResponse,
  groupSecretParamsBase64,
  groupMembersV2,
}: {
  groupId: string;
  groupSendEndorsementsResponse: Uint8Array<ArrayBuffer>;
  groupSecretParamsBase64: string;
  groupMembersV2: ReadonlyArray<GroupV2MemberType>;
}): GroupSendEndorsementsData {
  const idForLogging = `groupv2(${groupId})`;

  strictAssert(
    groupSendEndorsementsResponse != null,
    'Missing groupSendEndorsementsResponse'
  );

  strictAssert(
    groupSendEndorsementsResponse.byteLength > 0,
    'Received empty groupSendEndorsementsResponse'
  );

  const response = new GroupSendEndorsementsResponse(
    groupSendEndorsementsResponse
  );

  const expiration = response.getExpiration().getTime() / 1000;

  const localUser = Aci.parseFromServiceIdString(
    itemStorage.user.getCheckedAci()
  );

  const groupSecretParams = new GroupSecretParams(
    Bytes.fromBase64(groupSecretParamsBase64)
  );

  const serverPublicParams = new ServerPublicParams(
    Bytes.fromBase64(window.getServerPublicParams())
  );

  const groupMembers = groupMembersV2.map(member => {
    return Aci.parseFromServiceIdString(member.aci);
  });

  const receivedEndorsements = response.receiveWithServiceIds(
    groupMembers,
    localUser,
    groupSecretParams,
    serverPublicParams
  );

  const { combinedEndorsement, endorsements } = receivedEndorsements;

  strictAssert(
    endorsements.length === groupMembers.length,
    `Member endorsements must match input array (expected: ${groupMembers.length}, actual: ${endorsements.length})`
  );

  log.info(
    `decodeGroupSendEndorsementsResponse: Received endorsements (group: ${idForLogging}, expiration: ${expiration}, members: ${groupMembers.length})`
  );

  return parseStrict(groupSendEndorsementsDataSchema, {
    combinedEndorsement: {
      groupId,
      expiration,
      endorsement: combinedEndorsement.getContents(),
    },
    memberEndorsements: groupMembers.map((groupMember, index) => {
      const endorsement = endorsements.at(index);
      strictAssert(
        endorsement != null,
        `Missing endorsement at index ${index}`
      );

      return {
        groupId,
        memberAci: fromAciObject(groupMember),
        expiration,
        endorsement: endorsement.getContents(),
      };
    }),
  });
}

const TWO_DAYS = DurationInSeconds.fromDays(2);
const TWO_HOURS = DurationInSeconds.fromHours(2);

function logServiceIds(list: Iterable<string>) {
  const items = Array.from(list);
  if (items.length <= 5) {
    return items.join(', ');
  }
  return `${items.slice(0, 4).join(', ')}, and ${items.length - 4} others`;
}

export type ValidationResult =
  | { valid: true; reason?: never }
  | { valid: false; reason: string };

export function validateGroupSendEndorsementsExpiration(
  expiration: number,
  now: number = Date.now()
): ValidationResult {
  const expSeconds = DurationInSeconds.fromMillis(expiration);
  const nowSeconds = DurationInSeconds.fromMillis(now);
  const info = `now: ${nowSeconds}, exp: ${expSeconds}`;
  if (expSeconds <= nowSeconds) {
    return { valid: false, reason: `already expired, ${info}` };
  }
  // negative = exp is past, positive = exp is future
  const distance = Math.trunc(expSeconds - nowSeconds);
  if (distance <= TWO_HOURS) {
    return { valid: false, reason: `expires soon, ${info}` };
  }
  if (distance >= TWO_DAYS) {
    return { valid: false, reason: `expires too far in future, ${info}` };
  }
  return { valid: true };
}

export class GroupSendEndorsementState {
  readonly #logId: string;
  readonly #combinedEndorsement: GroupSendCombinedEndorsementRecord;
  readonly #memberEndorsements = new Map<
    ServiceIdString,
    GroupSendMemberEndorsementRecord
  >();
  readonly #memberEndorsementsAcis = new Set<ServiceIdString>();
  readonly #groupSecretParamsBase64: string;
  readonly #ourAci: ServiceIdString;
  readonly #endorsementCache = new WeakMap<
    Uint8Array<ArrayBuffer>,
    GroupSendEndorsement
  >();

  constructor(
    data: GroupSendEndorsementsData,
    groupSecretParamsBase64: string
  ) {
    this.#logId = `GroupSendEndorsementState/groupv2(${data.combinedEndorsement.groupId})`;
    this.#combinedEndorsement = data.combinedEndorsement;
    this.#groupSecretParamsBase64 = groupSecretParamsBase64;
    this.#ourAci = itemStorage.user.getCheckedAci();
    for (const endorsement of data.memberEndorsements) {
      this.#memberEndorsements.set(endorsement.memberAci, endorsement);
      this.#memberEndorsementsAcis.add(endorsement.memberAci);
    }
  }

  getExpiration(): Date {
    return new Date(this.#combinedEndorsement.expiration * 1000);
  }

  hasMember(serviceId: ServiceIdString): boolean {
    return this.#memberEndorsements.has(serviceId);
  }

  #toEndorsement(contents: Uint8Array<ArrayBuffer>): GroupSendEndorsement {
    let endorsement = this.#endorsementCache.get(contents);
    if (endorsement == null) {
      endorsement = new GroupSendEndorsement(contents);
      this.#endorsementCache.set(contents, endorsement);
    }
    return endorsement;
  }

  #toToken(endorsement: GroupSendEndorsement): GroupSendToken {
    const groupSecretParams = new GroupSecretParams(
      Bytes.fromBase64(this.#groupSecretParamsBase64)
    );

    const expiration = this.getExpiration();
    const result = validateGroupSendEndorsementsExpiration(
      expiration.getTime()
    );

    strictAssert(
      result.valid,
      `${this.#logId}: toToken: Endorsements are expired (${result.reason})`
    );

    const fullToken = endorsement.toFullToken(groupSecretParams, expiration);
    return toGroupSendToken(fullToken.serialize());
  }

  #getCombinedEndorsement(includesOurs: boolean) {
    const endorsement = this.#toEndorsement(
      this.#combinedEndorsement.endorsement
    );
    if (!includesOurs) {
      return endorsement;
    }
    return GroupSendEndorsement.combine([
      endorsement,
      this.#getMemberEndorsement(this.#ourAci),
    ]);
  }

  #getMemberEndorsement(serviceId: ServiceIdString) {
    const memberEndorsement = this.#memberEndorsements.get(serviceId);
    strictAssert(
      memberEndorsement,
      `${this.#logId}: getMemberEndorsement: Missing endorsement for ${serviceId}`
    );
    return this.#toEndorsement(memberEndorsement.endorsement);
  }

  // Strategy 1: Faster when we're sending to most of the group members
  // `combined.byRemoving(combine(difference(members, sends)))`
  #subtractMemberEndorsements(
    otherMembersServiceIds: Set<ServiceIdString>,
    includesOurs: boolean
  ): GroupSendEndorsement {
    strictAssert(
      !otherMembersServiceIds.has(this.#ourAci),
      `${this.#logId}: subtractMemberEndorsements: Cannot subtract our own aci from the combined endorsement`
    );
    return this.#getCombinedEndorsement(includesOurs).byRemoving(
      this.#combineMemberEndorsements(otherMembersServiceIds)
    );
  }

  // Strategy 2: Faster when we're not sending to most of the group members
  // `combine(sends)`
  #combineMemberEndorsements(
    serviceIds: Set<ServiceIdString>
  ): GroupSendEndorsement {
    return GroupSendEndorsement.combine(
      Array.from(serviceIds, serviceId => {
        return this.#getMemberEndorsement(serviceId);
      })
    );
  }

  #buildToken(serviceIds: Set<ServiceIdString>): GroupSendEndorsement {
    const sendCount = serviceIds.size;
    const memberCount = this.#memberEndorsements.size;
    const logId = `${this.#logId}: buildToken(${sendCount} of ${memberCount})`;

    // Fast path sending to one person
    if (serviceIds.size === 1) {
      const [serviceId] = serviceIds;
      strictAssert(serviceId, 'Missing serviceId');
      log.info(`${logId}: using single member endorsement (${serviceId})`);
      return this.#getMemberEndorsement(serviceId);
    }

    const missing = serviceIds.difference(this.#memberEndorsementsAcis);
    if (missing.size !== 0) {
      throw new Error(
        `${logId}: Attempted to build token with memberAcis we don't have endorsements for (${logServiceIds(missing)})`
      );
    }

    const difference = this.#memberEndorsementsAcis.difference(serviceIds);
    log.info(
      `${logId}: Endorsements without sends ${difference.size}: ${logServiceIds(difference)}`
    );

    const otherMembers = new Set(difference);
    const includesOurs = !otherMembers.delete(this.#ourAci);

    if (otherMembers.size === 0) {
      log.info(
        `${logId}: using combined endorsement (includesOurs: ${includesOurs})`
      );
      return this.#getCombinedEndorsement(includesOurs);
    }

    if (otherMembers.size < memberCount / 2) {
      log.info(
        `${logId}: subtracting missing members (includesOurs: ${includesOurs})`
      );
      return this.#subtractMemberEndorsements(otherMembers, includesOurs);
    }

    log.info(`${logId}: combining all members`);
    return this.#combineMemberEndorsements(serviceIds);
  }

  buildToken(serviceIds: Set<ServiceIdString>): GroupSendToken {
    return this.#toToken(this.#buildToken(new Set(serviceIds)));
  }
}

type MaybeCreateGroupSendEndorsementStateResult =
  | { state: GroupSendEndorsementState; didRefreshGroupState: false }
  | { state: null; didRefreshGroupState: boolean };

function validateGroupSendEndorsements(
  members: Set<ServiceIdString>,
  state: GroupSendEndorsementState | null
): ValidationResult {
  // Check if we should have endorsements (pending members should not)
  if (state == null) {
    const ourAci = itemStorage.user.getCheckedAci();
    if (members.has(ourAci)) {
      return { valid: false, reason: 'missing all endorsements' };
    }
    return { valid: true };
  }

  // Check for expired endorsements
  const expiration = state.getExpiration().getTime();
  const expirationResult = validateGroupSendEndorsementsExpiration(expiration);
  if (!expirationResult.valid) {
    return { valid: false, reason: expirationResult.reason };
  }

  // Check for missing endorsements
  const missing = new Set<ServiceIdString>();
  for (const member of members) {
    if (!state.hasMember(member)) {
      missing.add(member);
    }
  }
  if (missing.size > 0) {
    return {
      valid: false,
      reason: `missing ${missing.size} endorsements (${logServiceIds(missing)})`,
    };
  }

  return { valid: true };
}

export async function maybeCreateGroupSendEndorsementState(
  groupId: string,
  alreadyRefreshedGroupState: boolean,
  alreadyInQueue = false
): Promise<MaybeCreateGroupSendEndorsementStateResult> {
  const logId = `maybeCreateGroupSendEndorsementState/groupv2(${groupId})`;

  const conversation = window.ConversationController.get(groupId);
  strictAssert(conversation != null, `${logId}: Conversation not found`);
  strictAssert(
    isGroupV2(conversation.attributes),
    `${logId}: Conversation is not groupV2`
  );

  async function getEndorsementsState() {
    strictAssert(conversation != null, `${logId}: Conversation not found`);

    const secretParams = conversation.get('secretParams');
    const membersV2 = conversation.get('membersV2');
    strictAssert(secretParams, `${logId}: Must have secret params`);
    strictAssert(membersV2, `${logId}: Must have members`);

    const members = new Set(membersV2.map(member => member.aci));

    const data = await DataReader.getGroupSendEndorsementsData(groupId);
    const state =
      data != null ? new GroupSendEndorsementState(data, secretParams) : null;

    return { members, state };
  }

  const { members, state } = await getEndorsementsState();

  // Check if we need to refresh the group state.
  const result = validateGroupSendEndorsements(members, state);
  if (!result.valid) {
    // For terminated groups, the server will not return endorsements so just move on
    if (conversation.get('terminated')) {
      log.info(`${logId}: Group terminated, ignoring invalid endorsements`);
      return { state: null, didRefreshGroupState: false };
    }

    // If we've already refreshed the group state, we should log and move on.
    if (alreadyRefreshedGroupState) {
      return { state: null, didRefreshGroupState: false };
    }
    if (conversation.isBlocked()) {
      return { state: null, didRefreshGroupState: false };
    }

    log.info(
      `${logId}: Endorsements invalid, refreshing group: ${result.reason}`
    );

    // Wait for all incoming messages to be processed before refreshing
    // the group state to avoid incorrectly processing messages.
    await window.waitForEmptyEventQueue();

    if (alreadyInQueue) {
      await maybeUpdateGroup({ conversation, force: true });
      return { state: null, didRefreshGroupState: true };
    }

    // Refresh the group state all allow the caller to try again.
    await conversation.queueJob(
      'maybeCreateGroupSendEndorsementState',
      async () => {
        // Check again; we might have updated since we queued this job
        const { members: newMembers, state: newState } =
          await getEndorsementsState();

        const newResult = validateGroupSendEndorsements(newMembers, newState);
        if (newResult.valid) {
          log.info(
            `${logId}: Endorsements are valid; returning without updating group`
          );
          return;
        }

        await maybeUpdateGroup({ conversation, force: true });
      }
    );

    return { state: null, didRefreshGroupState: true };
  }

  return { state, didRefreshGroupState: false };
}
