// Core dependencies
import { AnchorProvider, Program } from "@coral-xyz/anchor";
import NodeWallet from "@coral-xyz/anchor/dist/cjs/nodewallet";
import { getAssociatedTokenAddressSync } from "@solana/spl-token";
import { AccountInfo, AddressLookupTableProgram, Connection, GetProgramAccountsFilter, GetProgramAccountsResponse, Keypair, PublicKey } from "@solana/web3.js";
import { Metaplex } from "@metaplex-foundation/js";

// Local imports
import { BASKETS_PROGRAM_ID, DEPOSIT_FEE_SEED, MEME_SOL, MINT_SEED, PYTH_SPONSORED_FEEDS, RAYDIUM_CPMM, RAYDIUM_LIQUIDITY_POOL_V4, REBALANCE_STATE_SEED, STATE_CREATOR_SEED, USDC_DECIMALS, USDC_MINT, WITHDRAW_FEE_SEED, WSOL_DECIMALS, WSOL_MINT } from "./constants";
import { IDL } from "../idl/idl";
import { BasketsProgram } from "../idl/types";
import { parse, v4 } from "uuid";
import { fetchPythSponsoredFeeds } from "../state/oracle";
import { BASKETS_STATE_SIZE, parseBasketState, ParsedBasketState } from "../state/basket";
import { BasketState } from "../state/basket";
import { WithdrawState, WITHDRAW_STATE_SIZE, parseWithdrawState, ParsedWithdrawState } from "../state/withdrawState";
import { PoolInfo } from "../state/oracle";

export function getBasketsProgram(
    connection: Connection,
): Program<BasketsProgram> {
    const program: Program<BasketsProgram> = new Program(
        IDL,
        new AnchorProvider(
            connection,
            new NodeWallet(Keypair.generate())
        )
    );
    return program;
}

export function getBasketPda(
    basket: PublicKey,
): PublicKey {
    return PublicKey.findProgramAddressSync(
        [Uint8Array.from(basket.toBuffer())],
        BASKETS_PROGRAM_ID,
    )[0];
}

export function getRebalanceStateAccount(): PublicKey {
    return PublicKey.findProgramAddressSync(
        [REBALANCE_STATE_SEED],
        BASKETS_PROGRAM_ID,
    )[0];
}

export function getStateCreatorAccount(): PublicKey {
    return PublicKey.findProgramAddressSync(
        [STATE_CREATOR_SEED],
        BASKETS_PROGRAM_ID,
    )[0];
}

export function getDepositFeesWallet(
    basket: PublicKey,
): PublicKey {
    return PublicKey.findProgramAddressSync(
        [DEPOSIT_FEE_SEED, basket.toBuffer()],
        BASKETS_PROGRAM_ID,
    )[0];
}

export function getWithdrawFeesWallet(
    basket: PublicKey,
): PublicKey {
    return PublicKey.findProgramAddressSync(
        [WITHDRAW_FEE_SEED, basket.toBuffer()],
        BASKETS_PROGRAM_ID,
    )[0];
}

export function getWithdrawStateAccount(
    withdrawStateSeed: number[],
): PublicKey {
    return PublicKey.findProgramAddressSync(
        [Uint8Array.from(withdrawStateSeed)],
        BASKETS_PROGRAM_ID,
    )[0];
}

export function getBasketTokenMintAccount(
    basket: PublicKey,
): PublicKey {
    return PublicKey.findProgramAddressSync(
        [
            MINT_SEED,
            basket.toBuffer()
        ],
        BASKETS_PROGRAM_ID,
    )[0];
}

export function getMetadataAccount(
    tokenMint: PublicKey
): PublicKey {
    const metaplex = Metaplex.make(new Connection("https://api.devnet.solana.com"));
    return metaplex.nfts().pdas().metadata({ mint: tokenMint });
}

export function getLookupTableAccount(
    creator: PublicKey,
    slot: number,
): PublicKey {
    const ixAndPubkey = AddressLookupTableProgram.createLookupTable({
        authority: creator,
        recentSlot: slot,
        payer: creator,
    });
    return ixAndPubkey[1];
}

