import { joinSignature, splitSignature } from 'ethers/lib/utils.js';

import { MerkleTreeHook__factory } from '@hyperlane-xyz/core';
import {
  ChainName,
  HyperlaneCore,
  IsmType,
  MerkleTreeHookConfig,
  MultisigIsmConfig,
  defaultMultisigConfigs,
  getValidatorFromStorageLocation,
} from '@hyperlane-xyz/sdk';
import {
  Address,
  Checkpoint,
  MerkleProof,
  S3CheckpointWithId,
  SignatureLike,
  WithAddress,
  assert,
  bytes32ToAddress,
  chunk,
  ensure0x,
  eqAddress,
  eqAddressEvm,
  fromHexString,
  mapAllSettled,
  rootLogger,
  strip0x,
  toHexString,
} from '@hyperlane-xyz/utils';

import {
  type MetadataBuilder,
  type MetadataContext,
  type MultisigMetadataBuildResult,
  type ValidatorInfo,
  ValidatorStatus,
} from './types.js';

interface MessageIdMultisigMetadata {
  type: typeof IsmType.MESSAGE_ID_MULTISIG;
  signatures: SignatureLike[];
  checkpoint: Omit<Checkpoint, 'mailbox_domain'>;
}

interface MerkleRootMultisigMetadata extends Omit<
  MessageIdMultisigMetadata,
  'type'
> {
  type: typeof IsmType.MERKLE_ROOT_MULTISIG;
  proof: MerkleProof;
}

const MerkleTreeInterface = MerkleTreeHook__factory.createInterface();

const SIGNATURE_LENGTH = 65;

export type MultisigMetadata =
  | MessageIdMultisigMetadata
  | MerkleRootMultisigMetadata;

type CheckpointValidator = Awaited<
  ReturnType<typeof getValidatorFromStorageLocation>
>;

export class MultisigMetadataBuilder implements MetadataBuilder {
  protected validatorCache: Record<
    ChainName,
    Record<string, CheckpointValidator | undefined>
  > = {};

  constructor(
    protected readonly core: HyperlaneCore,
    protected readonly logger = rootLogger.child({
      module: 'MultisigMetadataBuilder',
    }),
  ) {}

  /**
   * Get human-readable alias for a validator address from defaultMultisigConfigs
   */
  protected getValidatorAlias(
    validatorAddress: Address,
    chain: ChainName,
  ): string | undefined {
    const config = defaultMultisigConfigs[chain];
    if (!config) return undefined;
    const validator = config.validators.find((v) =>
      eqAddress(v.address, validatorAddress),
    );
    return validator?.alias;
  }

  protected getAnnouncedStorageLocations(
    originChain: ChainName,
    validators: string[],
  ): Promise<string[][]> {
    return this.core
      .getContracts(originChain)
      .validatorAnnounce.getAnnouncedStorageLocations(validators);
  }

  protected async checkpointValidators(
    originChain: ChainName,
    validators: string[],
  ): Promise<(CheckpointValidator | undefined)[]> {
    this.validatorCache[originChain] ??= {};
    const toFetch = validators.filter(
      (v) => !(v in this.validatorCache[originChain]),
    );

    if (toFetch.length > 0) {
      const storageLocations = await this.getAnnouncedStorageLocations(
        originChain,
        toFetch,
      );

      this.logger.debug({ storageLocations }, 'Fetched storage locations');

      // Use mapAllSettled to handle partial failures gracefully
      // Some validators may not have announced storage locations or may have invalid ones
      const { fulfilled, rejected } = await mapAllSettled(
        storageLocations,
        async (locations, index) => {
          const latestLocation = locations.slice(-1)[0];
          if (!latestLocation) {
            throw new Error(
              `No storage location announced for validator ${toFetch[index]}`,
            );
          }
          return getValidatorFromStorageLocation(latestLocation);
        },
      );

      // Log any failures for debugging
      rejected.forEach((error, index) => {
        this.logger.warn(
          { validator: toFetch[index], error: error.message },
          'Failed to initialize checkpoint validator',
        );
      });

      this.logger.debug(
        { fulfilled: fulfilled.size, rejected: rejected.size },
        'Fetched validators',
      );

      toFetch.forEach((validator, index) => {
        // Store undefined for failed validators so we don't retry them
        this.validatorCache[originChain][validator] = fulfilled.get(index);
      });
    }

    return validators.map((v) => this.validatorCache[originChain][v]);
  }

