import { PublicKey, TransactionSignature, AddressLookupTableAccount, TransactionInstruction, ComputeBudgetProgram, VersionedTransaction, TransactionMessage, Connection } from '@solana/web3.js';
import { Basket } from './basketState';
import { BasketError, TokenSettings, ADDITIONAL_UNITS, TransactionToSend } from './config';
import { buildUpdateCurrentWeightsIx, buildWithdrawBeforeRebalanceIx, buildDepositAfterRebalanceIx } from './instructionsBuilder';
import { generateJupTxData, getAddressLookupTableAccounts, signVersionedTransactions, sendSignedTransactions, getOraclePrices, delay } from './utils';
import { BN, Program } from '@coral-xyz/anchor';
import { BasketsIDL } from './basketsIDL';
import { getAssociatedTokenAddressSync } from './splTokenHelpers';


// Updated type definitions
type RebalanceInfo = { over: TokenRebalanceInfo[], under: TokenRebalanceInfo[] };
type TokenRebalanceInfo = { value: number, token: number, amount: number, iT?: boolean, rR?: boolean, uM: number };
type RebalanceAmounts = { from: number, to: number, tokenAmount: number, value: number };
type JupiterSwapData = {
    addressLookupTableAddresses: string[];
    swapInstruction: any;
    setupInstructions: any[];
    res: { inAmount: number };
    tokenAmount: number;
    swapValue: number,
};

/**
 * Fetches the current timestamp from the confirmed slot on the Solana blockchain.
 * This is used to determine when rebalances should occur based on basket settings.
 * @param connection The Solana connection
 * @returns The current timestamp or a default value
 */
export async function getConfirmedTimestamp(connection: Connection, basket: Basket): Promise<number> {
    const blockTime = await connection.getBlockTime(await connection.getSlot("confirmed")).catch((e) => null);
    return (blockTime !== null && basket.data.sellState.toNumber() == 0) ? blockTime : 2000000000;
}

/**
 * Fetches the lookup table account from the Solana blockchain.
 * Lookup tables are used to optimize transaction size and cost for complex operations like rebalancing.
 * @param connection The Solana connection
 * @returns The lookup table account
 */
export async function getLookupTableAccount(connection: Connection, lookupTableAccount: PublicKey): Promise<AddressLookupTableAccount> {
    const result = await connection.getAddressLookupTable(lookupTableAccount);
    //@ts-ignore
    return result.value;
}

/**
 * Determines if force rebalance is needed for actively managed baskets.
 * This allows basket managers to trigger rebalances regardless of other conditions.
 * @param basket The basket to check
 * @param wallet The wallet to use
 * @returns True if force rebalance is needed, false otherwise
 */
export function isForceRebalanceNeeded(basket: Basket, walletPublicKey: PublicKey): boolean {
    if (basket.data.sellState.toNumber() == 1 && basket.data.rebalanceSellState.toNumber() == 2)
        return true;
    return walletPublicKey.equals(basket.data.manager) && basket.data.activelyManaged.toNumber() === 1;
}

/**
 * Gets and sorts the rebalance info for a basket.
 * This function calculates which tokens are over or under their target weights,
 * considering the basket's rebalance threshold and current market prices.
 * @param basket The basket to rebalance
 * @param oraclePriceData Oracle price data for accurate token valuation
 * @param timestamp Current timestamp
 * @param tokenList The token list
 * @returns Sorted rebalance info
 */
export function getSortedRebalanceInfo(basket: Basket, oraclePriceData: number[], timestamp: number, tokenList: any): RebalanceInfo {
    const rebalanceInfos = getFlashRebalanceInfo(basket, tokenList, oraclePriceData, timestamp);
    rebalanceInfos.over.sort((a, b) => b.value - a.value);
    rebalanceInfos.under.sort((a, b) => b.value - a.value);
    return rebalanceInfos;
}

/**
 * Builds rebalance transactions for a basket.
 * This function creates transactions to adjust token weights to match their target weights,
 * considering factors like rebalance threshold, interval, and slippage settings.
 * @param basket The basket to rebalance
 * @param rebalanceInfos Information about tokens that need rebalancing
 * @param oraclePriceData Current oracle price data for tokens
 * @param forceRebalance Whether to force rebalance regardless of conditions
 * @param lookupTableAccount The lookup table account for the transaction
 * @param wallet The wallet to use for the transaction
 * @param connection The Solana connection
 * @param program The program instance
 * @param tokenList List of tokens in the basket
 * @param lamports Amount of lamports to use
 * @returns An array of transactions to send to the Solana blockchain
 */
