import {
  Address,
  address,
  AccountRole,
  getAddressDecoder,
  getAddressEncoder,
  getProgramDerivedAddress,
  Instruction,
  IAccountMeta,
  Rpc,
  SolanaRpcApi,
} from '@solana/kit';
import { TOKEN_PROGRAM_ADDRESS } from '@solana-program/token';

// ── Program IDs ──────────────────────────────────────────────────────
export const VAULT_MINT_PROGRAM_ID: Address = address('9WUyNREiPDMgwMh5Gt81Fd3JpiCKxpjZ5Dpq9Bo1RhMV');
export const VAULT_STAKE_PROGRAM_ID: Address = address('97V7JsExNC6yFWu5KjK1FLfVkNVvtMpAFL5QkLWKEGxY');

// ── Virtual share constants (must match on-chain) ────────────────────
const VIRTUAL_SHARES = 1_000_000n;
const VIRTUAL_ASSETS = 1_000_000n;

// ── Anchor discriminator for "deposit" instruction ───────────────────
// sha256("global:deposit")[0..8] - same for both programs since the method name is "deposit"
const DEPOSIT_DISCRIMINATOR = new Uint8Array([242, 35, 198, 137, 82, 225, 242, 182]);

// ── Types ────────────────────────────────────────────────────────────

export interface VaultMintConfig {
  vault: Address; // USDC mint
  mint: Address; // wYLDS mint
  vaultAuthority: Address;
  redeemVault: Address;
  bump: number;
  paused: boolean;
  allowedExternalMintProgram: Address;
}

export interface StakeConfig {
  vault: Address; // wYLDS mint
  mint: Address; // PRIME mint
  unbondingPeriod: bigint;
  bump: number;
  paused: boolean;
}

export interface StakeVaultTokenAccountConfig {
  vaultTokenAccount: Address;
  vaultAuthority: Address;
  bump: number;
}

export interface VaultTokenAccountConfig {
  vaultTokenAccount: Address;
  bump: number;
}

export interface HastraQuote {
  inputMint: Address;
  outputMint: Address;
  inputAmount: bigint;
  outputAmount: bigint;
  /** For USDC->wYLDS this is always 1:1. For wYLDS->PRIME it uses virtual share math. */
  exchangeRate: number;
}

// ── PDA derivation helpers ───────────────────────────────────────────

async function findPda(seeds: (Uint8Array | string)[], programId: Address): Promise<Address> {
  const seedBytes = seeds.map((s) => (typeof s === 'string' ? new TextEncoder().encode(s) : s));
  const [pda] = await getProgramDerivedAddress({ seeds: seedBytes, programAddress: programId });
  return pda;
}

async function findVaultMintConfigPda(): Promise<Address> {
  return findPda(['config'], VAULT_MINT_PROGRAM_ID);
}

async function findVaultMintMintAuthorityPda(): Promise<Address> {
  return findPda(['mint_authority'], VAULT_MINT_PROGRAM_ID);
}

async function findVaultTokenAccountConfigPda(configPda: Address): Promise<Address> {
  return findPda(
    ['vault_token_account_config', new Uint8Array(getAddressEncoder().encode(configPda))],
    VAULT_MINT_PROGRAM_ID,
  );
}

async function findStakeConfigPda(): Promise<Address> {
  return findPda(['stake_config'], VAULT_STAKE_PROGRAM_ID);
}

async function findStakeMintAuthorityPda(): Promise<Address> {
  return findPda(['mint_authority'], VAULT_STAKE_PROGRAM_ID);
}

async function findStakeVaultAuthorityPda(): Promise<Address> {
  return findPda(['vault_authority'], VAULT_STAKE_PROGRAM_ID);
}

async function findStakeVaultTokenAccountConfigPda(stakeConfigPda: Address): Promise<Address> {
  return findPda(
    ['stake_vault_token_account_config', new Uint8Array(getAddressEncoder().encode(stakeConfigPda))],
    VAULT_STAKE_PROGRAM_ID,
  );
}

// ── Account deserialization ──────────────────────────────────────────

function readPubkey(data: Uint8Array, offset: number): Address {
  return getAddressDecoder().decode(data.slice(offset, offset + 32));
}

function readU8(data: Uint8Array, offset: number): number {
  return data[offset];
}

function readBool(data: Uint8Array, offset: number): boolean {
  return data[offset] !== 0;
}

function readI64(data: Uint8Array, offset: number): bigint {
  const view = new DataView(data.buffer, data.byteOffset + offset, 8);
  return view.getBigInt64(0, true);
}

