import { CreatedPdaAccount, Pubkey, PubkeyUtil } from '@saturnbtcio/arch-sdk';
import { Token, tokenToBytes } from '@saturnbtcio/pool-serde-sdk';
import { IArchProvider } from '../providers/arch.provider';
import { PdaType } from './pda-type';

type ProtocolPdaAccountData = {
  pdaType: PdaType.Protocol;
  programId: Pubkey;
};

interface ConfigPdaAccountData {
  pdaType: PdaType.Config;
  feeTier: number;
  token0: Token;
  token1: Token;
  programId: Pubkey;
}

interface ShardPdaAccountData {
  pdaType: PdaType.Shard;
  shardIdx: number;
  feeTier: number;
  token0: Token;
  token1: Token;
  programId: Pubkey;
}

interface PositionPdaAccountData {
  pdaType: PdaType.Position;
  userPubKey: Pubkey;
  feeTier: number;
  token0: Token;
  token1: Token;
  programId: Pubkey;
}

type CreatePdaAccountData =
  | ProtocolPdaAccountData
  | ConfigPdaAccountData
  | ShardPdaAccountData
  | PositionPdaAccountData;

const getSeeds = (
  pdaType: PdaType,
  feeTier: number,
  sortedTokens: Array<Buffer>,
): Array<Buffer> => {
  return [
    Buffer.from('saturn', 'utf8'),
    Buffer.from(pdaType, 'utf-8'),
    Buffer.from(new Uint16Array([feeTier]).buffer),
    sortedTokens[0],
    sortedTokens[1],
  ];
};

const getSortedTokens = (token0: Token, token1: Token): [Buffer, Buffer] => {
  const token0Bytes = tokenToBytes(token0);
  const token1Bytes = tokenToBytes(token1);
  return Buffer.compare(token0Bytes, token1Bytes) <= 0
    ? [token0Bytes, token1Bytes]
    : [token1Bytes, token0Bytes];
};

export const createProtocolPda = (programId: Pubkey): [Pubkey, number] => {
  const seeds = [
    Buffer.from('saturn', 'utf-8'),
    Buffer.from(PdaType.Protocol, 'utf-8'),
  ];

  return PubkeyUtil.findProgramAddress(seeds, programId);
};

const createShardPda = (
  shardIdx: number,
  feeTier: number,
  token0: Token,
  token1: Token,
  programId: Pubkey,
): [Pubkey, number] => {
  const sortedTokens = getSortedTokens(token0, token1);
  const seeds = getSeeds(PdaType.Shard, feeTier, sortedTokens);
  const shardIdxBuffer = Buffer.alloc(2);
  shardIdxBuffer.writeUInt16LE(shardIdx, 0);
  seeds.push(shardIdxBuffer);
  return PubkeyUtil.findProgramAddress(seeds, programId);
};

const createConfigPda = (
  feeTier: number,
  token0: Token,
  token1: Token,
  programId: Pubkey,
): [Pubkey, number] => {
  const sortedTokens = getSortedTokens(token0, token1);
  const seeds = getSeeds(PdaType.Config, feeTier, sortedTokens);
  return PubkeyUtil.findProgramAddress(seeds, programId);
};

const createPositionPda = (
  userPubKey: Pubkey,
  feeTier: number,
  token0: Token,
  token1: Token,
  programId: Pubkey,
): [Pubkey, number] => {
  const sortedTokens = getSortedTokens(token0, token1);
  const seeds = getSeeds(PdaType.Position, feeTier, sortedTokens);
  seeds.push(Buffer.from(userPubKey));
  return PubkeyUtil.findProgramAddress(seeds, programId);
};

const toCreatedPdaAccount = async (
  pdaAccount: [Pubkey, number],
  rpcService: IArchProvider,
): Promise<CreatedPdaAccount> => {
  const [publicKey, _] = pdaAccount;
  return {
    pubkey: PubkeyUtil.toHex(publicKey),
    address: await rpcService.getAccountAddress(publicKey),
  };
};

export const createNewAccount = async (
  pdaAccountData: CreatePdaAccountData,
  rpcService: IArchProvider,
): Promise<CreatedPdaAccount> => {
  let pdaAccount: [Pubkey, number];
  switch (pdaAccountData.pdaType) {
    case PdaType.Protocol: {
      pdaAccount = createProtocolPda(pdaAccountData.programId);
      break;
    }
    case PdaType.Config: {
      pdaAccount = createConfigPda(
        pdaAccountData.feeTier,
        pdaAccountData.token0,
        pdaAccountData.token1,
        pdaAccountData.programId,
      );
      break;
    }
    case PdaType.Shard: {
      pdaAccount = createShardPda(
        pdaAccountData.shardIdx,
        pdaAccountData.feeTier,
        pdaAccountData.token0,
        pdaAccountData.token1,
        pdaAccountData.programId,
      );
      break;
    }
    case PdaType.Position: {
      pdaAccount = createPositionPda(
        pdaAccountData.userPubKey,
        pdaAccountData.feeTier,
        pdaAccountData.token0,
        pdaAccountData.token1,
        pdaAccountData.programId,
      );
      break;
    }
    default: {
      // Exhaustive check
      const _exhaustive: never = pdaAccountData;
      throw new Error(`Invalid PDA type: ${(pdaAccountData as any).pdaType}`);
    }
  }

  return await toCreatedPdaAccount(pdaAccount, rpcService);
};