import { PublicKey } from "@solana/web3.js";
import { ASSOCIATED_TOKEN_PROGRAM_ID, TOKEN_PROGRAM_ID } from "@solana/spl-token";
import { PROGRAM_ID } from "../constants.js";

/**
 * Derive the market state PDA for a given week number
 *
 * @param weekNumber - Week number (unique identifier for the market, must be valid u32)
 * @returns Market state PDA public key
 * @throws Error if weekNumber is invalid
 *
 * @example
 * ```typescript
 * const marketPda = getMarketAddress(1);
 * const accountInfo = await connection.getAccountInfo(marketPda);
 * ```
 */
export function getMarketAddress(weekNumber: number): PublicKey {
  if (!Number.isInteger(weekNumber) || weekNumber < 0 || weekNumber > 0xffffffff) {
    throw new Error("weekNumber must be a valid u32 (0 to 4294967295)");
  }

  const weekBytes = Buffer.alloc(4);
  weekBytes.writeUInt32LE(weekNumber);

  const [pda] = PublicKey.findProgramAddressSync([Buffer.from("market"), weekBytes], PROGRAM_ID);

  return pda;
}

/**
 * Derive the market state PDA with bump
 *
 * @param weekNumber - Week number (must be valid u32)
 * @returns Tuple of [PDA, bump]
 * @throws Error if weekNumber is invalid
 */
export function getMarketAddressWithBump(weekNumber: number): [PublicKey, number] {
  if (!Number.isInteger(weekNumber) || weekNumber < 0 || weekNumber > 0xffffffff) {
    throw new Error("weekNumber must be a valid u32 (0 to 4294967295)");
  }

  const weekBytes = Buffer.alloc(4);
  weekBytes.writeUInt32LE(weekNumber);

  return PublicKey.findProgramAddressSync([Buffer.from("market"), weekBytes], PROGRAM_ID);
}

/**
 * Derive the wSOL vault ATA for a market
 *
 * The vault is an Associated Token Account owned by the market PDA.
 *
 * @param marketPda - Market state PDA
 * @param wsolMint - wSOL mint address
 * @returns wSOL vault ATA public key
 */
export function getWsolVaultAddress(marketPda: PublicKey, wsolMint: PublicKey): PublicKey {
  const [ata] = PublicKey.findProgramAddressSync(
    [marketPda.toBuffer(), TOKEN_PROGRAM_ID.toBuffer(), wsolMint.toBuffer()],
    ASSOCIATED_TOKEN_PROGRAM_ID
  );

  return ata;
}

/**
 * Derive a user's ATA for a specific token mint
 *
 * @param owner - User's wallet public key
 * @param mint - Token mint address
 * @param programId - Token program ID (default: TOKEN_PROGRAM_ID, use TOKEN_2022_PROGRAM_ID for Hit/Miss tokens)
 * @returns User's ATA public key
 */
export function getUserTokenAddress(
  owner: PublicKey,
  mint: PublicKey,
  programId: PublicKey = TOKEN_PROGRAM_ID
): PublicKey {
  const [ata] = PublicKey.findProgramAddressSync(
    [owner.toBuffer(), programId.toBuffer(), mint.toBuffer()],
    ASSOCIATED_TOKEN_PROGRAM_ID
  );

  return ata;
}

/**
 * Derive the registry PDA that tracks all active market weeks
 *
 * The registry is a singleton account that stores an array of active week numbers.
 * It's updated when markets are initialized (add week) and when both pools are settled (remove week).
 *
 * @returns Registry state PDA public key
 *
 * @example
 * ```typescript
 * const registryPda = getRegistryAddress();
 * const accountInfo = await connection.getAccountInfo(registryPda);
 * ```
 */
export function getRegistryAddress(): PublicKey {
  const [pda] = PublicKey.findProgramAddressSync([Buffer.from("registry")], PROGRAM_ID);

  return pda;
}

/**
 * Derive the registry PDA with bump
 *
 * @returns Tuple of [PDA, bump]
 */
export function getRegistryAddressWithBump(): [PublicKey, number] {
  return PublicKey.findProgramAddressSync([Buffer.from("registry")], PROGRAM_ID);
}

