import { Program, BN, AnchorProvider } from '@coral-xyz/anchor';

import {
    Connection,
    Keypair,
    PublicKey,
    SystemProgram,
    StakeProgram,
} from '@solana/web3.js';
import {
    TOKEN_2022_PROGRAM_ID,
    ASSOCIATED_TOKEN_PROGRAM_ID,
    getAssociatedTokenAddress,
} from '@solana/spl-token';

import {
    LIQSOL_CORE,
    LIQSOL_TOKEN,
    PAY_RATE_SCALE_FACTOR,
    DEFAULT_AVERAGE_PAY_RATE,
    EPHEMERAL_RENT_EXEMPTION,
    LAMPORTS_PER_SOL,
    lamportsToSol,
    solToLamports,
    deriveDepositAuthorityPda,
    deriveLiqsolMintPda,
    deriveLiqsolMintAuthorityPda,
    deriveVaultPda,
    deriveReservePoolPda,
    deriveStakeControllerStatePda,
    derivePayoutStatePda,
    deriveBucketAuthorityPda,
    deriveDistributionStatePda,
    deriveUserRecordPda,
    derivePayRateHistoryPda,
    deriveStakeControllerVaultPda,
    deriveEphemeralStakeAddress,
    DEFAULT_PAY_RATE_LOOKBACK,
    deriveOutpostGlobalStatePda,
    deriveOutpostPoolAuthorityPda,
    deriveOutpostAccountPda,
    deriveTrancheStatePda,
    deriveUserPretokenRecordPda,
    CHAINLINK_FEED,
    CHAINLINK_PROGRAM,
    deriveExtraAccountMetaListPda,
} from './constants';

import liqsolCoreIDL from '../../assets/solana/idl/liqsol_core.json';
import { LiqsolCore } from '../../assets/solana/types/liqsol_core';
import { GlobalState, TrancheState } from './types';
import { TrancheLadderItem, TrancheSnapshot } from '../../types';
import { ChainID } from '@wireio/core';

// -----------------------------------------------------------------------------
// Tranche Support
// -----------------------------------------------------------------------------

const INDEX_SCALE = BigInt(1_000_000_000_000); // 1e12
const USD_SCALE = BigInt(100_000_000);         // 1e8
const BPS = BigInt(10_000);

/** BN | bigint -> bigint helper (keeps code readable) */
export function toBigint(x: any): bigint {
    if (typeof x === 'bigint') return x;
    if (typeof x === 'number') return BigInt(x);
    // Anchor BN
    return BigInt(x.toString());
}

/**
 * Convert token amount -> shares using the same rule as on-chain tests:
 * shares = ceil(amount * INDEX_SCALE / currentIndex)
 */
export function tokensToShares(amount: bigint, currentIndex: bigint): bigint {
    if (amount === BigInt(0)) return BigInt(0);
    const num = amount * INDEX_SCALE;
    const q = num / currentIndex;
    const r = num % currentIndex;
    return r === BigInt(0) ? q : q + BigInt(1);
}

/**
 * Apply one forward growth step: value * (BPS + growthBps) / BPS.
 * Simple integer round-half-up.
 */
function growOnce(value: bigint, growthBps: number): bigint {
    const g = BigInt(growthBps);
    return (value * (BPS + g) + BPS / BigInt(2)) / BPS;
}

/**
 * Apply one backward step: value * BPS / (BPS + growthBps).
 * Also integer round-half-up.
 */
function shrinkOnce(value: bigint, growthBps: number): bigint {
    const g = BigInt(growthBps);
    return (value * BPS + (BPS + g) / BigInt(2)) / (BPS + g);
}

/**
 * Build a local tranche ladder around the current tranche
 * using only on-chain config + current state.
 *
 * Rules (from liqsol_core tests):
 *  - past tranches are fully sold
 *  - current tranche has (initial - currentSupply) sold
 *  - future tranches start with 0 sold
 *  - supply/price grow by supplyGrowthBps/priceGrowthBps per tranche
 */
