import { AccountInfo, Connection, GetProgramAccountsFilter, GetProgramAccountsResponse, PublicKey } from '@solana/web3.js';
import { BASKETS_V3_PROGRAM_ID, MAX_MANAGERS_PER_BASKET, MAX_SUPPORTED_TOKENS_PER_BASKET } from '../constants';
import { FormattedWithdrawBasketFees, WithdrawBasketFees, WithdrawBasketFeesLayout } from '../layouts/basket';
import { getMultipleAccountsInfoBatched } from '../txUtils';

export function addFieldsToWithdrawBasketFees(withdrawBasketFees: WithdrawBasketFees): WithdrawBasketFees {
    // Filter out default/empty owners
    let numOwners = MAX_MANAGERS_PER_BASKET;
    while (numOwners > 0 && withdrawBasketFees.owners[numOwners - 1].equals(PublicKey.default)) {
        numOwners--;
    }
    withdrawBasketFees.owners = withdrawBasketFees.owners.slice(0, numOwners);
    withdrawBasketFees.ownersWeights = withdrawBasketFees.ownersWeights.slice(0, numOwners);

    // Filter out default/empty tokens
    let numTokens = MAX_SUPPORTED_TOKENS_PER_BASKET;
    while (numTokens > 0 && withdrawBasketFees.accumulatedTokens[numTokens - 1].equals(PublicKey.default)) {
        numTokens--;
    }
    withdrawBasketFees.accumulatedTokens = withdrawBasketFees.accumulatedTokens.slice(0, numTokens);
    withdrawBasketFees.accumulatedFees = withdrawBasketFees.accumulatedFees.slice(0, numTokens);

    let formatted: FormattedWithdrawBasketFees = {
        pubkey: withdrawBasketFees.ownAddress!.toBase58(),
        basket: withdrawBasketFees.basket.toBase58(),
        rent_payer: withdrawBasketFees.rentPayer.toBase58(),
        owners: withdrawBasketFees.owners.map(owner => owner.toBase58()),
        owners_weights: withdrawBasketFees.ownersWeights,
        accumulated_tokens: withdrawBasketFees.accumulatedTokens.map(token => token.toBase58()),
        accumulated_fees: withdrawBasketFees.accumulatedFees.map(fee => parseInt(fee.toString())),
    };
    withdrawBasketFees.formatted = formatted;
    return withdrawBasketFees;
}

export async function fetchWithdrawBasketFees(
    connection: Connection,
    withdrawBasketFeesAddress: PublicKey,
): Promise<WithdrawBasketFees> {
    const accountInfo = await connection.getAccountInfo(withdrawBasketFeesAddress);
    if (!accountInfo) throw new Error("WithdrawBasketFees account not found");
    let withdrawBasketFees = WithdrawBasketFeesLayout.decode(accountInfo.data.slice(8));
    withdrawBasketFees.ownAddress = withdrawBasketFeesAddress;
    withdrawBasketFees = addFieldsToWithdrawBasketFees(withdrawBasketFees);
    return withdrawBasketFees;
}

export async function fetchWithdrawBasketFeesMultiple(
    connection: Connection,
    withdrawBasketFeesAddresses: PublicKey[],
): Promise<Map<string, WithdrawBasketFees>> {
    let multipleAccountsInfo = await getMultipleAccountsInfoBatched(connection, withdrawBasketFeesAddresses);
    let withdrawBasketFeesList: WithdrawBasketFees[] = withdrawBasketFeesAddresses.map(address => {
        let ai = multipleAccountsInfo.get(address.toBase58());
        if (!ai) return null;
        let withdrawBasketFees = { ...WithdrawBasketFeesLayout.decode(ai.data.slice(8)), ownAddress: address };
        withdrawBasketFees = addFieldsToWithdrawBasketFees(withdrawBasketFees);
        return withdrawBasketFees;
    }).filter((item): item is WithdrawBasketFees => item !== null);
    
    let withdrawBasketFeesMap: Map<string, WithdrawBasketFees> = new Map();
    for (let withdrawBasketFees of withdrawBasketFeesList) {
        withdrawBasketFeesMap.set(withdrawBasketFees.ownAddress!.toBase58(), withdrawBasketFees);
    }
    return withdrawBasketFeesMap;
}

