import { PublicKey } from "@solana/web3.js";
import { BN, Program } from "@coral-xyz/anchor";
import { BasketsProgram } from "../idl/types";
import { loadOraclePrice, OraclePrice, OracleType } from "../utils/oracle";
import { MANAGERS_PER_BASKET, PYTHNET_CUSTODY_PRICE_USDC_ACCOUNT, TOTAL_WEIGHT } from "../utils/constants";
import { PYTHNET_CUSTODY_PRICE_SOL_ACCOUNT } from "../utils/constants";
import { USDC_DECIMALS } from "../utils/constants";
import { getAccountInfos } from "../utils/programAccounts";
import { WSOL_DECIMALS } from "../utils/constants";


export const BASKETS_STATE_SIZE = 28189;

export interface BasketState {
    version: number;
    ownAddress: PublicKey;
    basketType: number;
    basketPda: PublicKey;
    mint: PublicKey;
    supplyOutstanding: BN;
    lastPrice: BN;
    startingPrice: BN;
    highestPrice: BN;

    creator: PublicKey;
    creatorDepositFeeBps: number;
    creatorManagementFeeBps: number;
    creatorPerformanceFeeBps: number;

    host: PublicKey;
    hostDepositFeeBps: number;
    hostManagementFeeBps: number;
    hostPerformanceFeeBps: number;

    managers: PublicKey[];
    managersWeightBps: number[];
    managersAuthority: number[];
    managersDepositFeeBps: number;
    managersPerformanceFeeBps: number;
    managersManagementFeeBps: number;

    basketDepositFeeBps: number;
    basketWithdrawFeeBps: number;

    rebalanceIntervalSeconds: BN;
    rebalanceThresholdBps: number;
    rebalanceSlippageBps: number;
    lpThresholdBps: number;
    allowAutomation: number;
    allowLp: number;
    lamportsForAutomation: BN;

    symbolLength: number;
    symbol: number[];
    nameLength: number;
    name: number[];
    uriLength: number;
    uri: number[];
    metadataAccount: PublicKey;

    lookupTable1: PublicKey;
    lookupTable2: PublicKey;
    otherLookupTable1: PublicKey;
    otherLookupTable2: PublicKey;

    writeVersion: BN;
    numTokens: number;
    compositionMints: PublicKey[];
    compositionDecimals: number[];
    compositionOracleType: number[];
    compositionOracle1: PublicKey[];
    compositionOracle2: PublicKey[];
    compositionTargetWeights: number[];
    compositionAmounts: BN[];
    tokenPrices: BN[];
    tokenPriceUpdateTimestamps: BN[];
    lastRebalanceTimestamp: BN[];

    extraData: PublicKey[];
}

export interface ParsedBasketState {
    version: number;
    ownAddress: string;
    basketType: number;
    basketPda: string;
    mint: string;
    supplyOutstanding: number;
    lastPrice: number;
    startingPrice: number;
    highestPrice: number;

    creator: string;
    creatorDepositFeeBps: number;
    creatorManagementFeeBps: number;
    creatorPerformanceFeeBps: number;

    host: string;
    hostDepositFeeBps: number;
    hostManagementFeeBps: number;
    hostPerformanceFeeBps: number;

    managers: string[];
    managersWeightBps: number[];
    managersAuthority: number[];
    managersDepositFeeBps: number;
    managersPerformanceFeeBps: number;
    managersManagementFeeBps: number;

    basketDepositFeeBps: number;
    basketWithdrawFeeBps: number;

    rebalanceIntervalSeconds: number;
    rebalanceThresholdBps: number;
    rebalanceSlippageBps: number;
    lpThresholdBps: number;
    allowAutomation: number;
    allowLp: number;
    lamportsForAutomation: number;

    metadataAccount: string;

    lookupTable1: string;
    lookupTable2: string;
    otherLookupTable1: string;
    otherLookupTable2: string;

    writeVersion: number;
    numTokens: number;
    compositionMints: string[];
    compositionDecimals: number[];
    compositionOracleType: number[];
    compositionOracle1: string[];
    compositionOracle2: string[];
    compositionTargetWeights: number[];
    compositionAmounts: number[];
    tokenPrices: number[];
    tokenPriceUpdateTimestamps: number[];
    lastRebalanceTimestamp: number[];
    metadata: any;
    tvl: any;
    tokenValues: any;
}

export async function fetchBasketState(
    program: Program<BasketsProgram>,
    basket: PublicKey
): Promise<BasketState> {
    return await program.account.basketV200.fetch(basket);
}