export function buildSolanaTrancheLadder(options: {
    currentTranche: number;
    initialTrancheSupply: bigint;
    currentTrancheSupply: bigint;
    currentPriceUsd: bigint;
    supplyGrowthBps: number;
    priceGrowthBps: number;
    windowBefore?: number;
    windowAfter?: number;
}): TrancheLadderItem[] {
    const {
        currentTranche,
        initialTrancheSupply,
        currentTrancheSupply,
        currentPriceUsd,
        supplyGrowthBps,
        priceGrowthBps,
        windowBefore = 5,
        windowAfter = 5,
    } = options;

    const startId = Math.max(0, currentTranche - windowBefore);
    const endId = currentTranche + windowAfter;

    const capacity = new Map<number, bigint>();
    const price = new Map<number, bigint>();

    // Seed current
    capacity.set(currentTranche, initialTrancheSupply);
    price.set(currentTranche, currentPriceUsd);

    // Forward (future tranches)
    for (let id = currentTranche + 1; id <= endId; id++) {
        const prevCap = capacity.get(id - 1)!;
        const prevPrice = price.get(id - 1)!;
        capacity.set(id, growOnce(prevCap, supplyGrowthBps));
        price.set(id, growOnce(prevPrice, priceGrowthBps));
    }

    // Backward (past tranches)
    for (let id = currentTranche - 1; id >= startId; id--) {
        const nextCap = capacity.get(id + 1)!;
        const nextPrice = price.get(id + 1)!;
        capacity.set(id, shrinkOnce(nextCap, supplyGrowthBps));
        price.set(id, shrinkOnce(nextPrice, priceGrowthBps));
    }

    const ladder: TrancheLadderItem[] = [];
    for (let id = startId; id <= endId; id++) {
        const cap = capacity.get(id)!;
        let sold: bigint;
        if (id < currentTranche) {
            sold = cap;
        } else if (id === currentTranche) {
            sold = cap - currentTrancheSupply;
        } else {
            sold = BigInt(0);
        }

        ladder.push({
            id,
            capacity: cap,
            sold,
            remaining: cap - sold,
            priceUsd: price.get(id)!,
        });
    }

    return ladder;
}

/**
 * Turn raw liqsol_core accounts into a chain-agnostic TrancheSnapshot for SOL.
 * All math stays here; TokenClient just wires accounts + connection.
 */
export function buildSolanaTrancheSnapshot(options: {
    chainID: ChainID;
    globalState: GlobalState;
    trancheState: TrancheState;
    solPriceUsd?: bigint;
    nativePriceTimestamp?: number;
    ladderWindowBefore?: number;
    ladderWindowAfter?: number;
}): TrancheSnapshot {
    const {
        chainID,
        globalState,
        trancheState,
        solPriceUsd,
        nativePriceTimestamp,
        ladderWindowBefore,
        ladderWindowAfter,
    } = options;

    const currentIndex = toBigint(globalState.currentIndex);
    const totalShares = toBigint(globalState.totalShares);
    const currentTranche = trancheState.currentTrancheNumber.toNumber();
    const currentTrancheSupply = toBigint(trancheState.currentTrancheSupply);
    const initialTrancheSupply = toBigint(trancheState.initialTrancheSupply);
    const totalPretokensSold = toBigint(trancheState.totalPretokensSold);
    const currentPriceUsd = toBigint(trancheState.currentTranchePriceUsd);
    const supplyGrowthBps = trancheState.supplyGrowthBps;
    const priceGrowthBps = trancheState.priceGrowthBps;

    const ladder = buildSolanaTrancheLadder({
        currentTranche,
        initialTrancheSupply,
        currentTrancheSupply,
        currentPriceUsd,
        supplyGrowthBps,
        priceGrowthBps,
        windowBefore: ladderWindowBefore,
        windowAfter: ladderWindowAfter,
    });

    return {
        chainID,
        currentIndex,
        totalShares,
        currentTranche,
        currentPriceUsd,
        supplyGrowthBps,
        priceGrowthBps,
        totalPretokensSold,
        currentTrancheSupply,
        initialTrancheSupply,
        nativePriceUsd: solPriceUsd,
        nativePriceTimestamp,
        ladder,
    };
}


// -----------------------------------------------------------------------------
// Read-only liqsol_core Program helper
// -----------------------------------------------------------------------------

let _liqsolCoreProgram: Program<LiqsolCore> | null = null;

