import { Hex, HexBigInt, HexNumber, StatusCode } from "./schemas.js";
import { SinglePolicyId } from "./capabilities/paymaster.js";
import { object, union, withMetadata, type SchemaOptions, z } from "./zod.js";

const BASE58_REGEX = /^[1-9A-HJ-NP-Za-km-z]+$/;

export const SolanaChainId = union({ description: "Solana chain identifier" }, [
  union({ description: "Solana Mainnet" }, [
    z.literal("solana:mainnet"),
    z.literal("solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp"),
  ]),
  union({ description: "Solana Devnet" }, [
    z.literal("solana:devnet"),
    z.literal("solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1"),
  ]),
]);

export type SolanaChainId = z.input<typeof SolanaChainId>;

export const SolanaAddress = (options?: SchemaOptions) => {
  const message =
    options?.errorMessage ??
    "Must be a valid base58-encoded Solana address (32-44 characters)";
  return withMetadata(
    z
      .string({ error: message })
      .regex(BASE58_REGEX, { error: message })
      .min(32, { error: message })
      .max(44, { error: message }),
    {
      description: "Base58-encoded Solana address",
      ...options,
      errorMessage: message,
    },
  );
};

export type SolanaAddress = z.input<ReturnType<typeof SolanaAddress>>;

// Raw instruction is the default call type (no `type` field needed).
// Typed helpers are optional sugar discriminated by `type` field.
export const SolanaInstruction = object(
  {
    programId: SolanaAddress(),
    accounts: z
      .array(
        object({
          pubkey: SolanaAddress(),
          isSigner: z.boolean(),
          isWritable: z.boolean(),
        }),
      )
      .optional(),
    data: Hex({ description: "Hex-encoded instruction data" }),
  },
  { description: "Raw Solana instruction" },
);

export type SolanaInstruction = z.input<typeof SolanaInstruction>;

export const SolanaCall = SolanaInstruction;

export type SolanaCall = SolanaInstruction;

export const SolanaSignatureRequest = object(
  {
    type: z.literal("solana_signTransaction"),
    data: Hex({ description: "Full serialized transaction, hex-encoded" }),
  },
  { description: "Solana transaction signature request" },
);

export type SolanaSignatureRequest = z.input<typeof SolanaSignatureRequest>;

export const SolanaSig = object({
  signature: object({
    type: z.literal("ed25519"),
    data: z.string().meta({
      description: "64-byte Ed25519 signature, base58-encoded",
    }),
  }),
});

export type SolanaSig = z.input<typeof SolanaSig>;

export const SolanaV0TransactionData = object(
  {
    compiledTransaction: Hex({
      description:
        "Pre-compiled wire-format transaction bytes from prepareCalls. Contains all non-user signatures (e.g. paymaster) already embedded.",
    }),
    signer: SolanaAddress({ description: "User's signing address" }),
    version: z.literal("0"),
    lifetimeConstraint: object({
      blockHash: z.string().regex(BASE58_REGEX),
      // Optional to allow for durable nonce support later w/o a breaking change.
      lastValidBlockHeight: HexBigInt().optional(),
    }),
  },
  {
    description: "Solana v0 versioned transaction data (send-relevant fields)",
  },
);

export type SolanaV0TransactionData = z.input<typeof SolanaV0TransactionData>;

export const SolanaFeePayment = object(
  {
    sponsored: z.boolean(),
    feePayer: SolanaAddress({
      description: "Address paying transaction fees",
    }),
    estimatedFeeLamports: z
      .number()
      .meta({ description: "Estimated transaction fee in lamports" })
      .optional(),
    estimatedRentLamports: z
      .number()
      .meta({ description: "Estimated rent in lamports" })
      .optional(),
    prefundLamports: z
      .number()
      .meta({ description: "Total lamports pre-funded by paymaster" })
      .optional(),
  },
  { description: "Solana fee payment details" },
);

export type SolanaFeePayment = z.input<typeof SolanaFeePayment>;

export const PreparedCall_SolanaV0 = object(
  {
    type: z
      .literal("solana-transaction-v0")
      .meta({ description: "Solana v0 versioned transaction" }),
    chainId: SolanaChainId,
    signatureRequest: SolanaSignatureRequest,
    data: SolanaV0TransactionData,
    feePayment: SolanaFeePayment,
    details: object({
      type: z.literal("solana-transaction-v0"),
      data: object({
        calls: z.array(SolanaCall),
        addressLookupTableAddresses: z.array(SolanaAddress()).optional(),
      }),
    }),
  },
  {
    description: "Prepared Solana v0 transaction",
  },
);

export type PreparedCall_SolanaV0 = z.input<typeof PreparedCall_SolanaV0>;

export const PreparedCall_SolanaV0_Signed = z
  .intersection(
    PreparedCall_SolanaV0.omit({
      signatureRequest: true,
      feePayment: true,
      details: true,
    }),
    SolanaSig,
  )
  .meta({
    description: "Prepared Solana v0 transaction with signature",
  });

export type PreparedCall_SolanaV0_Signed = z.input<
  typeof PreparedCall_SolanaV0_Signed
>;

export const SolanaCommitmentLevel = union(
  {
    description: "Solana commitment level: processed, confirmed, or finalized.",
  },
  [z.literal("processed"), z.literal("confirmed"), z.literal("finalized")],
);

export type SolanaCommitmentLevel = z.input<typeof SolanaCommitmentLevel>;

export const SolanaReceipt = object(
  {
    signature: z.string().regex(BASE58_REGEX).meta({
      description: "Base58 transaction signature",
    }),
    blockTime: HexNumber({ description: "Unix timestamp" }).optional(),
    slot: HexNumber({ description: "Slot number" }),
  },
  { description: "Solana transaction receipt" },
);

export type SolanaReceipt = z.input<typeof SolanaReceipt>;

export const SolanaPaymasterCapability = z.intersection(
  SinglePolicyId,
  object({
    webhookData: z
      .string()
      .meta({
        description:
          "Additional data for custom policy rules, such as proof of humanity.",
      })
      .optional(),
    prefundRent: z
      .boolean()
      .meta({
        description:
          "Opt-in: ask the paymaster to simulate the transaction and prefund any lamports needed for CPI-created accounts (e.g. SystemProgram.createAccount).",
      })
      .optional(),
  }),
);

export type SolanaPaymasterCapability = z.input<
  typeof SolanaPaymasterCapability
>;

export const SolanaPrepareCallsParams = object({
  calls: z.array(SolanaCall),
  from: SolanaAddress(),
  chainId: SolanaChainId,
  addressLookupTableAddresses: z.array(SolanaAddress()).optional(),
  capabilities: object({
    paymasterService: SolanaPaymasterCapability.optional(),
  }).optional(),
});

export type SolanaPrepareCallsParams = z.input<typeof SolanaPrepareCallsParams>;

export const SolanaGetCallsStatusResponse = object(
  {
    id: Hex(),
    chainId: SolanaChainId,
    atomic: z.boolean(),
    status: StatusCode,
    commitment: SolanaCommitmentLevel.optional(),
    receipts: z.array(SolanaReceipt).optional(),
  },
  {
    description: "Solana response",
  },
);

export type SolanaGetCallsStatusResponse = z.input<
  typeof SolanaGetCallsStatusResponse
>;