function readU64(data: Uint8Array, offset: number): bigint {
  const view = new DataView(data.buffer, data.byteOffset + offset, 8);
  return view.getBigUint64(0, true);
}

function readVecPubkeys(data: Uint8Array, offset: number): { keys: Address[]; bytesRead: number } {
  const view = new DataView(data.buffer, data.byteOffset + offset, 4);
  const len = view.getUint32(0, true);
  if (offset + 4 + len * 32 > data.length) {
    throw new Error(
      `readVecPubkeys: expected ${len} pubkeys (${4 + len * 32} bytes) at offset ${offset}, but data is only ${data.length} bytes`,
    );
  }
  const keys: Address[] = [];
  let cursor = offset + 4;
  for (let i = 0; i < len; i++) {
    keys.push(readPubkey(data, cursor));
    cursor += 32;
  }
  return { keys, bytesRead: cursor - offset };
}

function deserializeVaultMintConfig(data: Uint8Array): VaultMintConfig {
  let offset = 8; // skip discriminator
  const vault = readPubkey(data, offset);
  offset += 32;
  const mint = readPubkey(data, offset);
  offset += 32;
  const { bytesRead: freezeBytes } = readVecPubkeys(data, offset);
  offset += freezeBytes;
  const { bytesRead: rewardsBytes } = readVecPubkeys(data, offset);
  offset += rewardsBytes;
  const vaultAuthority = readPubkey(data, offset);
  offset += 32;
  const redeemVault = readPubkey(data, offset);
  offset += 32;
  const bump = readU8(data, offset);
  offset += 1;
  const paused = readBool(data, offset);
  offset += 1;
  const allowedExternalMintProgram = readPubkey(data, offset);

  return { vault, mint, vaultAuthority, redeemVault, bump, paused, allowedExternalMintProgram };
}

function deserializeStakeConfig(data: Uint8Array): StakeConfig {
  let offset = 8;
  const vault = readPubkey(data, offset);
  offset += 32;
  const mint = readPubkey(data, offset);
  offset += 32;
  const unbondingPeriod = readI64(data, offset);
  offset += 8;
  const { bytesRead: freezeBytes } = readVecPubkeys(data, offset);
  offset += freezeBytes;
  const { bytesRead: rewardsBytes } = readVecPubkeys(data, offset);
  offset += rewardsBytes;
  const bump = readU8(data, offset);
  offset += 1;
  const paused = readBool(data, offset);

  return { vault, mint, unbondingPeriod, bump, paused };
}

function deserializeVaultTokenAccountConfig(data: Uint8Array): VaultTokenAccountConfig {
  let offset = 8;
  const vaultTokenAccount = readPubkey(data, offset);
  offset += 32;
  const bump = readU8(data, offset);
  return { vaultTokenAccount, bump };
}

function deserializeStakeVaultTokenAccountConfig(data: Uint8Array): StakeVaultTokenAccountConfig {
  let offset = 8;
  const vaultTokenAccount = readPubkey(data, offset);
  offset += 32;
  const vaultAuthority = readPubkey(data, offset);
  offset += 32;
  const bump = readU8(data, offset);
  return { vaultTokenAccount, vaultAuthority, bump };
}

// ── Token account helpers ────────────────────────────────────────────

interface TokenAccountInfo {
  mint: Address;
  owner: Address;
  amount: bigint;
}

function deserializeTokenAccount(data: Uint8Array): TokenAccountInfo {
  return {
    mint: readPubkey(data, 0),
    owner: readPubkey(data, 32),
    amount: readU64(data, 64),
  };
}

interface MintInfo {
  supply: bigint;
  decimals: number;
}

function deserializeMint(data: Uint8Array): MintInfo {
  // SPL Mint layout: 4 bytes mintAuthorityOption, 32 bytes mintAuthority, 8 bytes supply, 1 byte decimals
  const supply = readU64(data, 36);
  const decimals = readU8(data, 44);
  return { supply, decimals };
}

// ── Instruction encoding ─────────────────────────────────────────────

function encodeU64LE(value: bigint): Uint8Array {
  const buf = new Uint8Array(8);
  const view = new DataView(buf.buffer);
  view.setBigUint64(0, value, true);
  return buf;
}

function buildInstruction(
  programId: Address,
  keys: IAccountMeta[],
  discriminator: Uint8Array,
  amount: bigint,
): Instruction {
  const amountBytes = encodeU64LE(amount);
  const data = new Uint8Array(discriminator.length + amountBytes.length);
  data.set(discriminator, 0);
  data.set(amountBytes, discriminator.length);

  return {
    programAddress: programId,
    accounts: keys,
    data,
  };
}