export function getLiqsolCoreProgram(
    connection: Connection,
): Program<LiqsolCore> {
    if (_liqsolCoreProgram && _liqsolCoreProgram.provider.connection === connection) {
        return _liqsolCoreProgram;
    }

    // Dummy wallet – we're only doing read-only account fetches here.
    const tmpKeypair = Keypair.generate();
    const wallet: any = {
        publicKey: tmpKeypair.publicKey,
        signAllTransactions: async () => [],
        signTransaction: async () => tmpKeypair,
    };

    const provider = new AnchorProvider(connection, wallet, {
        commitment: 'confirmed',
    });

    const program = new Program(
        liqsolCoreIDL,
        provider,
    ) as Program<LiqsolCore>;

    _liqsolCoreProgram = program;
    return program;
}


// -----------------------------------------------------------------------------
// Balance + state getters (UI-friendly)
// -----------------------------------------------------------------------------

export async function getUserLiqSolBalance(
    connection: Connection,
    user: PublicKey,
): Promise<number> {
    const liqsolMint = deriveLiqsolMintPda();
    const ata = await getAssociatedTokenAddress(
        liqsolMint,
        user,
        false,
        TOKEN_2022_PROGRAM_ID,
    );

    try {
        const bal = await connection.getTokenAccountBalance(ata);
        return Number(bal.value.amount); // raw lamports, not decimals-adjusted
    } catch {
        return 0;
    }
}

export async function getBucketLiqSolBalance(
    connection: Connection,
): Promise<number> {
    const liqsolMint = deriveLiqsolMintPda();
    const bucketAuthority = deriveBucketAuthorityPda();
    const ata = await getAssociatedTokenAddress(
        liqsolMint,
        bucketAuthority,
        true,
        TOKEN_2022_PROGRAM_ID,
    );

    try {
        const bal = await connection.getTokenAccountBalance(ata);
        return Number(bal.value.amount);
    } catch {
        return 0;
    }
}

export async function getReservePoolBalance(
    connection: Connection,
): Promise<number> {
    const reservePool = deriveReservePoolPda();
    try {
        const lamports = await connection.getBalance(reservePool);
        return lamports;
    } catch {
        return 0;
    }
}

/**
 * Raw account info for the stake_controller state (Anchor decode is left
 * to the caller so the SDK can stay IDL-agnostic at this layer).
 */
export async function getStakeControllerStateRaw(
    connection: Connection,
): Promise<Uint8Array | null> {
    const pda = deriveStakeControllerStatePda();
    const info = await connection.getAccountInfo(pda);
    return info?.data ?? null;
}

/**
 * Raw account info for the payout-state account.
 */
export async function getPayoutStateRaw(
    connection: Connection,
): Promise<Uint8Array | null> {
    const pda = derivePayoutStatePda();
    const info = await connection.getAccountInfo(pda);
    return info?.data ?? null;
}

/**
 * Raw account info for a user's distribution user_record.
 *
 * NOTE: user_record is keyed by the TOKEN ACCOUNT (ATA), not the wallet.
 * Here we derive the user's liqSOL ATA and then its user_record PDA.
 */
export async function getUserRecordRaw(
    connection: Connection,
    user: PublicKey,
): Promise<Uint8Array | null> {
    const liqsolMint = deriveLiqsolMintPda();
    const userAta = await getAssociatedTokenAddress(
        liqsolMint,
        user,
        false,
        TOKEN_2022_PROGRAM_ID,
    );
    const pda = deriveUserRecordPda(userAta);
    const info = await connection.getAccountInfo(pda);
    return info?.data ?? null;
}


// -----------------------------------------------------------------------------
// Pay-rate & fee utilities (on-chain compatible)
// -----------------------------------------------------------------------------