export function parseBasketState(
    basketState: BasketState
): ParsedBasketState {
    const managers = [];
    for (let i = 0; i < MANAGERS_PER_BASKET; i++) {
        if (basketState.managers[i].equals(PublicKey.default)) {
            break;
        }
        managers.push(basketState.managers[i]);
    }
    return {
        version: basketState.version,
        ownAddress: basketState.ownAddress.toBase58(),
        basketType: basketState.basketType,
        basketPda: basketState.basketPda.toBase58(),
        mint: basketState.mint.toBase58(),
        supplyOutstanding: parseInt(basketState.supplyOutstanding.toString()),
        lastPrice: parseInt(basketState.lastPrice.toString()),
        startingPrice: parseInt(basketState.startingPrice.toString()),
        highestPrice: parseInt(basketState.highestPrice.toString()),
    
        creator: basketState.creator.toBase58(),
        creatorDepositFeeBps: basketState.creatorDepositFeeBps,
        creatorManagementFeeBps: basketState.creatorManagementFeeBps,
        creatorPerformanceFeeBps: basketState.creatorPerformanceFeeBps,
    
        host: basketState.host.toBase58(),
        hostDepositFeeBps: basketState.hostDepositFeeBps,
        hostManagementFeeBps: basketState.hostManagementFeeBps,
        hostPerformanceFeeBps: basketState.hostPerformanceFeeBps,
    
        managers: basketState.managers.slice(0, managers.length).map((manager) => manager.toBase58()),
        managersWeightBps: basketState.managersWeightBps.slice(0, managers.length),
        managersAuthority: basketState.managersAuthority.slice(0, managers.length),
        managersDepositFeeBps: basketState.managersDepositFeeBps,
        managersPerformanceFeeBps: basketState.managersPerformanceFeeBps,
        managersManagementFeeBps: basketState.managersManagementFeeBps,
    
        basketDepositFeeBps: basketState.basketDepositFeeBps,
        basketWithdrawFeeBps: basketState.basketWithdrawFeeBps,
    
        rebalanceIntervalSeconds: parseInt(basketState.rebalanceIntervalSeconds.toString()),
        rebalanceThresholdBps: basketState.rebalanceThresholdBps,
        rebalanceSlippageBps: basketState.rebalanceSlippageBps,
        lpThresholdBps: basketState.lpThresholdBps,
        allowAutomation: basketState.allowAutomation,
        allowLp: basketState.allowLp,
        lamportsForAutomation: parseInt(basketState.lamportsForAutomation.toString()),
    
        // symbolLength: basketState.symbolLength,
        // symbol: basketState.symbol,
        // nameLength: basketState.nameLength,
        // name: basketState.name,
        // uriLength: basketState.uriLength,
        // uri: basketState.uri,
        metadataAccount: basketState.metadataAccount.toBase58(),
    
        lookupTable1: basketState.lookupTable1.toBase58(),
        lookupTable2: basketState.lookupTable2.toBase58(),
        otherLookupTable1: basketState.otherLookupTable1.toBase58(),
        otherLookupTable2: basketState.otherLookupTable2.toBase58(),
    
        writeVersion: parseInt(basketState.writeVersion.toString()),
        numTokens: basketState.numTokens,
        compositionMints: basketState.compositionMints.slice(0, basketState.numTokens).map((mint) => mint.toBase58()),
        compositionDecimals: basketState.compositionDecimals.slice(0, basketState.numTokens),
        compositionOracleType: basketState.compositionOracleType.slice(0, basketState.numTokens),
        compositionOracle1: basketState.compositionOracle1.slice(0, basketState.numTokens).map((oracle) => oracle.toBase58()),
        compositionOracle2: basketState.compositionOracle2.slice(0, basketState.numTokens).map((oracle) => oracle.toBase58()),
        compositionTargetWeights: basketState.compositionTargetWeights.slice(0, basketState.numTokens),
        compositionAmounts: basketState.compositionAmounts.slice(0, basketState.numTokens).map(x => parseInt(x.toString())),
        tokenPrices: basketState.tokenPrices.slice(0, basketState.numTokens).map(x => parseInt(x.toString())),
        tokenPriceUpdateTimestamps: basketState.tokenPriceUpdateTimestamps.slice(0, basketState.numTokens).map(x => parseInt(x.toString())),
        lastRebalanceTimestamp: basketState.lastRebalanceTimestamp.slice(0, basketState.numTokens).map(x => parseInt(x.toString())),
        
        metadata: null,
        tvl: null,
        tokenValues: null,
    };
}