export async function buildRebalanceTransactions(
    basket: Basket,
    rebalanceInfos: RebalanceInfo,
    oraclePriceData: number[],
    forceRebalance: boolean,
    lookups: AddressLookupTableAccount[],
    maxAllowedAccounts: number,
    walletPublicKey: PublicKey,
    connection: Connection,
    program: Program<BasketsIDL>,
    tokenList: TokenSettings[],
    lamports: number,
    updateOraclesTxData: TransactionToSend[],
    softCap: number,
    hardCap: number,
    underTokens: number,
    overTokens: number,
    jupAPIkey: string,
): Promise<TransactionToSend[]> {
    let txsToSend: TransactionToSend[] = updateOraclesTxData;
    let size = txsToSend.length;

    if (basket.data.sellState.toNumber() != 0)
        { overTokens = rebalanceInfos.over.length; underTokens = rebalanceInfos.under.length }
    // Iterate over overweighted and underweighted tokens
    for (const over of rebalanceInfos.over.slice(0, overTokens)) {
        for (const under of rebalanceInfos.under.slice(0, underTokens)) {
            // Check if we should process this rebalance based on thresholds and intervals
            if (!shouldProcessRebalance(over, under, forceRebalance)) continue;

            // Calculate the amounts to rebalance
            const { from, to, tokenAmount, value } = calculateRebalanceAmounts(
                over,
                under,
                oraclePriceData,
                tokenList,
                hardCap,
            );
            if (value <= softCap && !forceRebalance && basket.data.sellState.toNumber() == 0) continue;

            // Get Jupiter swap data for the rebalance
            const jupData = await getJupiterSwapData(
                from, to, tokenAmount, maxAllowedAccounts,
                basket, oraclePriceData, walletPublicKey, tokenList, jupAPIkey
            );
            if (!jupData) continue;

            if (jupData.swapValue <= softCap && !forceRebalance && basket.data.sellState.toNumber() == 0) continue;

            // Build the rebalance transaction
            const tx = await buildRebalanceTransaction(
                basket, from, to, jupData.tokenAmount, jupData,
                lookups, walletPublicKey, connection,
                program, tokenList, lamports,
            );
            txsToSend.push(tx);

            // Update rebalance info values
            over.value -= jupData.swapValue;
            under.value -= jupData.swapValue;
        }
    }

    if (txsToSend.length == size) return [];
    return txsToSend;
}

/**
 * Determines if a rebalance should be processed.
 * This function checks if the rebalance is necessary based on the basket's settings and current state.
 * @param over Over-weighted token info
 * @param under Under-weighted token info
 * @param forceRebalance Whether to force rebalance
 * @returns True if rebalance should be processed, false otherwise
 */
export function shouldProcessRebalance(over: TokenRebalanceInfo, under: TokenRebalanceInfo, forceRebalance: boolean): boolean {
    if (over.value <= 0 || under.value <= 0) return false;
    if (!forceRebalance && (over.iT && under.iT)) return false; // Both tokens are within threshold
    if (!forceRebalance && (over.rR && under.rR)) return false; // Both tokens were recently rebalanced
    return true;
}

/**
 * Calculates rebalance amounts.
 * This function determines how much of each token should be swapped to bring them closer to their target weights.
 * @param over Over-weighted token info
 * @param under Under-weighted token info
 * @param oraclePriceData Oracle price data
 * @param tokenList The token list
 * @returns Calculated rebalance amounts
 */
export function calculateRebalanceAmounts(
    over: TokenRebalanceInfo,
    under: TokenRebalanceInfo,
    oraclePriceData: number[],
    tokenList: any,
    hardCap: number
): RebalanceAmounts {
    const from = over.token;
    const to = under.token;
    const maxAmountUn = Math.min(hardCap, under.value) / oraclePriceData[from] * 10 ** tokenList[from].decimals;
    const maxAmountOv = Math.min(hardCap, over.value) / oraclePriceData[from] * 10 ** tokenList[from].decimals;
    let tokenAmount = Math.min(Math.floor(Math.min(maxAmountUn, maxAmountOv) * 0.995), over.amount);
    if (maxAmountUn > maxAmountOv && over.uM !== 0) tokenAmount = over.uM;
    const value = tokenAmount * oraclePriceData[from] / 10 ** tokenList[from].decimals;
    return { from, to, tokenAmount, value };
}

/**
 * Gets Jupiter swap data.
 * This function prepares the data needed for a token swap using Jupiter DEX aggregator.
 * @param from From token
 * @param to To token
 * @param tokenAmount Token amount
 * @param basket Basket
 * @param oraclePriceData Oracle price data
 * @param wallet The wallet to use
 * @param tokenList The token list
 * @returns Jupiter swap data or null if failed
 */