export function getDeactivatedLookupTableAccount(
    lookupTable: PublicKey,
): PublicKey {
    return PublicKey.findProgramAddressSync(
        [lookupTable.toBuffer()],
        BASKETS_PROGRAM_ID,
    )[0];
}

export function getAta(
    wallet: PublicKey,
    tokenMint: PublicKey,
): PublicKey {
    return getAssociatedTokenAddressSync(tokenMint, wallet, true);
}

export function getRandomSeed(): number[] {
    return Array.from(parse(v4()));
}

export async function getAccountInfos(
    connection: Connection,
    keys: PublicKey[],
): Promise<(AccountInfo<Buffer> | null)[]> {
    const allAccounts: (AccountInfo<Buffer>|null)[] = [];
    const batchSize = 100;
    for (let i = 0; i < keys.length; i += batchSize) {
        const batch = keys.slice(i, i + batchSize);
        const batchAccounts = await connection.getMultipleAccountsInfo(batch);
        allAccounts.push(...batchAccounts);
    }
    return allAccounts;
}


async function getRaydiumPools(
    connection: Connection,
    tokenMint: PublicKey,
    programId: PublicKey,
    dataSize: number,
    baseTokenOffset: number,
    quoteTokenOffset: number,
    baseMintOffset: number,
    quoteMintOffset: number,
): Promise<PoolInfo[]> {
    // Get all pools that pair tokenMint with either SOL or USDC
    const [solQuotePairs, solBasePairs, usdcQuotePairs, usdcBasePairs, memeQuotePairs, memeBasePairs] = await Promise.all([
        // Get pools where tokenMint is quote token and SOL is base
        connection.getProgramAccounts(programId, {
            filters: [
                { dataSize },
                { memcmp: { offset: baseMintOffset, bytes: WSOL_MINT.toBase58() } },
                { memcmp: { offset: quoteMintOffset, bytes: tokenMint.toBase58() } },
            ]
        }),
        // Get pools where tokenMint is base token and SOL is quote
        connection.getProgramAccounts(programId, {
            filters: [
                { dataSize },
                { memcmp: { offset: quoteMintOffset, bytes: WSOL_MINT.toBase58() } },
                { memcmp: { offset: baseMintOffset, bytes: tokenMint.toBase58() } },
            ]
        }),
        // Get pools where tokenMint is quote token and USDC is base
        connection.getProgramAccounts(programId, {
            filters: [
                { dataSize },
                { memcmp: { offset: baseMintOffset, bytes: USDC_MINT.toBase58() } },
                { memcmp: { offset: quoteMintOffset, bytes: tokenMint.toBase58() } },
            ]
        }),
        // Get pools where tokenMint is base token and USDC is quote
        connection.getProgramAccounts(programId, {
            filters: [
                { dataSize },
                { memcmp: { offset: quoteMintOffset, bytes: USDC_MINT.toBase58() } },
                { memcmp: { offset: baseMintOffset, bytes: tokenMint.toBase58() } },
            ]
        }),
        // Get pools where tokenMint is quote token and MEME SOL    is base
        connection.getProgramAccounts(programId, {
            filters: [
                { dataSize },
                { memcmp: { offset: baseMintOffset, bytes: MEME_SOL.toBase58() } },
                { memcmp: { offset: quoteMintOffset, bytes: tokenMint.toBase58() } },
            ]
        }),
        // Get pools where tokenMint is base token and MEME SOL is quote
        connection.getProgramAccounts(programId, {
            filters: [
                { dataSize },
                { memcmp: { offset: quoteMintOffset, bytes: MEME_SOL.toBase58() } },
                { memcmp: { offset: baseMintOffset, bytes: tokenMint.toBase58() } },
            ]
        }),
    ]);

    // Combine all pool accounts
    const allPoolAccounts = [
        ...solQuotePairs,
        ...solBasePairs,
        ...usdcQuotePairs,
        ...usdcBasePairs,
        ...memeQuotePairs,
        ...memeBasePairs,
    ];

    // Process pool accounts
    const pools = allPoolAccounts
        .filter(account => {
            // For V4 pools, only include active pools (status === 6)
            if (programId === RAYDIUM_LIQUIDITY_POOL_V4) {
                const status = parseInt(account.account.data.readBigUInt64LE(0).toString());
                return status === 6;
            }
            return true;
        })
        .map(account => {
            // Extract pool data
            const poolData = {
                liquidity: 0,
                poolType: 0,
                pool: account.pubkey.toBase58(),
                baseMint: new PublicKey(account.account.data.slice(baseMintOffset, baseMintOffset + 32)).toBase58(),
                quoteMint: new PublicKey(account.account.data.slice(quoteMintOffset, quoteMintOffset + 32)).toBase58(),
                baseTokenAccount: new PublicKey(account.account.data.slice(baseTokenOffset, baseTokenOffset + 32)).toBase58(),
                quoteTokenAccount: new PublicKey(account.account.data.slice(quoteTokenOffset, quoteTokenOffset + 32)).toBase58(),
                baseBalance: 0,
                quoteBalance: 0,
                baseDecimals: 0,
                quoteDecimals: 0,
            };

            // Ensure tokenMint is always base token
            if (poolData.baseMint !== tokenMint.toBase58()) {
                [poolData.quoteMint, poolData.baseMint] = [poolData.baseMint, poolData.quoteMint];
                [poolData.quoteTokenAccount, poolData.baseTokenAccount] = [poolData.baseTokenAccount, poolData.quoteTokenAccount];
            }

            return poolData;
        });

    // Collect token accounts for balance lookup
    const tokenAccountsToLookup = [
        ...pools.map(pool => new PublicKey(pool.quoteTokenAccount)),
        ...pools.map(pool => new PublicKey(pool.baseTokenAccount)),
        ...pools.map(pool => new PublicKey(pool.quoteMint)),
        ...pools.map(pool => new PublicKey(pool.baseMint))
    ];

    // Get account info for all token and mint accounts
    const accountInfos = await getAccountInfos(connection, tokenAccountsToLookup);

    // Update pool balances and decimals
    pools.forEach((pool, i) => {
        pool.quoteBalance = parseInt(accountInfos[i]?.data.readBigUInt64LE(64).toString() ?? "0");
        pool.baseBalance = parseInt(accountInfos[pools.length + i]?.data.readBigUInt64LE(64).toString() ?? "0");
        pool.quoteDecimals = accountInfos[2 * pools.length + i]?.data[44] ?? 0;
        pool.baseDecimals = accountInfos[3 * pools.length + i]?.data[44] ?? 0;
    });

    return pools;
}

