import { AnchorProvider, Program, BN } from '@coral-xyz/anchor';
import { PublicKey, SystemProgram, TransactionInstruction } from '@solana/web3.js';

import { LiqsolCoreClientIdl, SolanaProgramService } from '../program';
import type { DistributionState, DistributionUserRecord, GlobalConfig, PayRateEntry, PayRateHistory } from '../types';
import { getAssociatedTokenAddressSync, TOKEN_2022_PROGRAM_ID } from '@solana/spl-token';
import { ceilDiv } from '../utils';

const INDEX_SCALE_BN = new BN('1000000000000');

/**
 * Distribution client – wraps the distribution portion of the liqsol_core
 * program in the *new* shares-only model.
 *
 * Responsibilities:
 *  - Read DistributionState (global index/share pool)
 *  - Read per-user UserRecord (shares)
 *  - Provide simple helpers for querying user share positions
 *
 * There is no longer any notion of:
 *  - tracked vs actual balances
 *  - mismatch candidates
 *  - updateUser / correct-&-register flow
 *
 * On-chain truth:
 *  - User balances live in their liqSOL ATA.
 *  - DistributionState + UserRecord track *shares* for yield.
 */
export class DistributionClient {
    private program: Program<LiqsolCoreClientIdl>;

    constructor(
        private readonly provider: AnchorProvider,
        private readonly pgs: SolanaProgramService
    ) {
        this.program = pgs.getProgram('liqsolCore');
    }

    get connection() {
        return this.provider.connection;
    }

    private async getTokenBalance(ata: PublicKey): Promise<BN> {
        try {
            const bal = await this.connection.getTokenAccountBalance(ata);
            return new BN(bal.value.amount);
        } catch {
            return new BN(0);
        }
    }

    /**
     * Fetch the global distribution state account.
     *
     * IDL account name: "distributionState"
     */
    async getDistributionState(): Promise<DistributionState | null> {
        const pda = this.pgs.deriveDistributionStatePda();
        const program = this.program;
        try {
            return await program.account.distributionState.fetch(pda);
        } catch {
            return null;
        }
    }


    /**
     * Fetch the global payRateHistory account (circular buffer of scaled pay rates).
     *
     * IDL account name: "payRateHistory"
     */
    async getPayRateHistory(): Promise<PayRateHistory | null> {
        const pda = this.pgs.derivePayRateHistoryPda();
        const program = this.program;
        try {
            // Anchor types map directly onto our PayRateHistory TS type
            return (await program.account.payRateHistory.fetch(
                pda,
            )) as PayRateHistory;
        } catch {
            return null;
        }
    }

    /**
     * Fetch the globalConfig account (contains depositFeeMultiplier, etc).
     *
     * IDL account name: "globalConfig"
     */
    async getGlobalConfig(): Promise<GlobalConfig | null> {
        const pda = this.pgs.deriveGlobalConfigPda();
        const program = this.program;
        try {
            return (await program.account.globalConfig.fetch(
                pda,
            )) as GlobalConfig;
        } catch {
            return null;
        }
    }
    /**
     * Fetch a user's distribution userRecord (or null if missing).
     *
     * IDL account name: "userRecord"
     */
    async getUserRecord(
        ownerOrAta: PublicKey,
    ): Promise<DistributionUserRecord | null> {
        const program = this.program;
        const liqsolMint = this.pgs.deriveLiqsolMintPda();

        // 1) Try treating ownerOrAta as a WALLET: derive its liqSOL ATA, then userRecord.
        try {
            const ata = getAssociatedTokenAddressSync(
                liqsolMint,
                ownerOrAta,
                true,
                TOKEN_2022_PROGRAM_ID,
            );
            const pdaFromWallet = this.pgs.deriveUserRecordPda(ata);
            const recFromWallet =
                await program.account.userRecord.fetchNullable(pdaFromWallet);
            if (recFromWallet) {
                return recFromWallet;
            }
        } catch {
            // fall through
        }

        // 2) Fallback: treat ownerOrAta as already being the TOKEN ACCOUNT.
        try {
            const pdaFromAta = this.pgs.deriveUserRecordPda(ownerOrAta);
            return await program.account.userRecord.fetchNullable(pdaFromAta);
        } catch {
            return null;
        }
    }

    /**
     * Convenience: get the user's current distribution shares.
     *
     * Returns BN(0) if the user has no userRecord yet.
     */
    async getUserShares(user: PublicKey): Promise<BN> {
        const record = await this.getUserRecord(user);
        return record ? record.shares : new BN(0);
    }

    /**
     * Convenience: get both the user's shares and totalShares, plus a
     * floating-point ratio (userShares / totalShares).
     *
     * NOTE: ratio is a JS number and may lose precision for very large values,
     * but is fine for display / percentage UIs.
     */
    async getUserSharePosition(user: PublicKey): Promise<{
        shares: BN;
        totalShares: BN;
        ratio: number;
    }> {
        const [dist, userShares] = await Promise.all([
            this.getDistributionState(),
            this.getUserShares(user),
        ]);

        const totalShares = dist ? dist.totalShares : new BN(0);

        let ratio = 0;
        if (!totalShares.isZero()) {
            // use toNumber() for UI-friendly ratio; if you need exact math,
            // do it with BN or bigint at the call site instead.
            ratio = userShares.toNumber() / totalShares.toNumber();
        }

        return { shares: userShares, totalShares, ratio };
    }