  /**
   * Shared helper to fetch checkpoints from validators.
   * Returns the raw results from mapAllSettled for both getS3Checkpoints and getValidatorInfos to use.
   */
  protected async fetchValidatorCheckpoints(
    validators: Address[],
    match: {
      origin: number;
      merkleTree: Address;
      messageId: string;
      index: number;
    },
  ): Promise<{
    originChain: ChainName;
    fulfilled: Map<number, S3CheckpointWithId | void>;
    rejected: Map<number, Error>;
  }> {
    this.logger.debug({ match, validators }, 'Fetching validator checkpoints');

    const originChain = this.core.multiProvider.getChainName(match.origin);
    const checkpointValidators = await this.checkpointValidators(
      originChain,
      validators,
    );

    const { fulfilled, rejected } = await mapAllSettled(
      validators,
      async (_, index) => {
        const checkpointValidator = checkpointValidators[index];
        if (!checkpointValidator) {
          throw new Error('No valid storage location for validator');
        }
        return checkpointValidator.getCheckpoint(match.index);
      },
    );

    // Log any errors from failed checkpoint fetches
    rejected.forEach((error, index) => {
      this.logger.warn(
        { validator: validators[index], error: error.message },
        'Failed to fetch checkpoint',
      );
    });

    this.logger.debug(
      { fulfilled: fulfilled.size, rejected: rejected.size },
      'Fetched validator checkpoints',
    );

    return { originChain, fulfilled, rejected };
  }

  /**
   * Helper to check if a checkpoint matches the expected values.
   */
  protected checkpointMatches(
    checkpoint: S3CheckpointWithId,
    match: {
      origin: number;
      merkleTree: Address;
      messageId: string;
      index: number;
    },
  ): boolean {
    return (
      eqAddress(
        bytes32ToAddress(checkpoint.value.checkpoint.merkle_tree_hook_address),
        match.merkleTree,
      ) &&
      checkpoint.value.message_id === match.messageId &&
      checkpoint.value.checkpoint.index === match.index &&
      checkpoint.value.checkpoint.mailbox_domain === match.origin
    );
  }

  async getS3Checkpoints(
    validators: Address[],
    match: {
      origin: number;
      merkleTree: Address;
      messageId: string;
      index: number;
    },
  ): Promise<S3CheckpointWithId[]> {
    const { fulfilled } = await this.fetchValidatorCheckpoints(
      validators,
      match,
    );

    // Filter to only valid, matching checkpoints
    const checkpoints = [...fulfilled.values()].filter(
      (value): value is S3CheckpointWithId => value !== undefined,
    );

    const matchingCheckpoints = checkpoints.filter((checkpoint) =>
      this.checkpointMatches(checkpoint, match),
    );

    if (matchingCheckpoints.length !== checkpoints.length) {
      this.logger.warn(
        {
          matchingCheckpoints: matchingCheckpoints.length,
          checkpoints: checkpoints.length,
          match,
        },
        'Mismatched checkpoints',
      );
    }

    return matchingCheckpoints;
  }

  /**
   * Get detailed status for each validator including signature status.
   * Returns ValidatorInfo for each validator with 'signed', 'pending', or 'error' status.
   */
  async getValidatorInfos(
    validators: Address[],
    match: {
      origin: number;
      merkleTree: Address;
      messageId: string;
      index: number;
    },
  ): Promise<{
    validatorInfos: ValidatorInfo[];
    checkpoint?: Checkpoint;
  }> {
    const { originChain, fulfilled, rejected } =
      await this.fetchValidatorCheckpoints(validators, match);

    let firstMatchingCheckpoint: Checkpoint | undefined;
    const validatorInfos: ValidatorInfo[] = validators.map((address, index) => {
      const alias = this.getValidatorAlias(address, originChain);

      // Check if this validator failed (no storage location or fetch failed)
      if (rejected.has(index)) {
        return {
          address,
          alias,
          status: ValidatorStatus.Error,
          error: `Failed to fetch checkpoint: ${rejected.get(index)?.message}`,
        };
      }

      const checkpoint = fulfilled.get(index);

      // Check if checkpoint doesn't exist (validator hasn't signed this index yet)
      if (!checkpoint) {
        return {
          address,
          alias,
          status: ValidatorStatus.Pending,
        };
      }

      // Check if checkpoint matches our expected values
      if (!this.checkpointMatches(checkpoint, match)) {
        return {
          address,
          alias,
          status: ValidatorStatus.Pending,
        };
      }

      // Valid matching checkpoint found
      if (!firstMatchingCheckpoint) {
        firstMatchingCheckpoint = checkpoint.value.checkpoint;
      }

      return {
        address,
        alias,
        status: ValidatorStatus.Signed,
        signature: checkpoint.signature,
        checkpointIndex: checkpoint.value.checkpoint.index,
      };
    });

    return {
      validatorInfos,
      checkpoint: firstMatchingCheckpoint,
    };
  }