// ── Main client ──────────────────────────────────────────────────────

/** Config cache TTL in seconds — refetch on-chain configs after this period */
const CONFIG_TTL_SECONDS = 5 * 60; // 5 minutes

export class HastraVaultClient {
  private rpc: Rpc<SolanaRpcApi>;

  private vaultMintConfig: VaultMintConfig | null = null;
  private vaultMintConfigFetchedAt = 0;
  private stakeConfig: StakeConfig | null = null;
  private stakeConfigFetchedAt = 0;
  private vaultTokenAccountConfig: VaultTokenAccountConfig | null = null;
  private stakeVaultTokenAccountConfig: StakeVaultTokenAccountConfig | null = null;

  private vaultMintConfigPda: Address | null = null;
  private vaultMintMintAuthorityPda: Address | null = null;
  private vaultTokenAccountConfigPda: Address | null = null;
  private stakeConfigPda: Address | null = null;
  private stakeMintAuthorityPda: Address | null = null;
  private stakeVaultAuthorityPda: Address | null = null;
  private stakeVaultTokenAccountConfigPda: Address | null = null;

  constructor(rpc: Rpc<SolanaRpcApi>) {
    this.rpc = rpc;
  }

  // ── Lazy PDA loaders ───────────────────────────────────────────────

  private async getVaultMintConfigPda(): Promise<Address> {
    if (!this.vaultMintConfigPda) {
      this.vaultMintConfigPda = await findVaultMintConfigPda();
    }
    return this.vaultMintConfigPda;
  }

  private async getVaultMintMintAuthorityPda(): Promise<Address> {
    if (!this.vaultMintMintAuthorityPda) {
      this.vaultMintMintAuthorityPda = await findVaultMintMintAuthorityPda();
    }
    return this.vaultMintMintAuthorityPda;
  }

  private async getVaultTokenAccountConfigPda(): Promise<Address> {
    if (!this.vaultTokenAccountConfigPda) {
      const configPda = await this.getVaultMintConfigPda();
      this.vaultTokenAccountConfigPda = await findVaultTokenAccountConfigPda(configPda);
    }
    return this.vaultTokenAccountConfigPda;
  }

  private async getStakeConfigPda(): Promise<Address> {
    if (!this.stakeConfigPda) {
      this.stakeConfigPda = await findStakeConfigPda();
    }
    return this.stakeConfigPda;
  }

  private async getStakeMintAuthorityPda(): Promise<Address> {
    if (!this.stakeMintAuthorityPda) {
      this.stakeMintAuthorityPda = await findStakeMintAuthorityPda();
    }
    return this.stakeMintAuthorityPda;
  }

  private async getStakeVaultAuthorityPda(): Promise<Address> {
    if (!this.stakeVaultAuthorityPda) {
      this.stakeVaultAuthorityPda = await findStakeVaultAuthorityPda();
    }
    return this.stakeVaultAuthorityPda;
  }

  private async getStakeVaultTokenAccountConfigPda(): Promise<Address> {
    if (!this.stakeVaultTokenAccountConfigPda) {
      const stakeConfigPda = await this.getStakeConfigPda();
      this.stakeVaultTokenAccountConfigPda = await findStakeVaultTokenAccountConfigPda(stakeConfigPda);
    }
    return this.stakeVaultTokenAccountConfigPda;
  }

  // ── Account fetching ───────────────────────────────────────────────

  private async fetchAccountData(addr: Address): Promise<Uint8Array> {
    const resp = await this.rpc.getAccountInfo(addr, { encoding: 'base64' }).send();
    if (!resp.value) {
      throw new Error(`Account not found: ${addr}`);
    }
    const b64 = resp.value.data[0];
    return new Uint8Array(Buffer.from(b64, 'base64'));
  }

  async getVaultMintConfig(): Promise<VaultMintConfig> {
    const now = Date.now() / 1000;
    if (!this.vaultMintConfig || now - this.vaultMintConfigFetchedAt > CONFIG_TTL_SECONDS) {
      const configPda = await this.getVaultMintConfigPda();
      const data = await this.fetchAccountData(configPda);
      this.vaultMintConfig = deserializeVaultMintConfig(data);
      this.vaultMintConfigFetchedAt = now;
    }
    return this.vaultMintConfig;
  }