export async function getJupiterSwapData(
    from: number,
    to: number,
    tokenAmount: number,
    maxAllowedAccounts: number,
    basket: Basket,
    oraclePriceData: number[],
    walletPublicKey: PublicKey,
    tokenList: any,
    jupAPIkey: string,
): Promise<JupiterSwapData | null> {
    const data = await generateJupTxData(
        walletPublicKey,
        tokenList[from].tokenMint,
        tokenList[to].tokenMint,
        tokenAmount,
        maxAllowedAccounts,
        basket.data.rebalanceSlippage.toNumber(),
        oraclePriceData[from] / 10 ** tokenList[from].decimals,
        oraclePriceData[to] / 10 ** tokenList[to].decimals,
        jupAPIkey
    ).catch((e) => {
        console.log("---------- Error ------------");
        console.log("Jup Tx Data", e.message);
        console.log("---------- End Error ------------");
        return null;
    });

    return data;
}

/**
 * Builds a rebalance transaction.
 * This function creates a transaction that will perform the actual rebalancing of tokens in the basket on the Solana blockchain.
 * @param basket The basket to rebalance
 * @param from From token
 * @param to To token
 * @param tokenAmount Token amount
 * @param jupData Jupiter swap data
 * @param lookupTableAccount Lookup table account
 * @param wallet The wallet to use
 * @param connection The Solana connection
 * @param program The program to use
 * @param tokenList The token list
 * @param lamports The lamports to use
 * @returns A transaction to send to the Solana blockchain
 */
export async function buildRebalanceTransaction(
    basket: Basket,
    from: number,
    to: number,
    tokenAmount: number,
    jupData: JupiterSwapData,
    lookups: AddressLookupTableAccount[],
    walletPublicKey: PublicKey,
    connection: Connection,
    program: Program<BasketsIDL>,
    tokenList: TokenSettings[],
    lamports: number
): Promise<TransactionToSend> {
    const { addressLookupTableAddresses, swapInstruction, setupInstructions } = jupData;

    const processedSetupInstructions = processInstructions(setupInstructions);
    const processedSwapInstruction = processInstruction(swapInstruction);

    const lookupTableAccounts = await getAddressLookupTableAccounts(connection, addressLookupTableAddresses);

    return {
        payerKey: walletPublicKey,
        instructions: [
            await buildUpdateCurrentWeightsIx(program, basket, tokenList),
            ...processedSetupInstructions,
            await buildWithdrawBeforeRebalanceIx(
                program, walletPublicKey, basket, tokenList, from, to, tokenAmount
            ),
            processedSwapInstruction,
            await buildDepositAfterRebalanceIx(
                program, walletPublicKey, basket, tokenList, to
            ),
            ComputeBudgetProgram.setComputeUnitLimit({units: ADDITIONAL_UNITS}),
            ComputeBudgetProgram.setComputeUnitPrice({microLamports: lamports})
        ],
        lookupTables: [...lookupTableAccounts, ...lookups]
    };
}

/**
 * Processes instructions by converting pubkeys to PublicKey objects.
 * @param instructions Instructions to process
 * @returns Processed instructions
 */
export function processInstructions(instructions: any[]): TransactionInstruction[] {
    return instructions.map(processInstruction);
}

/**
 * Processes a single instruction by converting pubkeys to PublicKey objects.
 * @param instruction Instruction to process
 * @returns Processed instruction
 */
export function processInstruction(instruction: any): TransactionInstruction {
    return {
        programId: new PublicKey(instruction.programId),
        keys: instruction.accounts.map((a: any) => ({ ...a, pubkey: new PublicKey(a.pubkey) })),
        data: Buffer.from(instruction.data, "base64")
    };
}

/**
 * Signs and sends transactions to the Solana blockchain.
 * @param txsToSend Transactions to send
 * @param connection The Solana connection
 * @param wallet The wallet to use
 * @returns An array of transaction signatures
 */
export async function signAndSendTransactions(
    txsToSend: TransactionToSend[],
    connection: Connection,
    wallet: any,
    confirmFirst: number,
): Promise<TransactionSignature[]> {
    if (txsToSend.length == 0) return [];
    const blockhash = (await connection.getLatestBlockhash("confirmed")).blockhash;
    const signedTransactions = await signVersionedTransactions(
        wallet,
        txsToSend.map(tx => new VersionedTransaction(
            new TransactionMessage({
                payerKey: tx.payerKey,
                recentBlockhash: blockhash,
                instructions: tx.instructions,
            }).compileToV0Message(tx.lookupTables)
        ))
    ).catch((e) => { console.log("Sign V Transactions", e); return []; });
    return sendSignedTransactions(connection, signedTransactions, confirmFirst);
}