  async build(
    context: MetadataContext<
      WithAddress<MultisigIsmConfig>,
      WithAddress<MerkleTreeHookConfig>
    >,
  ): Promise<MultisigMetadataBuildResult> {
    // Currently only MESSAGE_ID_MULTISIG is supported
    assert(
      context.ism.type === IsmType.MESSAGE_ID_MULTISIG ||
        context.ism.type === IsmType.STORAGE_MESSAGE_ID_MULTISIG,
      'Merkle proofs are not yet supported',
    );

    const merkleTree = context.hook.address;

    // Find the merkle tree insertion event for this message
    const matchingInsertion = context.dispatchTx.logs
      .filter((log) => eqAddressEvm(log.address, merkleTree))
      .map((log) => MerkleTreeInterface.parseLog(log))
      .find((event) => event.args.messageId === context.message.id);

    assert(
      matchingInsertion,
      `No merkle tree insertion of ${context.message.id} to ${merkleTree} found in dispatch tx`,
    );
    this.logger.debug({ matchingInsertion }, 'Found matching insertion event');

    const checkpointIndex = matchingInsertion.args.index;

    // Get detailed validator status
    const { validatorInfos, checkpoint } = await this.getValidatorInfos(
      context.ism.validators,
      {
        origin: context.message.parsed.origin,
        messageId: context.message.id,
        merkleTree,
        index: checkpointIndex,
      },
    );

    // Count signed validators
    const signedValidators = validatorInfos.filter(
      (v) => v.status === ValidatorStatus.Signed,
    );
    const signedCount = signedValidators.length;
    const quorumMet = signedCount >= context.ism.threshold;

    this.logger.debug(
      { signedCount, threshold: context.ism.threshold, quorumMet },
      `Validator signature status for message ${context.message.id}`,
    );

    // Build the result
    const result: MultisigMetadataBuildResult = {
      type: context.ism.type,
      ismAddress: context.ism.address,
      threshold: context.ism.threshold,
      validators: validatorInfos,
      checkpointIndex,
    };

    // Only encode metadata if quorum is met
    if (quorumMet && checkpoint) {
      const signatures = signedValidators
        .map((v) => v.signature!)
        .slice(0, context.ism.threshold);

      this.logger.debug(
        { signatures: signatures.length, ism: context.ism },
        `Taking ${signatures.length} (threshold) signatures for message ${context.message.id}`,
      );

      const metadata: MessageIdMultisigMetadata = {
        type: IsmType.MESSAGE_ID_MULTISIG,
        checkpoint,
        signatures,
      };

      result.metadata = MultisigMetadataBuilder.encode(metadata);
    } else {
      this.logger.debug(
        { signedCount, threshold: context.ism.threshold },
        `Quorum not met for message ${context.message.id}, metadata not buildable`,
      );
    }

    return result;
  }

  protected static encodeSimplePrefix(
    metadata: MessageIdMultisigMetadata,
  ): string {
    const checkpoint = metadata.checkpoint;
    const buf = Buffer.alloc(68);
    buf.write(strip0x(checkpoint.merkle_tree_hook_address), 0, 32, 'hex');
    buf.write(strip0x(checkpoint.root), 32, 32, 'hex');
    buf.writeUInt32BE(checkpoint.index, 64);
    return toHexString(buf);
  }