export interface WithdrawBasketFeesFilter {
    type: "basket" | "manager" | "creator" | "host" | "symmetry";
    pubkey: string;
}

export async function fetchWithdrawBasketFeesList(
    connection: Connection,
    filter?: WithdrawBasketFeesFilter,
): Promise<WithdrawBasketFees[]> {

    const MANAGERS_ARRAY_OFFSET = 8 + 32 + 32;
    const MANAGER_PUBKEY_SIZE = 32;

    if (filter?.type === "manager") {
        try {
            const managerPubkey = new PublicKey(filter.pubkey);
            const managerBytes = Buffer.from(managerPubkey.toBytes());
            const managerSliceLength = MAX_MANAGERS_PER_BASKET * MANAGER_PUBKEY_SIZE;

            // Request only manager bytes for each basket (single GPA), then filter locally.
            const managerSlices = await connection.getProgramAccounts(BASKETS_V3_PROGRAM_ID, {
                commitment: "confirmed",
                filters: [{ dataSize: 8 + WithdrawBasketFeesLayout.getSpan() }],
                dataSlice: {
                    offset: MANAGERS_ARRAY_OFFSET,
                    length: managerSliceLength,
                },
                encoding: 'base64',
            });

            const matchedPubkeys: PublicKey[] = [];
            for (const basketAccount of managerSlices) {
                const managersData = basketAccount.account.data;
                for (let i = 0; i < MAX_MANAGERS_PER_BASKET; i++) {
                    const start = i * MANAGER_PUBKEY_SIZE;
                    const end = start + MANAGER_PUBKEY_SIZE;
                    if (managersData.subarray(start, end).equals(managerBytes)) {
                        matchedPubkeys.push(basketAccount.pubkey);
                        break;
                    }
                }
            }

            if (matchedPubkeys.length === 0) return [];
            const matchedAccountsInfo = await getMultipleAccountsInfoBatched(connection, matchedPubkeys);

            let withdrawBasketFeesList: WithdrawBasketFees[] = [];
            for (let pubkey of matchedPubkeys) {
                const accountInfo = matchedAccountsInfo.get(pubkey.toBase58());
                if (!accountInfo) continue;
                let withdrawBasketFees = WithdrawBasketFeesLayout.decode(accountInfo.data.slice(8));
                withdrawBasketFees.ownAddress = pubkey;
                withdrawBasketFees = addFieldsToWithdrawBasketFees(withdrawBasketFees);
                withdrawBasketFeesList.push(withdrawBasketFees);
            }
            return withdrawBasketFeesList;
        } catch (error) {
            // Fallback for RPC providers that might not support/serve this path reliably.
        }
    }

    let filters: GetProgramAccountsFilter[] = [{ dataSize: 8 + WithdrawBasketFeesLayout.getSpan() }];

    if (filter?.type === "creator" || filter?.type === "host" || filter?.type === "symmetry")
        filters.push({
            memcmp: {
                offset: 8 + 32 + 32,
                bytes: filter.pubkey
            }
        });

    if (filter?.type === "basket")
        filters.push({
            memcmp: {
                offset: 8,
                bytes: filter.pubkey
            }
        });

    let results = await connection.getProgramAccounts(BASKETS_V3_PROGRAM_ID, {
        commitment: "confirmed",
        filters: filters,
        encoding: 'base64'
    })

    let withdrawBasketFeesList = results.map(account => {
        let withdrawBasketFees = WithdrawBasketFeesLayout.decode(account.account.data.slice(8));
        withdrawBasketFees.ownAddress = account.pubkey;
        withdrawBasketFees = addFieldsToWithdrawBasketFees(withdrawBasketFees);
        return withdrawBasketFees;
    });
    return withdrawBasketFeesList;
}