/**
 * Calculates the flash rebalance amounts for a basket of tokens.
 * This function determines which tokens need to be rebalanced based on their current weights,
 * target weights, and the basket's rebalance threshold.
 * 
 * @param numTokens - The number of tokens in the basket
 * @param timestamp - The current timestamp
 * @param lastRebalanceTime - Array of timestamps for the last rebalance of each token
 * @param rebalanceInterval - The interval between rebalances
 * @param currentCompToken - Array of current token compositions
 * @param currentCompAmount - Array of current token amounts
 * @param targetWeights - Array of target weights for each token
 * @param weightSum - The sum of all target weights
 * @param tokenList - List of token settings
 * @param rebalanceThreshold - The threshold for rebalancing
 * @param oraclePriceData - Array of current oracle prices for each token
 * 
 * @returns An object containing arrays of over-weighted and under-weighted tokens
 */
export function calculateFlashRebalanceAmounts(
    numTokens: number,
    timestamp: number,
    lastRebalanceTime: number[],
    rebalanceInterval: number,
    currentCompToken: number[],
    currentCompAmount: number[],
    targetWeights: number[],
    weightSum: number,
    tokenList: TokenSettings[],
    rebalanceThreshold: number,
    oraclePriceData: number[],
): {
    over: TokenRebalanceInfo[]
    under: TokenRebalanceInfo[]
} {
    const currentValues: number[] = [];
    let basketWorth = 0;

    // Calculate current values and total basket worth
    for (let i = 0; i < numTokens; i++) {
        const price = oraclePriceData[currentCompToken[i]];
        const tokenAmount = currentCompAmount[i] / Math.pow(10, tokenList[currentCompToken[i]].decimals);
        const tokenValue = price * tokenAmount;
        currentValues.push(tokenValue);
        basketWorth += tokenValue;
    }

    // Return empty arrays if basket is worth nothing
    if (basketWorth === 0) return { over: [], under: [] };

    const res: { over: TokenRebalanceInfo[], under: TokenRebalanceInfo[] } = { over: [], under: [] };

    // Determine over-weighted and under-weighted tokens
    for (let i = 0; i < numTokens; i++) {
        const currentPercentage = (currentValues[i] / basketWorth) * 10000;
        const targetPercentage = Math.floor((targetWeights[i] / weightSum) * 10000);
        const recentlyRebalanced = (lastRebalanceTime[i] + rebalanceInterval > timestamp);
        const inThresholds = (
            currentPercentage >= targetPercentage * (1 - rebalanceThreshold / 10000) &&
            currentPercentage <= targetPercentage * (1 + rebalanceThreshold / 10000)
        );
        const diffOnChain = Math.floor(currentPercentage) == 0 ? 0 :
            Math.floor(
                currentCompAmount[i] * 
                Math.floor(Math.floor(currentPercentage) - targetPercentage) /
                Math.floor(currentPercentage)
            );

        // Create TokenRebalanceInfo object
        const item: TokenRebalanceInfo = {
            token: currentCompToken[i],
            value: Math.abs((currentPercentage - targetPercentage) * basketWorth) / 10000,
            amount: targetPercentage == 0 ? currentCompAmount[i] : diffOnChain,
            rR: recentlyRebalanced,
            iT: inThresholds,
            uM: targetPercentage === 0 ? currentCompAmount[i] : 0
        };

        // Categorize as over-weighted or under-weighted
        if (currentValues[i] >= targetPercentage * basketWorth / 10000) {
            res.over.push(item);
        } else {
            res.under.push(item);
        }
    }

    return res;
}

/**
 * Retrieves flash rebalance information for a given basket.
 * 
 * @param basket - The basket to analyze
 * @param tokenList - List of token settings
 * @param oraclePriceData - Array of current oracle prices for each token
 * @param timestamp - The current timestamp
 * 
 * @returns An object containing arrays of over-weighted and under-weighted tokens
 */
export function getFlashRebalanceInfo(
    basket: Basket,
    tokenList: TokenSettings[],
    oraclePriceData: number[],
    timestamp: number,
): {
    over: TokenRebalanceInfo[]
    under: TokenRebalanceInfo[]
} {
    // Extract relevant data from the basket
    const {
        currentCompToken,
        currentCompAmount,
        targetWeight,
        numOfTokens,
        rebalanceThreshold,
        weightSum,
        rebalanceInterval,
        lastRebalanceTime
    } = basket.data;

    // Call calculateFlashRebalanceAmounts with extracted and parsed data
    return calculateFlashRebalanceAmounts(
        parseInt(numOfTokens.toString()),
        timestamp,
        lastRebalanceTime.map((x: BN) => parseInt(x.toString())),
        parseInt(rebalanceInterval.toString()),
        currentCompToken.map((x: BN) => parseInt(x.toString())),
        currentCompAmount.map((x: BN) => parseInt(x.toString())),
        targetWeight.map((x: BN) => parseInt(x.toString())),
        parseInt(weightSum.toString()),
        tokenList,
        parseInt(rebalanceThreshold.toString()),
        oraclePriceData,
    );
}