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

import { SolanaProgramService } from '../program';
import type { LiqsolCore } from '../../../assets/solana/types/liqsol_core';
import {
    deriveDistributionStatePda,
    deriveGlobalConfigPda,
    deriveLiqsolMintPda,
    derivePayRateHistoryPda,
    deriveUserRecordPda,
} from '../constants';
import type { DistributionState, DistributionUserRecord, GlobalConfig, PayRateEntry, PayRateHistory } from '../types';
import { getAssociatedTokenAddressSync, TOKEN_2022_PROGRAM_ID } from '@solana/spl-token';
import { ceilDiv } from '../utils';

/**
 * 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<LiqsolCore>;

    constructor(private provider: AnchorProvider) {
        const svc = new SolanaProgramService(provider);
        this.program = svc.getProgram('liqsolCore');
    }

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

    /**
     * Fetch the global distribution state account.
     *
     * IDL account name: "distributionState"
     */
    async getDistributionState(): Promise<DistributionState | null> {
        const pda = deriveDistributionStatePda();
        try {
            return await this.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 = derivePayRateHistoryPda();
        try {
            // Anchor types map directly onto our PayRateHistory TS type
            return (await this.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 = deriveGlobalConfigPda();
        try {
            return (await this.program.account.globalConfig.fetch(
                pda,
            ));
        } 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 liqsolMint = deriveLiqsolMintPda();

        // 1) Try treating ownerOrAta as a WALLET: derive its liqSOL ATA, then userRecord.
        try {
            const ata = getAssociatedTokenAddressSync(
                liqsolMint,
                ownerOrAta,
                false,
                TOKEN_2022_PROGRAM_ID,
            );
            const pdaFromWallet = deriveUserRecordPda(ata);
            const recFromWallet =
                await this.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 = deriveUserRecordPda(ownerOrAta);
            return await this.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 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: any = 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));
    }
}