// Core dependencies
import { Program } from "@coral-xyz/anchor";
import { Connection, Keypair, MessageV0, PublicKey, TransactionMessage, TransactionSignature, VersionedTransaction } from "@solana/web3.js";

// Local imports
import { JUPITER_API_KEY, MAX_JUPITER_ACCOUNTS, PRIORITY_FEE, PYTHNET_CUSTODY_PRICE_SOL_ACCOUNT, PYTHNET_CUSTODY_PRICE_USDC_ACCOUNT, USDC_DECIMALS, WSOL_DECIMALS } from "./utils/constants";
import { BasketsProgram } from "./idl/types";
import { getAccountInfos, getAllBaskets, getBasketsByCreator, getBasketsProgram, getPythSponsoredFeeds, getRandomSeed, getRaydiumCpmmPools, getRaydiumV4Pools, getWithdrawStateAccount, getWithdrawStatesByUser } from "./utils/programAccounts";
import { addLamportsForAutomationHandler, addNewTokenHandler, createBasketHandler, CreateBasketParams, createPythSponsoredFeedsHandler, EditBasketParams, editBasketSettingsHandler, removeTokenHandler, updatePythSponsoredFeedsHandler, updateTokenWeightsHandler } from "./basketManager";
import { buyBasketHandler, claimTokensHandler, sellBasketHandler, sellRebalanceHandler } from "./basketTrade";
import { rebalanceBasketTokensHandler, swapTokensHandler } from "./basketRebalance";
import { getQuoteResponseHandler } from "./instructions/jup";
import { fetchBasketState, getBasketTvl, parseBasketState, ParsedBasketState } from "./state/basket";
import { sendV0Transactions, VersionedTxs } from "./utils/txUtils";
import { BasketState } from "./state/basket";
import { fetchWithdrawState, ParsedWithdrawState, parseWithdrawState, WithdrawState } from "./state/withdrawState";
import { parseMetadata } from "./utils/metadataUtils";
import { loadOraclePrice, OracleType } from "./utils/oracle";
import { PoolInfo } from "./state/oracle";
import { parseRebalanceEvent } from "./utils/events";
export {
    BasketState,
    WithdrawState,
    VersionedTxs,
}

export class BasketsSDK {

    private sdkParams: {
        payer: PublicKey,
        connection: Connection,
        program: Program<BasketsProgram>,
        priorityFee: number,
        jupiterApiKey: string;
        maxAllowedAccounts: number;
    };

    constructor(params: {
        connection: Connection,
        payer?: PublicKey,
        priorityFee?: number,
        jupiterApiKey?: string,
        maxAllowedAccounts?: number,
    }) {
        this.sdkParams = {
            connection: params.connection,
            payer: params.payer ?? Keypair.generate().publicKey,
            program: getBasketsProgram(params.connection),
            priorityFee: params.priorityFee ?? PRIORITY_FEE,
            jupiterApiKey: params.jupiterApiKey ?? JUPITER_API_KEY,
            maxAllowedAccounts: params.maxAllowedAccounts ?? MAX_JUPITER_ACCOUNTS,
        };
    }

    async setPayer(payer: PublicKey) {
        this.sdkParams.payer = payer;
    }

    async createBasket(params: CreateBasketParams): Promise<{
        blockhash: string;
        lastValidBlockHeight: number;
        versionedTxs: VersionedTransaction[];
        batches: number[];
        address: PublicKey;
    }> {
        const basketKeypair = Keypair.generate();
        return {
            ...(await createBasketHandler(this.sdkParams, params, basketKeypair)),
            address: basketKeypair.publicKey,
        };
    }

    async editBasketSettings(params: EditBasketParams): Promise<VersionedTxs> {
        return await editBasketSettingsHandler(this.sdkParams, params);
    }

    async addLamportsForAutomation(params: {
        basket: PublicKey;
        amount: number;
    }): Promise<VersionedTxs> {
        return await addLamportsForAutomationHandler(this.sdkParams, params);
    }