export async function getAveragePayRate(
    connection: Connection,
    lookback: number = DEFAULT_PAY_RATE_LOOKBACK,
): Promise<bigint> {
    const program = getLiqsolCoreProgram(connection);
    const payRateHistoryPda = derivePayRateHistoryPda();

    try {
        const anyProgram = program;
        const payRateHistoryAccount = await anyProgram.account.payRateHistory.fetch(
            payRateHistoryPda,
        );

        const entries = payRateHistoryAccount.entries ?? [];
        const totalEntriesAdded = Number(
            payRateHistoryAccount.totalEntriesAdded ?? 0,
        );
        const currentIndex: number = payRateHistoryAccount.currentIndex ?? 0;
        const maxEntries: number =
            payRateHistoryAccount.maxEntries ?? entries.length;

        if (!entries.length) {
            return DEFAULT_AVERAGE_PAY_RATE;
        }

        const entriesToFetch = Math.min(lookback, maxEntries, entries.length);

        let idx: number;
        if (totalEntriesAdded === 0) {
            idx = 0;
        } else if (currentIndex === 0) {
            idx = maxEntries - 1;
        } else {
            idx = currentIndex - 1;
        }

        let sum = BigInt(0);
        let validCount = BigInt(0);

        for (let i = 0; i < entriesToFetch; i++) {
            const entry = entries[idx];
            if (entry && typeof entry.scaledRate !== 'undefined') {
                const rate = BigInt(entry.scaledRate.toString());
                if (rate > BigInt(0)) {
                    sum += rate;
                    validCount += BigInt(1);
                }
            }

            if (totalEntriesAdded === 0) {
                idx = (idx + 1) % maxEntries;
            } else {
                idx = idx === 0 ? maxEntries - 1 : idx - 1;
            }
        }

        if (validCount === BigInt(0)) {
            return DEFAULT_AVERAGE_PAY_RATE;
        }

        return sum / validCount;
    } catch (err) {
        // Fallback to default when history missing/unavailable
        return DEFAULT_AVERAGE_PAY_RATE;
    }
}

/**
 * On-chain fee formula:
 *   fee = (average_pay_rate * 4 * deposit_amount_lamports) / 10^12
 */
export function calculateExpectedFee(
    depositAmountLamports: BN,
    averagePayRate: BN,
): BN {
    return averagePayRate
        .mul(new BN(4))
        .mul(depositAmountLamports)
        .div(new BN(1_000_000_000_000)); // 10^12
}

/**
 * Convenience helper to preview how a deposit splits between user + bucket
 * and how much goes into the reserve, assuming the simple model:
 *
 * - userLiqSol = amount - fee + EPHEMERAL_RENT_EXEMPTION
 * - bucketLiqSol = fee
 * - reserveLamports = amount + EPHEMERAL_RENT_EXEMPTION
 */
export function previewDepositEffects(params: {
    depositAmountLamports: BN;
    averagePayRate: BN;
    rentExemptionLamports?: number;
}): {
    feeLamports: BN;
    userLiqSolLamports: BN;
    bucketLiqSolLamports: BN;
    reserveIncreaseLamports: BN;
} {
    const { depositAmountLamports, averagePayRate } = params;
    const rent = new BN(
        params.rentExemptionLamports ?? EPHEMERAL_RENT_EXEMPTION,
    );

    const fee = calculateExpectedFee(depositAmountLamports, averagePayRate);
    const userLiqSol = depositAmountLamports.sub(fee).add(rent);
    const bucketLiqSol = fee;
    const reserveIncrease = depositAmountLamports.add(rent);

    return {
        feeLamports: fee,
        userLiqSolLamports: userLiqSol,
        bucketLiqSolLamports: bucketLiqSol,
        reserveIncreaseLamports: reserveIncrease,
    };
}

/**
 * ---------------------------------------------------------------------------
 * Outpost-specific helpers
 * ---------------------------------------------------------------------------
 *
 * This is the SDK equivalent of `deriveWirePdas` in capital-staking tests.
 * It aggregates all PDAs / ATAs needed for Outpost pretokens + synd/desynd.
 */
export interface OutpostAccounts {
    user: PublicKey;

    // Outpost / pretoken wiring
    globalState: PublicKey;
    poolAuthority: PublicKey;
    liqsolMint: PublicKey;
    liqsolPoolAta: PublicKey;
    userAta: PublicKey;
    outpostAccount: PublicKey;
    userPretokenRecord: PublicKey;