export async function getRaydiumV4Pools(
    connection: Connection,
    tokenMint: PublicKey,
    usdcPrice: number,
    solPrice: number,
): Promise<PoolInfo[]> {
    const pools = await getRaydiumPools(
        connection,
        tokenMint,
        RAYDIUM_LIQUIDITY_POOL_V4,
        752,
        336,
        368,
        400,
        432
    );
    pools.forEach(pool => {
        pool.poolType = 2;
        const price = (pool.quoteMint == WSOL_MINT.toBase58() || pool.quoteMint == MEME_SOL.toBase58()) ? solPrice : usdcPrice;
        const decimals = (pool.quoteMint == WSOL_MINT.toBase58() || pool.quoteMint == MEME_SOL.toBase58()) ? WSOL_DECIMALS : USDC_DECIMALS;
        pool.liquidity = price * pool.quoteBalance / 10 ** decimals * 2;
    });
    pools.sort((a, b) => b.liquidity - a.liquidity);
    return pools.slice(0, 2);
}

export async function getRaydiumCpmmPools(
    connection: Connection,
    tokenMint: PublicKey,
    usdcPrice: number,
    solPrice: number,
): Promise<PoolInfo[]> {
    const pools = await getRaydiumPools(
        connection,
        tokenMint,
        RAYDIUM_CPMM,
        637,
        72,
        104,
        168,
        200
    );
    pools.forEach(pool => {
        pool.poolType = 1;
        const price = (pool.quoteMint == WSOL_MINT.toBase58() || pool.quoteMint == MEME_SOL.toBase58()) ? solPrice : usdcPrice;
        const decimals = (pool.quoteMint == WSOL_MINT.toBase58() || pool.quoteMint == MEME_SOL.toBase58()) ? WSOL_DECIMALS : USDC_DECIMALS;
        pool.liquidity = price * pool.quoteBalance / 10 ** decimals * 2;
    });
    pools.sort((a, b) => b.liquidity - a.liquidity);
    return pools.slice(0, 2);
    return pools;
}