/**
 * Derive a mint PDA for a specific market token
 *
 * Seeds: ["mint", week_number (4 bytes LE), pool_type (1 byte), strike_index (1 byte), is_hit (1 byte)]
 *
 * @param weekNumber - Market week number (must be valid u32)
 * @param poolType - 0 for High pool, 1 for Low pool
 * @param strikeIndex - Strike index (0, 1, or 2)
 * @param isHit - true for HIT token, false for MISS token
 * @returns Mint PDA public key
 * @throws Error if parameters are invalid
 *
 * @example
 * ```typescript
 * // Get High pool HIT token mint at strike 0 for week 1
 * const mintPda = getMintAddress(1, 0, 0, true);
 *
 * // Get Low pool MISS token mint at strike 2 for week 1
 * const mintPda = getMintAddress(1, 1, 2, false);
 * ```
 */
export function getMintAddress(
  weekNumber: number,
  poolType: number,
  strikeIndex: number,
  isHit: boolean
): PublicKey {
  if (!Number.isInteger(weekNumber) || weekNumber < 0 || weekNumber > 0xffffffff) {
    throw new Error("weekNumber must be a valid u32 (0 to 4294967295)");
  }
  if (poolType !== 0 && poolType !== 1) {
    throw new Error("poolType must be 0 (High) or 1 (Low)");
  }
  if (strikeIndex < 0 || strikeIndex > 2) {
    throw new Error("strikeIndex must be 0, 1, or 2");
  }

  const weekBytes = Buffer.alloc(4);
  weekBytes.writeUInt32LE(weekNumber);

  const [pda] = PublicKey.findProgramAddressSync(
    [
      Buffer.from("mint"),
      weekBytes,
      Buffer.from([poolType]),
      Buffer.from([strikeIndex]),
      Buffer.from([isHit ? 1 : 0]),
    ],
    PROGRAM_ID
  );

  return pda;
}

/**
 * Derive a mint PDA with bump
 *
 * @param weekNumber - Market week number
 * @param poolType - 0 for High, 1 for Low
 * @param strikeIndex - Strike index (0, 1, 2)
 * @param isHit - true for HIT, false for MISS
 * @returns Tuple of [PDA, bump]
 */
export function getMintAddressWithBump(
  weekNumber: number,
  poolType: number,
  strikeIndex: number,
  isHit: boolean
): [PublicKey, number] {
  if (!Number.isInteger(weekNumber) || weekNumber < 0 || weekNumber > 0xffffffff) {
    throw new Error("weekNumber must be a valid u32 (0 to 4294967295)");
  }
  if (poolType !== 0 && poolType !== 1) {
    throw new Error("poolType must be 0 (High) or 1 (Low)");
  }
  if (strikeIndex < 0 || strikeIndex > 2) {
    throw new Error("strikeIndex must be 0, 1, or 2");
  }

  const weekBytes = Buffer.alloc(4);
  weekBytes.writeUInt32LE(weekNumber);

  return PublicKey.findProgramAddressSync(
    [
      Buffer.from("mint"),
      weekBytes,
      Buffer.from([poolType]),
      Buffer.from([strikeIndex]),
      Buffer.from([isHit ? 1 : 0]),
    ],
    PROGRAM_ID
  );
}

/**
 * Get all 12 mint PDAs for a market in the order expected by initialize_market
 *
 * Order:
 * - [0-2]: High pool HIT mints (strikes 0, 1, 2)
 * - [3-5]: High pool MISS mints (strikes 0, 1, 2)
 * - [6-8]: Low pool HIT mints (strikes 0, 1, 2)
 * - [9-11]: Low pool MISS mints (strikes 0, 1, 2)
 *
 * @param weekNumber - Market week number
 * @returns Array of 12 mint PDAs
 */
export function getAllMintAddresses(weekNumber: number): PublicKey[] {
  const mints: PublicKey[] = [];

  // High pool HIT mints (strikes 0, 1, 2)
  for (let i = 0; i < 3; i++) {
    mints.push(getMintAddress(weekNumber, 0, i, true));
  }
  // High pool MISS mints (strikes 0, 1, 2)
  for (let i = 0; i < 3; i++) {
    mints.push(getMintAddress(weekNumber, 0, i, false));
  }
  // Low pool HIT mints (strikes 0, 1, 2)
  for (let i = 0; i < 3; i++) {
    mints.push(getMintAddress(weekNumber, 1, i, true));
  }
  // Low pool MISS mints (strikes 0, 1, 2)
  for (let i = 0; i < 3; i++) {
    mints.push(getMintAddress(weekNumber, 1, i, false));
  }

  return mints;
}