export async function getBasketTokenPrices(
    program: Program<BasketsProgram>,
    basketState: BasketState,
): Promise<OraclePrice[]> {
    // Collect oracle accounts to fetch
    const oracleAccounts: PublicKey[] = [
        PYTHNET_CUSTODY_PRICE_USDC_ACCOUNT,
        PYTHNET_CUSTODY_PRICE_SOL_ACCOUNT,
    ];
    for (let i = 0; i < basketState.numTokens; i++) {
        oracleAccounts.push(basketState.compositionOracle1[i]);
        
        const oracle2 = basketState.compositionOracle2[i];
        if (!oracle2.equals(PublicKey.default)) {
            oracleAccounts.push(oracle2);
        }
    }
    const oracleAccountInfos = await getAccountInfos(program.provider.connection, oracleAccounts);
    const usdcPrice = loadOraclePrice(USDC_DECIMALS, OracleType.Pyth, oracleAccountInfos[0], null, 0, 0, 0).avgPrice;
    const solPrice = loadOraclePrice(WSOL_DECIMALS, OracleType.Pyth, oracleAccountInfos[1], null, 0, 0, 0).avgPrice;

    // Calculate prices for each token
    const oraclePrices: OraclePrice[] = [];
    let currentIndex = 2;

    for (let i = 0; i < basketState.numTokens; i++) {
        const oracleAccount1 = oracleAccountInfos[currentIndex];
        let oracleAccount2 = oracleAccountInfos[currentIndex];
        currentIndex++;

        if (!basketState.compositionOracle2[i].equals(PublicKey.default)) {
            oracleAccount2 = oracleAccountInfos[currentIndex];
            currentIndex++;
        }

        const oraclePrice = loadOraclePrice(
            basketState.compositionDecimals[i],
            basketState.compositionOracleType[i],
            oracleAccount1,
            oracleAccount2,
            solPrice,
            usdcPrice,
            0,
        );

        oraclePrices.push(oraclePrice);
    }

    return oraclePrices;
}

export interface RebalanceInfo {
    token: PublicKey,
    tokenDecimals: number,
    tokenPrice: number,
    index: number,
    currentAmount: number,
    currentWeight: number,
    currentValue: number,
    targetWeight: number,
    targetValue: number,
    valueDiff: number,
    maxSpendAmount: number,
}

export function computeRebalanceInfos(
    params: {
        basketState: BasketState,
        oraclePrices: OraclePrice[],
    }
): {
    tvl: number, 
    tokenValues: any[],
    rebalanceInfos: RebalanceInfo[],
} {
    let basketValue = 0;
    const values: number[] = [];
    for (let i = 0; i < params.basketState.numTokens; i++) {
        const tokenPrice = params.oraclePrices[i].avgPrice;
        const currentAmount = parseInt(params.basketState.compositionAmounts[i].toString());
        const decimals = params.basketState.compositionDecimals[i];
        const currentValue = tokenPrice * currentAmount / (10 ** decimals);
        basketValue += currentValue;
        values.push(currentValue);
    };
    const tokenValues: any[] = [];
    const rebalanceInfos: RebalanceInfo[] = [];
    for (let i = 0; i < params.basketState.numTokens; i++) {
        const tokenPrice = params.oraclePrices[i].avgPrice;
        const currentAmount = parseInt(params.basketState.compositionAmounts[i].toString());
        const decimals = params.basketState.compositionDecimals[i];
        const currentValue = tokenPrice * currentAmount / (10 ** decimals);
        const currentWeight = Math.floor(currentValue * TOTAL_WEIGHT / basketValue);
        const targetWeight = params.basketState.compositionTargetWeights[i];
        const targetValue = basketValue * targetWeight / TOTAL_WEIGHT;
        let maxSpendAmount = 0;
        if (currentWeight > targetWeight) {
            maxSpendAmount = Math.floor(currentAmount * (currentWeight - targetWeight) / currentWeight);
        }
        rebalanceInfos.push({
            token: params.basketState.compositionMints[i],
            tokenDecimals: decimals,
            tokenPrice: tokenPrice,
            index: i,
            currentAmount: currentAmount,
            currentWeight: currentWeight,
            currentValue: currentValue,
            targetWeight: targetWeight,
            targetValue: targetValue,
            valueDiff: targetValue - currentValue,
            maxSpendAmount: maxSpendAmount,
        });
        tokenValues.push([Number(values[i].toFixed(6)), (currentWeight/100).toFixed(2) + "% -> " + (targetWeight/100).toFixed(2)+ "%"]);
    }

    return {
        tvl: basketValue,
        tokenValues: tokenValues,
        rebalanceInfos: rebalanceInfos,
    };
}

export async function getBasketTvl(
    program: Program<BasketsProgram>,
    basketState: BasketState,
) {
    const oraclePrices = await getBasketTokenPrices(program, basketState);
    return computeRebalanceInfos({
        basketState: basketState,
        oraclePrices: oraclePrices,
    });
}