import type { SlotNumber } from '@aztec/foundation/branded-types';

import { z } from 'zod';

import { CheckpointAttestation } from '../p2p/checkpoint_attestation.js';
import { type ApiSchemaFor, optional, schemas } from '../schemas/index.js';
import { Tx } from '../tx/tx.js';
import { TxHash } from '../tx/tx_hash.js';
import { MAX_RPC_TXS_LEN } from './api_limit.js';

export type PeerInfo =
  | { status: 'connected'; score: number; id: string }
  | { status: 'dialing'; dialStatus: string; id: string; addresses: string[] }
  | { status: 'cached'; id: string; addresses: string[]; enr: string; dialAttempts: number };

const PeerInfoSchema = z.discriminatedUnion('status', [
  z.object({ status: z.literal('connected'), score: z.number(), id: z.string() }),
  z.object({ status: z.literal('dialing'), dialStatus: z.string(), id: z.string(), addresses: z.array(z.string()) }),
  z.object({
    status: z.literal('cached'),
    id: z.string(),
    addresses: z.array(z.string()),
    enr: z.string(),
    dialAttempts: z.number(),
  }),
]);

/** Exposed API to the P2P module. */
export interface P2PApi {
  /**
   * Returns all pending transactions in the transaction pool.
   * @param limit - The number of items to returns
   * @param after - The last known pending tx. Used for pagination
   * @returns An array of Txs.
   */
  getPendingTxs(limit?: number, after?: TxHash): Promise<Tx[]>;

  /** Returns the number of pending txs in the p2p tx pool. */
  getPendingTxCount(): Promise<number>;

  /**
   * Returns the ENR for this node, if any.
   */
  getEncodedEnr(): Promise<string | undefined>;

  /**
   * Returns info for all connected, dialing, and cached peers.
   */
  getPeers(includePending?: boolean): Promise<PeerInfo[]>;

  /**
   * Queries the Attestation pool for checkpoint attestations for the given slot
   *
   * @param slot - the slot to query
   * @param proposalId - the proposal id to query, or undefined to query all proposals for the slot
   * @returns CheckpointAttestations
   */
  getCheckpointAttestationsForSlot(slot: SlotNumber, proposalId?: string): Promise<CheckpointAttestation[]>;
}

export interface P2PClient extends P2PApi {
  /** Manually adds checkpoint attestations to the p2p client attestation pool. */
  addOwnCheckpointAttestations(attestations: CheckpointAttestation[]): Promise<void>;
}

export const P2PApiSchema: ApiSchemaFor<P2PApi> = {
  getCheckpointAttestationsForSlot: z
    .function()
    .args(schemas.SlotNumber, optional(z.string()))
    .returns(z.array(CheckpointAttestation.schema)),
  getPendingTxs: z
    .function()
    .args(optional(z.number().gte(1).lte(MAX_RPC_TXS_LEN).default(MAX_RPC_TXS_LEN)), optional(TxHash.schema))
    .returns(z.array(Tx.schema)),

  getPendingTxCount: z.function().returns(schemas.Integer),
  getEncodedEnr: z.function().returns(z.string().optional()),
  getPeers: z.function().args(optional(z.boolean())).returns(z.array(PeerInfoSchema)),
};