    /**
     * Compute claimable liqSOL for a wallet:
     * claimable = max(0, entitled - actual)
     * entitled = floor(shares * syncedIndex / 1e12)
     *
     * `syncedIndex` mirrors on-chain `sync_index` behavior by incorporating any
     * unsynced bucket delta (bucket_balance - last_bucket_balance).
     */
    async getClaimableLiqsol(user: PublicKey): Promise<BN> {
        const liqsolMint = this.pgs.deriveLiqsolMintPda();
        const bucketAuthority = this.pgs.deriveBucketAuthorityPda();

        const userAta = getAssociatedTokenAddressSync(
            liqsolMint,
            user,
            true,
            TOKEN_2022_PROGRAM_ID,
        );
        const bucketTokenAccount = getAssociatedTokenAddressSync(
            liqsolMint,
            bucketAuthority,
            true,
            TOKEN_2022_PROGRAM_ID,
        );

        const [distributionState, userRecord, actualBalance, bucketBalance] = await Promise.all([
            this.getDistributionState(),
            this.getUserRecord(user),
            this.getTokenBalance(userAta),
            this.getTokenBalance(bucketTokenAccount),
        ]);

        if (!distributionState || !userRecord) {
            return new BN(0);
        }

        let syncedIndex = new BN(distributionState.currentIndex.toString());
        const totalShares = new BN(distributionState.totalShares.toString());
        const lastBucketBalance = new BN(distributionState.lastBucketBalance.toString());

        if (totalShares.gt(new BN(0)) && bucketBalance.gt(lastBucketBalance)) {
            const delta = bucketBalance.sub(lastBucketBalance);
            const indexDelta = delta.mul(INDEX_SCALE_BN).div(totalShares);
            if (indexDelta.gt(new BN(0))) {
                syncedIndex = syncedIndex.add(indexDelta);
            }
        }

        const shares = new BN(userRecord.shares.toString());
        const entitled = shares.mul(syncedIndex).div(INDEX_SCALE_BN);

        if (entitled.lte(actualBalance)) {
            return new BN(0);
        }
        return entitled.sub(actualBalance);
    }

    /**
     * Build claim_rewards instruction for a wallet.
     */
    async buildClaimRewardsIx(user: PublicKey): Promise<TransactionInstruction> {
        const program = this.program;
        const liqsolMint = this.pgs.deriveLiqsolMintPda();
        const distributionState = this.pgs.deriveDistributionStatePda();
        const bucketAuthority = this.pgs.deriveBucketAuthorityPda();

        const userAta = getAssociatedTokenAddressSync(
            liqsolMint,
            user,
            true,
            TOKEN_2022_PROGRAM_ID,
        );
        const bucketTokenAccount = getAssociatedTokenAddressSync(
            liqsolMint,
            bucketAuthority,
            true,
            TOKEN_2022_PROGRAM_ID,
        );

        const userRecord = this.pgs.deriveUserRecordPda(userAta);
        const bucketUserRecord = this.pgs.deriveUserRecordPda(bucketTokenAccount);
        const extraAccountMetaList = this.pgs.deriveExtraAccountMetaListPda(liqsolMint);

        return program.methods
            .claimRewards()
            .accounts({
                user,
                userAta,
                userRecord,
                bucketUserRecord,
                distributionState,
                extraAccountMetaList,
                liqsolCoreProgram: this.pgs.PROGRAM_IDS.LIQSOL_CORE,
                transferHookProgram: this.pgs.PROGRAM_IDS.TRANSFER_HOOK,
                liqsolMint,
                bucketAuthority,
                bucketTokenAccount,
                tokenProgram: TOKEN_2022_PROGRAM_ID,
                systemProgram: SystemProgram.programId
            })
            .instruction();
    }

    /**
     * Compute an average scaled pay rate over the most recent `windowSize`
     * valid entries in the pay-rate history circular buffer.
     *
     * Returns a BN scaled by 1e12 (same as on-chain).
     */
    async getAverageScaledPayRate(windowSize = 5): Promise<BN> {
        const history = await this.getPayRateHistory();
        if (!history) {
            return new BN(0);
        }

        const entries: PayRateEntry[] = history.entries ?? [];
        if (!entries.length) {
            return new BN(0);
        }

        const maxEntries: number =
            typeof history.maxEntries === 'number'
                ? history.maxEntries
                : entries.length;

        const rawTotalAdded = history.totalEntriesAdded ?? 0;
        const totalAddedBn = new BN(rawTotalAdded.toString());

        // No valid entries written yet
        if (totalAddedBn.isZero()) {
            return new BN(0);
        }

        const totalAdded = Math.min(
            totalAddedBn.toNumber(),
            maxEntries,
            entries.length,
        );

        if (!totalAdded) {
            return new BN(0);
        }

        const COUNT = Math.max(1, Math.min(windowSize, totalAdded));

        const currentIndexNum = Number(history.currentIndex ?? 0);

        // Most recently written entry is at currentIndex - 1 (mod maxEntries)
        let idx =
            currentIndexNum === 0
                ? maxEntries - 1
                : currentIndexNum - 1;

        let sum = new BN(0);
        let valid = 0;
        const zero = new BN(0);

        for (let i = 0; i < COUNT; i++) {
            const entry : any = entries[idx];
            if (entry) {
                // Support both camelCase and snake_case (for safety)
                const rawScaled =
                    entry.scaledRate ??
                    entry.scaled_rate ??
                    0;

                const rate = new BN(rawScaled.toString());
                if (rate.gt(zero)) {
                    sum = sum.add(rate);
                    valid += 1;
                }
            }

            // Walk backwards through the circular buffer
            idx = idx === 0 ? maxEntries - 1 : idx - 1;
        }

        if (!valid) {
            return new BN(0);
        }

        // Same behavior as the dashboard: use a ceiling-like average
        return ceilDiv(sum, new BN(valid));
    }
}