  async getStakeConfig(): Promise<StakeConfig> {
    const now = Date.now() / 1000;
    if (!this.stakeConfig || now - this.stakeConfigFetchedAt > CONFIG_TTL_SECONDS) {
      const stakeConfigPda = await this.getStakeConfigPda();
      const data = await this.fetchAccountData(stakeConfigPda);
      this.stakeConfig = deserializeStakeConfig(data);
      this.stakeConfigFetchedAt = now;
    }
    return this.stakeConfig;
  }

  async getVaultTokenAccountConfig(): Promise<VaultTokenAccountConfig> {
    if (!this.vaultTokenAccountConfig) {
      const pda = await this.getVaultTokenAccountConfigPda();
      const data = await this.fetchAccountData(pda);
      this.vaultTokenAccountConfig = deserializeVaultTokenAccountConfig(data);
    }
    return this.vaultTokenAccountConfig;
  }

  async getStakeVaultTokenAccountConfig(): Promise<StakeVaultTokenAccountConfig> {
    if (!this.stakeVaultTokenAccountConfig) {
      const pda = await this.getStakeVaultTokenAccountConfigPda();
      const data = await this.fetchAccountData(pda);
      this.stakeVaultTokenAccountConfig = deserializeStakeVaultTokenAccountConfig(data);
    }
    return this.stakeVaultTokenAccountConfig;
  }

  private async getTokenAccountInfo(addr: Address): Promise<TokenAccountInfo> {
    const data = await this.fetchAccountData(addr);
    return deserializeTokenAccount(data);
  }

  private async getMintInfo(addr: Address): Promise<MintInfo> {
    const data = await this.fetchAccountData(addr);
    return deserializeMint(data);
  }

  // ── Quote: USDC -> wYLDS ──────────────────────────────────────────

  /**
   * Quote for minting wYLDS from USDC via vault-mint.
   * The exchange rate is always 1:1 (deposit amount = minted amount).
   * Returns outputAmount 0 if the vault is paused.
   */
  async quoteUsdcToWylds(usdcAmount: bigint): Promise<HastraQuote> {
    const config = await this.getVaultMintConfig();
    if (config.paused || usdcAmount === 0n) {
      return {
        inputMint: config.vault,
        outputMint: config.mint,
        inputAmount: usdcAmount,
        outputAmount: 0n,
        exchangeRate: 0,
      };
    }
    return {
      inputMint: config.vault,
      outputMint: config.mint,
      inputAmount: usdcAmount,
      outputAmount: usdcAmount, // 1:1
      exchangeRate: 1.0,
    };
  }

  // ── Quote: wYLDS -> PRIME ─────────────────────────────────────────

  /**
   * Quote for staking wYLDS to receive PRIME via vault-stake.
   * Uses virtual share accounting to calculate the output.
   * Returns outputAmount 0 if the vault is paused.
   */
  async quoteWyldsToPrime(wyldsAmount: bigint): Promise<HastraQuote> {
    const stakeConfig = await this.getStakeConfig();
    if (stakeConfig.paused) {
      console.debug('HastraVaultClient: vault-stake is paused, wYLDS -> PRIME staking is disabled');
      return {
        inputMint: stakeConfig.vault,
        outputMint: stakeConfig.mint,
        inputAmount: wyldsAmount,
        outputAmount: 0n,
        exchangeRate: 0,
      };
    }

    const stakeVaultTaConfig = await this.getStakeVaultTokenAccountConfig();

    const [vaultAccountInfo, primeMintInfo] = await Promise.all([
      this.getTokenAccountInfo(stakeVaultTaConfig.vaultTokenAccount),
      this.getMintInfo(stakeConfig.mint),
    ]);

    const vaultBalance = vaultAccountInfo.amount;
    const totalShares = primeMintInfo.supply;

    const sharesOut = calculateAssetsToShares(wyldsAmount, totalShares, vaultBalance);

    const exchangeRate =
      totalShares === 0n ? 1.0 : Number(vaultBalance + VIRTUAL_ASSETS) / Number(totalShares + VIRTUAL_SHARES);

    return {
      inputMint: stakeConfig.vault,
      outputMint: stakeConfig.mint,
      inputAmount: wyldsAmount,
      outputAmount: sharesOut,
      exchangeRate,
    };
  }

  // ── Instructions: USDC -> wYLDS ────────────────────────────────────