export async function getPythSponsoredFeeds(
    program: Program<BasketsProgram>,
    tokenMint: PublicKey,
): Promise<PoolInfo[]> {
    const pythSponsoredFeeds = await fetchPythSponsoredFeeds(program, PYTH_SPONSORED_FEEDS);
    const pools = [];
    for (let i = 0; i < pythSponsoredFeeds.numTokens; i++)
        if (pythSponsoredFeeds.mints[i].toBase58() == tokenMint.toBase58()) {
            pools.push({
                liquidity: 999999999,
                poolType: 0,
                pool: pythSponsoredFeeds.ownAddress.toBase58(),
                baseMint: pythSponsoredFeeds.mints[i].toBase58(),
                quoteMint: USDC_MINT.toBase58(),
                baseTokenAccount: pythSponsoredFeeds.feeds[i].toBase58(),
                quoteTokenAccount: PublicKey.default.toBase58(),
                baseBalance: 0,
                quoteBalance: 0,
                baseDecimals: (await getAccountInfos(program.provider.connection, [tokenMint]))[0]?.data[44] ?? 0,
                quoteDecimals: 6,
            });
        }
    return pools;
}

export async function getAllBaskets(
    program: Program<BasketsProgram>,
): Promise<ParsedBasketState[]> {
    const accounts: BasketState[] = (await program.account.basketV200.all()).map(account => account.account);
    return accounts.map(account => parseBasketState(account));
}

export async function getBasketsByCreator(
    program: Program<BasketsProgram>,
    creator: PublicKey,
): Promise<ParsedBasketState[]> {
    const accountFilters: GetProgramAccountsFilter[] = [
        {
            dataSize: BASKETS_STATE_SIZE + 8,
        },
        {
            memcmp: {
                offset: 8 + 1 + 32 + 1 + 32 + 32 + 8 + 8 + 8 + 8,
                bytes: creator.toBase58(),
            },
        }
    ]
    const accounts: GetProgramAccountsResponse = await program.provider.connection
        .getProgramAccounts(
            BASKETS_PROGRAM_ID,
            {
                commitment: "confirmed",
                filters: accountFilters,
                encoding: 'base64'
            }
        );
    const baskets: BasketState[] = accounts.map(account => 
        program.coder.accounts.decode("basketV200", account.account.data)
    );
    return await Promise.all(baskets.map(basket => parseBasketState(basket)));
}

export async function getWithdrawStatesByUser(
    program: Program<BasketsProgram>,
    user: PublicKey,
): Promise<ParsedWithdrawState[]> {
    const accountFilters: GetProgramAccountsFilter[] = [
        {
            dataSize: WITHDRAW_STATE_SIZE + 8,
        },
        {
            memcmp: {
                offset: 8 + 32 + 16 + 32 ,
                bytes: user.toBase58(),
            },
        }
    ]
    const accounts: GetProgramAccountsResponse = await program.provider.connection
        .getProgramAccounts(
            BASKETS_PROGRAM_ID,
            {
                commitment: "confirmed",
                filters: accountFilters,
                encoding: 'base64'
            }
        );
    const withdrawStates: WithdrawState[] = accounts.map(account => 
        program.coder.accounts.decode("withdrawStateV200", account.account.data)
    );
    return withdrawStates.map(withdrawState => parseWithdrawState(withdrawState));
}