    // Distribution / yield context (user_record keyed by token accounts)
    distributionState: PublicKey;
    bucketAuthority: PublicKey;
    bucketTokenAccount: PublicKey;
    bucketUserRecord: PublicKey;
    liqsolPoolUserRecord: PublicKey;
    userUserRecord: PublicKey;

    // Tranche / pricing
    trancheState: PublicKey;

    // Chainlink
    chainLinkFeed: PublicKey;
    chainLinkProgram: PublicKey;

    // Transfer hook extra accounts
    extraAccountMetaList: PublicKey;
}

export async function buildOutpostAccounts(
    connection: Connection,
    user: PublicKey,
): Promise<OutpostAccounts> {
    const globalState = deriveOutpostGlobalStatePda();
    const poolAuthority = deriveOutpostPoolAuthorityPda();
    const liqsolMint = deriveLiqsolMintPda();
    const distributionState = deriveDistributionStatePda();
    const bucketAuthority = deriveBucketAuthorityPda();
    const trancheState = deriveTrancheStatePda();

    // Outpost / pretoken PDAs
    const outpostAccount = deriveOutpostAccountPda(user);
    const userPretokenRecord = deriveUserPretokenRecordPda(user);

    // ATAs
    const liqsolPoolAta = await getAssociatedTokenAddress(
        liqsolMint,
        poolAuthority,
        true,
        TOKEN_2022_PROGRAM_ID,
    );

    const bucketTokenAccount = await getAssociatedTokenAddress(
        liqsolMint,
        bucketAuthority,
        true,
        TOKEN_2022_PROGRAM_ID,
    );

    const userAta = await getAssociatedTokenAddress(
        liqsolMint,
        user,
        false,
        TOKEN_2022_PROGRAM_ID,
    );

    // user_record PDAs keyed by TOKEN ACCOUNTS (matches capital-staking)
    const bucketUserRecord = deriveUserRecordPda(bucketTokenAccount);
    const liqsolPoolUserRecord = deriveUserRecordPda(liqsolPoolAta);
    const userUserRecord = deriveUserRecordPda(userAta);

    const extraAccountMetaList = deriveExtraAccountMetaListPda(liqsolMint);

    // Chainlink program feeds
    let chainLinkFeed = CHAINLINK_FEED;
    let chainLinkProgram = CHAINLINK_PROGRAM;

    try {
        const program = getLiqsolCoreProgram(connection);
        const ts: TrancheState =
            await program.account.trancheState.fetch(trancheState);

        // If on-chain config overrides are present, prefer those.
        if (ts.chainlinkFeed && ts.chainlinkProgram) {
            chainLinkFeed = ts.chainlinkFeed as PublicKey;
            chainLinkProgram = ts.chainlinkProgram as PublicKey;
        }
    } catch {
        // If trancheState isn't initialized yet, we fall back to the constants.
        // In that case, purchase instructions will still fail, which is the
        // correct behavior until admin initializes TrancheState on-chain.
    }

    void connection; // reserved for future extra validation

    return {
        user,
        globalState,
        poolAuthority,
        liqsolMint,
        liqsolPoolAta,
        userAta,
        outpostAccount,
        userPretokenRecord,
        distributionState,
        bucketAuthority,
        bucketTokenAccount,
        bucketUserRecord,
        liqsolPoolUserRecord,
        userUserRecord,
        trancheState,
        chainLinkFeed,
        chainLinkProgram,
        extraAccountMetaList
    };
}
// -----------------------------------------------------------------------------
// Epoch / scheduling helpers
// -----------------------------------------------------------------------------

export type EpochSnapshot = {
    epoch: number;
    slot: number;
    firstSlot: number;
    slotsInEpoch: number;
    slotMs: number;
};

export async function getEpochSnapshot(
    connection: Connection,
): Promise<EpochSnapshot> {
    const info = await connection.getEpochInfo();

    // Fallback slot time
    let slotTimeMs = Number(process.env.SLOT_TIME_MS_FALLBACK ?? 400);

    try {
        const samples = await connection.getRecentPerformanceSamples(16);
        if (samples.length) {
            const avgMs =
                samples.reduce(
                    (acc, s) => acc + (s.samplePeriodSecs * 1000) / s.numSlots,
                    0,
                ) / samples.length;
            if (isFinite(avgMs) && avgMs > 0) {
                slotTimeMs = avgMs;
            }
        }
    } catch {
        // ignore; keep fallback
    }

    return {
        epoch: info.epoch,
        slot: info.slotIndex,
        firstSlot: info.absoluteSlot - info.slotIndex,
        slotsInEpoch: info.slotsInEpoch,
        slotMs: slotTimeMs,
    };
}