  static decodeSimplePrefix(metadata: string): {
    signatureOffset: number;
    type: IsmType;
    checkpoint: {
      root: string;
      index: number;
      merkle_tree_hook_address: string;
    };
  } {
    const buf = fromHexString(metadata);
    const merkleTree = toHexString(buf.subarray(0, 32));
    const root = toHexString(buf.subarray(32, 64));
    const index = buf.readUint32BE(64);
    const checkpoint = {
      root,
      index,
      merkle_tree_hook_address: merkleTree,
    };
    return {
      signatureOffset: 68,
      type: IsmType.MESSAGE_ID_MULTISIG,
      checkpoint,
    };
  }

  static encodeProofPrefix(metadata: MerkleRootMultisigMetadata): string {
    const checkpoint = metadata.checkpoint;
    const buf = Buffer.alloc(1096);
    buf.write(strip0x(checkpoint.merkle_tree_hook_address), 0, 32, 'hex');
    buf.writeUInt32BE(metadata.proof.index, 32);
    buf.write(strip0x(metadata.proof.leaf.toString()), 36, 32, 'hex');
    const branchEncoded = metadata.proof.branch
      .map((b) => strip0x(b.toString()))
      .join('');
    buf.write(branchEncoded, 68, 32 * 32, 'hex');
    buf.writeUint32BE(checkpoint.index, 1092);
    return toHexString(buf);
  }

  static decodeProofPrefix(metadata: string): {
    signatureOffset: number;
    type: IsmType;
    checkpoint: {
      root: string;
      index: number;
      merkle_tree_hook_address: string;
    };
    proof: MerkleProof;
  } {
    const buf = fromHexString(metadata);
    const merkleTree = toHexString(buf.subarray(0, 32));
    const messageIndex = buf.readUint32BE(32);
    const signedMessageId = toHexString(buf.subarray(36, 68));
    const branchEncoded = buf.subarray(68, 1092).toString('hex');
    const branch = chunk(branchEncoded, 32 * 2).map((v) => ensure0x(v));
    const signedIndex = buf.readUint32BE(1092);
    const checkpoint = {
      root: '',
      index: messageIndex,
      merkle_tree_hook_address: merkleTree,
    };
    const proof: MerkleProof = {
      branch,
      leaf: signedMessageId,
      index: signedIndex,
    };
    return {
      signatureOffset: 1096,
      type: IsmType.MERKLE_ROOT_MULTISIG,
      checkpoint,
      proof,
    };
  }

  static encode(metadata: MultisigMetadata): string {
    let encoded =
      metadata.type === IsmType.MESSAGE_ID_MULTISIG
        ? this.encodeSimplePrefix(metadata)
        : this.encodeProofPrefix(metadata);

    metadata.signatures.forEach((signature) => {
      const encodedSignature = joinSignature(signature);
      assert(
        fromHexString(encodedSignature).byteLength === SIGNATURE_LENGTH,
        'Invalid signature length',
      );
      encoded += strip0x(encodedSignature);
    });

    return encoded;
  }

  static signatureAt(
    metadata: string,
    offset: number,
    index: number,
  ): SignatureLike | undefined {
    const buf = fromHexString(metadata);
    const start = offset + index * SIGNATURE_LENGTH;
    const end = start + SIGNATURE_LENGTH;
    if (end > buf.byteLength) {
      return undefined;
    }

    return toHexString(buf.subarray(start, end));
  }

  static decode(
    metadata: string,
    type:
      | typeof IsmType.MERKLE_ROOT_MULTISIG
      | typeof IsmType.MESSAGE_ID_MULTISIG
      | typeof IsmType.STORAGE_MERKLE_ROOT_MULTISIG
      | typeof IsmType.STORAGE_MESSAGE_ID_MULTISIG,
  ): MultisigMetadata {
    const prefix: any =
      type === IsmType.MERKLE_ROOT_MULTISIG ||
      type === IsmType.STORAGE_MERKLE_ROOT_MULTISIG
        ? this.decodeProofPrefix(metadata)
        : this.decodeSimplePrefix(metadata);

    const { signatureOffset: offset, ...values } = prefix;

    const signatures: SignatureLike[] = [];
    for (let i = 0; this.signatureAt(metadata, offset, i); i++) {
      const { r, s, v } = splitSignature(
        this.signatureAt(metadata, offset, i)!,
      );
      signatures.push({ r, s, v });
    }

    return {
      signatures,
      ...values,
    };
  }
}