  /**
   * Build instructions to deposit USDC and mint wYLDS (1:1).
   *
   * @param user - The user's wallet address (signer)
   * @param userUsdcTokenAccount - The user's USDC token account
   * @param userWyldsTokenAccount - The user's wYLDS token account
   * @param amount - Amount of USDC to deposit (in smallest units, e.g. 1_000_000 = 1 USDC)
   */
  async getDepositUsdcIxs(
    user: Address,
    userUsdcTokenAccount: Address,
    userWyldsTokenAccount: Address,
    amount: bigint,
  ): Promise<Instruction[]> {
    const [configPda, vtaConfigPda, vtaConfig, config, mintAuthorityPda] = await Promise.all([
      this.getVaultMintConfigPda(),
      this.getVaultTokenAccountConfigPda(),
      this.getVaultTokenAccountConfig(),
      this.getVaultMintConfig(),
      this.getVaultMintMintAuthorityPda(),
    ]);

    const keys: IAccountMeta[] = [
      { address: configPda, role: AccountRole.READONLY },
      { address: vtaConfigPda, role: AccountRole.READONLY },
      { address: vtaConfig.vaultTokenAccount, role: AccountRole.WRITABLE },
      { address: config.mint, role: AccountRole.WRITABLE },
      { address: mintAuthorityPda, role: AccountRole.READONLY },
      { address: user, role: AccountRole.READONLY_SIGNER },
      { address: userUsdcTokenAccount, role: AccountRole.WRITABLE },
      { address: userWyldsTokenAccount, role: AccountRole.WRITABLE },
      { address: TOKEN_PROGRAM_ADDRESS, role: AccountRole.READONLY },
    ];

    return [buildInstruction(VAULT_MINT_PROGRAM_ID, keys, DEPOSIT_DISCRIMINATOR, amount)];
  }

  // ── Instructions: wYLDS -> PRIME ───────────────────────────────────

  /**
   * Build instructions to stake wYLDS and receive PRIME.
   *
   * @param user - The user's wallet address (signer)
   * @param userWyldsTokenAccount - The user's wYLDS token account
   * @param userPrimeTokenAccount - The user's PRIME token account
   * @param amount - Amount of wYLDS to stake (in smallest units)
   */
  async getStakeWyldsIxs(
    user: Address,
    userWyldsTokenAccount: Address,
    userPrimeTokenAccount: Address,
    amount: bigint,
  ): Promise<Instruction[]> {
    const [stakeConfigPda, stakeConfig, svtaConfigPda, svtaConfig, mintAuthorityPda, vaultAuthorityPda] =
      await Promise.all([
        this.getStakeConfigPda(),
        this.getStakeConfig(),
        this.getStakeVaultTokenAccountConfigPda(),
        this.getStakeVaultTokenAccountConfig(),
        this.getStakeMintAuthorityPda(),
        this.getStakeVaultAuthorityPda(),
      ]);

    const keys: IAccountMeta[] = [
      { address: stakeConfigPda, role: AccountRole.READONLY },
      { address: svtaConfigPda, role: AccountRole.READONLY },
      { address: svtaConfig.vaultTokenAccount, role: AccountRole.WRITABLE },
      { address: vaultAuthorityPda, role: AccountRole.READONLY },
      { address: stakeConfig.mint, role: AccountRole.WRITABLE },
      { address: stakeConfig.vault, role: AccountRole.WRITABLE },
      { address: mintAuthorityPda, role: AccountRole.READONLY },
      { address: user, role: AccountRole.READONLY_SIGNER },
      { address: userWyldsTokenAccount, role: AccountRole.WRITABLE },
      { address: userPrimeTokenAccount, role: AccountRole.WRITABLE },
      { address: TOKEN_PROGRAM_ADDRESS, role: AccountRole.READONLY },
    ];

    return [buildInstruction(VAULT_STAKE_PROGRAM_ID, keys, DEPOSIT_DISCRIMINATOR, amount)];
  }
}

// ── Pure math (matches on-chain virtual share logic) ─────────────────

export function calculateAssetsToShares(assets: bigint, totalShares: bigint, vaultBalance: bigint): bigint {
  if (totalShares === 0n) {
    return (assets * VIRTUAL_SHARES) / VIRTUAL_ASSETS;
  }
  return (assets * (totalShares + VIRTUAL_SHARES)) / (vaultBalance + VIRTUAL_ASSETS);
}

export function calculateSharesToAssets(shares: bigint, totalShares: bigint, vaultBalance: bigint): bigint {
  if (totalShares === 0n) {
    return 0n;
  }
  return (shares * (vaultBalance + VIRTUAL_ASSETS)) / (totalShares + VIRTUAL_SHARES);
}

export function calculateExchangeRate(totalShares: bigint, vaultBalance: bigint): bigint {
  const SCALE = 1_000_000_000n;
  if (totalShares === 0n) {
    return SCALE;
  }
  return ((vaultBalance + VIRTUAL_ASSETS) * SCALE) / (totalShares + VIRTUAL_SHARES);
}