export function msToEpochEnd(snapshot: EpochSnapshot): number {
    const remainingSlots = snapshot.slotsInEpoch - snapshot.slot;
    return remainingSlots * snapshot.slotMs;
}

/**
 * Generic "execute around epoch boundaries" helper:
 *
 * - If current progress is within [early, late], run immediately.
 * - If too early, sleep until `early` percentage into the epoch.
 * - If too late, sleep until `early` percentage into the *next* epoch.
 *
 * This is generic; you can wrap any instruction builder in here (including
 * deposit flows) without baking in program-specific logic.
 */
export async function scheduledInstruction<T>(
    connection: Connection,
    config: ScheduleConfig,
    instruction: () => Promise<T>,
): Promise<T> {
    const early = config.early ?? 0.10;
    const late = config.late ?? 0.90;

    const snapshot = await getEpochSnapshot(connection);
    const progress = snapshot.slot / snapshot.slotsInEpoch;

    // Case 1: Already in safe zone
    if (progress >= early && progress <= late) {
        return instruction();
    }

    // Case 2: Early in epoch → wait until `early`
    if (progress < early) {
        const targetSlot = snapshot.slotsInEpoch * early;
        const slotsRemaining = targetSlot - snapshot.slot;
        const msToWait = slotsRemaining * snapshot.slotMs;

        console.log(
            `Epoch early (${(progress * 100).toFixed(
                2,
            )}%). Sleeping ${(msToWait / 1000).toFixed(1)}s until ${early * 100
            }%...`,
        );

        await sleep(Math.max(1000, msToWait));
        return instruction();
    }

    // Case 3: Late in epoch → wait for next epoch + early window
    const msToNextEpoch = msToEpochEnd(snapshot);
    const earlyWaitTime = snapshot.slotsInEpoch * early * snapshot.slotMs;
    const totalWaitTime = msToNextEpoch + earlyWaitTime + 1000;

    console.log(
        `Epoch late (${(progress * 100).toFixed(
            2,
        )}%). Sleeping ${(totalWaitTime / 1000).toFixed(
            1,
        )}s until next epoch + ${early * 100}%...`,
    );

    await sleep(totalWaitTime);
    return instruction();
}


// -----------------------------------------------------------------------------
// Generic helpers (used in tests + can be useful in apps)
// -----------------------------------------------------------------------------

export function getErrorMessage(error: any): string {
    return (
        error?.error?.errorCode?.code ||
        error?.error?.errorMessage ||
        error?.message ||
        ''
    );
}

export function generateRandomDepositAmount(
    minSol = 2,
    maxSol = 100,
): BN {
    const randomSol = Math.random() * (maxSol - minSol) + minSol;
    return new BN(solToLamports(randomSol));
}

export function generateTestKeypair(): Keypair {
    return Keypair.generate();
}

export async function airdropSol(
    connection: Connection,
    publicKey: PublicKey,
    amountSol: number,
): Promise<void> {
    const lamports = solToLamports(amountSol);
    const sig = await connection.requestAirdrop(publicKey, Number(lamports));
    await connection.confirmTransaction(sig, 'confirmed');
}

export async function waitForConfirmation(
    connection: Connection,
    signature: string,
): Promise<void> {
    await connection.confirmTransaction(signature, 'confirmed');
}

export function sleep(ms: number): Promise<void> {
    return new Promise((resolve) => setTimeout(resolve, ms));
}

/**
 * Simple helper used in tests: wait until safe zone, no-op.
 */
export async function waitUntilSafeToExecuteFunction(
    connection: Connection,
    config: ScheduleConfig = {},
): Promise<void> {
    await scheduledInstruction(connection, config, async () => {
        // no-op
        return;
    });
}

export interface ScheduleConfig {
    early?: number; // fraction of epoch, default 0.10
    late?: number;  // fraction of epoch, default 0.90
}