    async findPoolsForToken(params: {
        token: PublicKey;
    }): Promise<PoolInfo[]> {
        const oracleAccounts: PublicKey[] = [
            PYTHNET_CUSTODY_PRICE_USDC_ACCOUNT,
            PYTHNET_CUSTODY_PRICE_SOL_ACCOUNT,
        ];
        const oracleAccountInfos = await getAccountInfos(this.sdkParams.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;
    
        const raydiumV4Pools = await getRaydiumV4Pools(this.sdkParams.connection, params.token, usdcPrice, solPrice);
        const raydiumCpmmPools = await getRaydiumCpmmPools(this.sdkParams.connection, params.token, usdcPrice, solPrice);
        const pythSponsoredFeeds = await getPythSponsoredFeeds(this.sdkParams.program, params.token);

        const allPools = [...raydiumV4Pools, ...raydiumCpmmPools, ...pythSponsoredFeeds];
        allPools.sort((a, b) => b.liquidity - a.liquidity);
        return allPools;
    }

    async createPythSponsoredFeed(): Promise<VersionedTxs> {
        return await createPythSponsoredFeedsHandler(this.sdkParams);
    }

    async updatePythSponsoredFeed(params: {
        tokenMint: PublicKey,
        feedAccount: PublicKey,
        isActive: boolean,
    }): Promise<VersionedTxs> {
        return await updatePythSponsoredFeedsHandler(this.sdkParams, params);
    }

    async addNewToken(params: {
        basket: PublicKey;
        token: PublicKey;
        tokenWeight: number;
        oracleType: number;
        oraclePool: PublicKey;
        oracle1: PublicKey;
        oracle2: PublicKey;
    }): Promise<VersionedTxs> {
       return await addNewTokenHandler(this.sdkParams, params);
    }

    async removeToken(params: {
        basket: PublicKey;
        token: PublicKey;
    }): Promise<VersionedTxs> {
        return await removeTokenHandler(this.sdkParams, params);
    }

    async updateTokenWeights(params: {
        basket: PublicKey;
        tokenWeights: number[];
        writeVersion: number;
    }): Promise<VersionedTxs> {
        return await updateTokenWeightsHandler(this.sdkParams, params);
    }

    async buyBasketBackend(params: {
        basket: PublicKey,
        depositAmount: number,
        depositMint: PublicKey,
    }): Promise<VersionedTxs> {
        return await buyBasketHandler(this.sdkParams, params);
    }

    async buyBasket(params: {
        basket: PublicKey,
        depositAmount: number,
        depositMint: PublicKey,
    }): Promise<VersionedTxs> {
        const response = await fetch(
            `https://api.symmetry.fi/baskets/v2/tx/buy` +
            `?payer=${this.sdkParams.payer.toBase58()}` +
            `&basket=${params.basket.toBase58()}` +
            `&depositMint=${params.depositMint.toBase58()}` +
            `&depositAmount=${params.depositAmount}` +
            `&priorityFee=${this.sdkParams.priorityFee}`
        );
        const data = await response.json();
        if (data.error)
            throw new Error(data.error);
        const { blockhash, lastValidBlockHeight, versionedTxs: rawTxs, batches } = data;
        const versionedTxs = rawTxs.map(
            (tx: any) => VersionedTransaction.deserialize(
                Uint8Array.from(Buffer.from(tx, "base64"))
            )
        );
        return {
            blockhash,
            lastValidBlockHeight,
            versionedTxs,
            batches,
        };
    }

    async sellBasketBackend(params: {
        basket: PublicKey;
        amountToWithdraw: number,
        destinationMint: PublicKey,
        rebalance: boolean,
    }): Promise<{
        blockhash: string;
        lastValidBlockHeight: number;
        versionedTxs: VersionedTransaction[];
        batches: number[];
        address: PublicKey;
    }> {
        const withdrawStateSeed = getRandomSeed();
        const address = getWithdrawStateAccount(withdrawStateSeed);
        return {
            ...(await sellBasketHandler(this.sdkParams, {
                ...params,
                withdrawStateSeed,
            })),
            address,
        };
    }

    async sellBasket(params: {
        basket: PublicKey;
        amountToWithdraw: number,
        destinationMint: PublicKey,
        rebalance: boolean,
    }): Promise<{
        blockhash: string;
        lastValidBlockHeight: number;
        versionedTxs: VersionedTransaction[];
        batches: number[];
        address: PublicKey;
    }> {
        const response = await fetch(
            `https://api.symmetry.fi/baskets/v2/tx/sell` +
            `?payer=${this.sdkParams.payer.toBase58()}` +
            `&basket=${params.basket.toBase58()}` +
            `&amountToWithdraw=${params.amountToWithdraw}` +
            `&destinationMint=${params.destinationMint.toBase58()}` +
            `&rebalance=${params.rebalance}` +
            `&priorityFee=${this.sdkParams.priorityFee}`
        );
        const data = await response.json();
        if (data.error)
            throw new Error(data.error);
        const { address, blockhash, lastValidBlockHeight, versionedTxs: rawTxs, batches } = data;
        const versionedTxs = rawTxs.map(
            (tx: any) => VersionedTransaction.deserialize(
                Uint8Array.from(Buffer.from(tx, "base64"))
            )
        );
        return {
            blockhash,
            lastValidBlockHeight,
            versionedTxs,
            batches,
            address,
        };
    }

    async rebalanceSellState(params: {
        withdrawState: PublicKey;
    }): Promise<VersionedTxs> {
        return await sellRebalanceHandler(this.sdkParams, params);
    }

    async claimTokens(params: {
        withdrawState: PublicKey;
    }): Promise<VersionedTxs> {
        return await claimTokensHandler(this.sdkParams, params);
    }

    async generateSwapQuote(params: {
        fromToken: PublicKey;
        toToken: PublicKey;
        amount: number;
        slippageBps: number;
    }): Promise<any> {
        return await getQuoteResponseHandler({
            jupiterApiKey: this.sdkParams.jupiterApiKey,
            maxAllowedAccounts: this.sdkParams.maxAllowedAccounts,
            ...params,
        });
    }

    async swapTokens(params: {
        basket: PublicKey;
        fromToken: PublicKey;
        toToken: PublicKey;
        fromAmount: number;
        quoteResponse: any;
        fromTokenWeight?: number;
        toTokenWeight?: number;
    }): Promise<VersionedTxs> {
        return await swapTokensHandler(this.sdkParams, params);
    }

    async rebalanceBasketTokensBackend(params: {
        basket: PublicKey;
        fromToken?: PublicKey;
        toToken?: PublicKey;
        minSwapValue?: number;
        maxSellValuePerToken?: number;
        maxNumberOfSwaps?: number;
    }): Promise<VersionedTxs> {
        return await rebalanceBasketTokensHandler(this.sdkParams, params);
    }

    async rebalanceBasketTokens(params: {
        basket: PublicKey;
        fromToken?: PublicKey;
        toToken?: PublicKey;
        minSwapValue?: number;
        maxSellValuePerToken?: number;
        maxNumberOfSwaps?: number;
    }): Promise<VersionedTxs> {
        let str = `https://api.symmetry.fi/baskets/v2/tx/rebalance` +
            `?payer=${this.sdkParams.payer.toBase58()}` +
            `&basket=${params.basket.toBase58()}` +
            `&priorityFee=${this.sdkParams.priorityFee}`;
        if (params.fromToken)
            str += `&fromToken=${params.fromToken.toBase58()}`;
        if (params.toToken)
            str += `&toToken=${params.toToken.toBase58()}`;
        if (params.minSwapValue)
            str += `&minSwapValue=${params.minSwapValue}`;
        if (params.maxSellValuePerToken)
            str += `&maxSellValuePerToken=${params.maxSellValuePerToken}`;
        if (params.maxNumberOfSwaps)
            str += `&maxNumberOfSwaps=${params.maxNumberOfSwaps}`;
        const response = await fetch(str);
        const data = await response.json();
        if (data.error)
            throw new Error(data.error);
        const { blockhash, lastValidBlockHeight, versionedTxs: rawTxs, batches } = data;
        const versionedTxs = rawTxs.map(
            (tx: any) => VersionedTransaction.deserialize(
                Uint8Array.from(Buffer.from(tx, "base64"))
            )
        );
        return {
            blockhash,
            lastValidBlockHeight,
            versionedTxs,
            batches,
        };
    }

    async getBasket(params: {
        basket: PublicKey;
        requestTvl?: boolean;
    }): Promise<ParsedBasketState> {
        const basketState = await fetchBasketState(this.sdkParams.program, params.basket);
        const metadata = await parseMetadata(this.sdkParams.connection, basketState.metadataAccount);
        let tvl = null, tokenValues = null;
        if (params.requestTvl)
            ({tvl, tokenValues} = await getBasketTvl(this.sdkParams.program, basketState));
        return {
            ...parseBasketState(basketState),
            metadata: metadata,
            tokenValues: tokenValues,
            tvl: tvl,
        }
    }

    async getWithdrawState(params: {
        withdrawState: PublicKey;
    }): Promise<ParsedWithdrawState> {
        const withdrawStateData = await fetchWithdrawState(this.sdkParams.program, params.withdrawState);
        return parseWithdrawState(withdrawStateData);
    }

    async getAllBaskets(): Promise<ParsedBasketState[]> {
        return await getAllBaskets(this.sdkParams.program);
    }

    async getBasketsByCreator(
        creator: PublicKey,
    ): Promise<ParsedBasketState[]> {
        return await getBasketsByCreator(this.sdkParams.program, creator);
    }

    async getWithdrawStatesByUser(
        user: PublicKey,
    ): Promise<ParsedWithdrawState[]> {
        return await getWithdrawStatesByUser(this.sdkParams.program, user);
    }

    async sendSignedVersionedTxs(
        txs: VersionedTxs,
        simulateTransactions: boolean = false,
    ): Promise<TransactionSignature[]> {
        return await sendV0Transactions(
            this.sdkParams.connection,
            txs,
            simulateTransactions
        );
    }

    addSwapListener(
        callback: (event: any, slot: number, signature: string) => void
    ) {
        return this.sdkParams.program.addEventListener(
            "rebalance",
            (event, slot, signature) => {
                let processedEvent = parseRebalanceEvent(event);
                try { callback(processedEvent, slot, signature) } catch (e: any) {
                    console.log(e.message);
                }
            }
        );
    }

    removeEventListener(id: number) {
        this.sdkParams.program.removeEventListener(id);
    }
}
