import { config } from 'dotenv';
import { Connection, Keypair, PublicKey } from '@solana/web3.js';
import { AnchorProvider, Wallet } from '@coral-xyz/anchor';
import { getAssociatedTokenAddress } from '@solana/spl-token';
import bs58 from 'bs58';
import fs from 'fs';
import { SystemProgram } from '@solana/web3.js';
import { ASSOCIATED_TOKEN_PROGRAM_ID, TOKEN_PROGRAM_ID } from '@solana/spl-token';
import { bigintToBN, MerkleDistributor, validateTimestamps } from '../src/index';

// Import the JITO merkle tree implementation (double hashing)
import { createJitoMerkleTree, AirdropRecipient } from '../src/utils/merkle-tree';

// Load environment variables
config();

// Custom PDA derivation that matches the SDK class
function getCorrectDistributorPDA(mint: PublicKey, version: bigint, programId: PublicKey): [PublicKey, number] {
  const versionBuffer = Buffer.alloc(8);
  versionBuffer.writeBigUInt64LE(version);
  return PublicKey.findProgramAddressSync(
    [
      Buffer.from('MerkleDistributor'),
      mint.toBuffer(),
      versionBuffer
    ],
    programId
  );
}

async function createDistributorV6Fixed() {
  console.log('🚀 Creating Merkle Distributor V6 with FIXED V2 (Double Hashing) Merkle Proofs\n');

  // Validate environment variables
  const privateKey = process.env.PRIVATE_KEY;
  const rpcEndpoint = process.env.RPC_ENDPOINT || 'https://mainnet.helius-rpc.com/?api-key=616ef0ca-f8ff-499b-8ca2-cd967fb07ef2';
  const usdcTokenMint = "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v"; // USDC mint address on mainnet
  const walletAddress = process.env.WALLET_ADDRESS || "CoJebSiqLWbXmSCmSSis8NSjva4ag93isJV7dxcz8x5q";
  const distributorVersion = BigInt(8); // V6 with fixed V2 (double hashing) merkle proofs

  if (!privateKey) {
    throw new Error('PRIVATE_KEY not found in environment variables');
  }

  console.log('📋 Configuration:');
  console.log(`RPC Endpoint: ${rpcEndpoint}`);
  console.log(`USDC Token Mint: ${usdcTokenMint}`);
  console.log(`Claimee Address: ${walletAddress}`);
  console.log(`Distributor Version: ${distributorVersion} (V6 - Fixed V2 Double Hashing)\n`);

  // Setup connection and wallet
  const connection = new Connection(rpcEndpoint, 'confirmed');
  const keypair = Keypair.fromSecretKey(bs58.decode(privateKey));
  const wallet = new Wallet(keypair);
  const provider = new AnchorProvider(connection, wallet, { commitment: 'confirmed' });

  console.log(`Connected wallet: ${wallet.publicKey.toString()}`);
  console.log(`Expected wallet: ${walletAddress}`);
  
  if (wallet.publicKey.toString() !== walletAddress) {
    throw new Error('Wallet address mismatch! Check your PRIVATE_KEY and WALLET_ADDRESS');
  }

  // Check wallet balance
  const balance = await connection.getBalance(wallet.publicKey);
  console.log(`Wallet SOL balance: ${balance / 1e9} SOL`);
  
  if (balance < 0.1 * 1e9) { // 0.1 SOL minimum
    console.warn('⚠️  Low SOL balance! You may need more SOL for transaction fees.');
  }

  // Initialize SDK
  const sdk = new MerkleDistributor(provider);
  const mint = new PublicKey(usdcTokenMint);

  // Create airdrop recipients (exactly 1 USDC to the specified address)
  const recipients: AirdropRecipient[] = [
    {
      address: new PublicKey("CoJebSiqLWbXmSCmSSis8NSjva4ag93isJV7dxcz8x5q"), // Specific claimee address
      unlockedAmount: 1000000, // 1 USDC (6 decimals = 1,000,000 units)
      lockedAmount: 0,    // No locked/vesting tokens
    }
  ];

  console.log('\n🌳 Creating FIXED V2 (Double Hashing) Merkle Tree:');
  console.log(`Recipients: ${recipients.length}`);
  recipients.forEach((r, i) => {
    console.log(`  ${i + 1}. ${r.address.toString()}`);
    console.log(`     Unlocked: ${r.unlockedAmount / 1000000} USDC`);
    if (r.lockedAmount > 0) {
      console.log(`     Locked: ${r.lockedAmount / 1000000} USDC`);
    }
  });

  // Generate merkle tree using JITO implementation (double hashing)
  const { tree, root, recipients: treeRecipients } = createJitoMerkleTree(recipients);
  console.log(`✅ FIXED V2 Merkle Root: ${Buffer.from(root).toString('hex')}`);
  console.log('🔧 This root uses DOUBLE HASHING to exactly match Rust program expectations!');
  console.log('💡 V2 matches the Rust program\'s two-step hashing process:');
  console.log('   1. hashv([claimant, amount_unlocked, amount_locked])');
  console.log('   2. hashv([LEAF_PREFIX, result_from_step_1])');

  // Calculate timestamps - set in future to account for transaction processing
  const now = BigInt(Math.floor(Date.now() / 1000));
  const startVesting = now + 10n; // Start 10 seconds from now
  const endVesting = now + 11n;   // End 11 seconds from now (minimal vesting period)
  const clawbackStart = now + 86400n * 30n; // 30 days from now for clawback

  console.log('\n⏰ Token Distribution Schedule:');
  console.log(`Distribution starts: ${new Date(Number(startVesting) * 1000).toLocaleString()}`);
  console.log(`All tokens unlocked: ${new Date(Number(endVesting) * 1000).toLocaleString()} (1 second later)`);
  console.log(`Clawback available: ${new Date(Number(clawbackStart) * 1000).toLocaleString()}`);

  // Validate timestamps
  const validation = validateTimestamps(startVesting, endVesting, clawbackStart);
  if (!validation.valid) {
    throw new Error(`Invalid timestamps: ${validation.error}`);
  }

  // Get distributor PDA using the correct derivation
  const [distributorPDA] = getCorrectDistributorPDA(mint, distributorVersion, sdk.programId);
  console.log(`\n🔑 Distributor PDA: ${distributorPDA.toString()}`);

  // Check if distributor already exists
  try {
    const existingDistributor = await sdk.getDistributor(distributorPDA);
    console.log('⚠️  Distributor V6 already exists!');
    console.log('Existing distributor details:', {
      version: existingDistributor.version.toString(),
      root: Buffer.from(existingDistributor.root).toString('hex'),
      maxTotalClaim: existingDistributor.maxTotalClaim.toString(),
      admin: existingDistributor.admin.toString(),
    });
    
    // Save distributor info for claiming
    const distributorInfo = {
      distributorPDA: distributorPDA.toString(),
      mint: mint.toString(),
      version: distributorVersion.toString(),
      merkleRoot: Buffer.from(root).toString('hex'),
      recipients: treeRecipients,
      timestamps: {
        startVesting: startVesting.toString(),
        endVesting: endVesting.toString(),
        clawbackStart: clawbackStart.toString(),
      },
      isFixedV2: true // Mark as using fixed V2 implementation
    };

    fs.writeFileSync('./distributor-v6-fixed-v2-info.json', JSON.stringify(distributorInfo, null, 2));
    console.log('\n💾 Distributor V6 (Fixed V2) info saved to distributor-v6-fixed-v2-info.json');
    console.log('You can use this file to claim tokens later with FIXED V2 proofs.');
    return;
  } catch (error) {
    console.log('✅ Distributor V6 does not exist yet, proceeding with creation...');
  }

  try {
    console.log('\n🚀 Creating distributor V6 with FIXED V2 (double hashing) merkle proofs...');
    
    // Get clawback receiver token account (ATA for USDC) - this should be for the wallet
    const clawbackReceiverATA = await getAssociatedTokenAddress(
      mint,
      wallet.publicKey
    );
    
    // Calculate the token vault (ATA for the distributor PDA)
    const tokenVault = await getAssociatedTokenAddress(
      mint,
      distributorPDA,
      true // allowOwnerOffCurve
    );
    
    console.log(`📍 Clawback receiver ATA: ${clawbackReceiverATA.toString()}`);
    console.log(`📍 Token vault ATA: ${tokenVault.toString()}`);
    
    const signature = await sdk.program.methods
      .newDistributor(
        bigintToBN(distributorVersion),
        Array.from(root),
        bigintToBN(BigInt(1000000)), // Total 1 USDC
        bigintToBN(BigInt(recipients.length)),
        bigintToBN(startVesting),
        bigintToBN(endVesting),
        bigintToBN(clawbackStart)
      )
      .accounts({
        clawbackReceiver: clawbackReceiverATA,
        mint: mint,
        tokenVault: tokenVault,
        admin: wallet.publicKey,
        systemProgram: SystemProgram.programId,
        associatedTokenProgram: ASSOCIATED_TOKEN_PROGRAM_ID,
        tokenProgram: TOKEN_PROGRAM_ID,
      })
      .rpc();

    console.log(`✅ Distributor V6 (FIXED V2) created successfully!`);
    console.log(`Transaction signature: ${signature}`);
    console.log(`View on Solana Explorer: https://explorer.solana.com/tx/${signature}?cluster=mainnet-beta`);

    // Save distributor info for claiming
    const distributorInfo = {
      distributorPDA: distributorPDA.toString(),
      mint: mint.toString(),
      version: distributorVersion.toString(),
      merkleRoot: Buffer.from(root).toString('hex'),
      recipients: treeRecipients,
      timestamps: {
        startVesting: startVesting.toString(),
        endVesting: endVesting.toString(),
        clawbackStart: clawbackStart.toString(),
      },
      transactionSignature: signature,
      isFixedV2: true // Mark as using fixed V2 implementation
    };

    fs.writeFileSync('./distributor-v6-fixed-v2-info.json', JSON.stringify(distributorInfo, null, 2));
    console.log('\n💾 Distributor V6 (Fixed V2) info saved to distributor-v6-fixed-v2-info.json');

    // Verify the creation
    console.log('\n🔍 Verifying distributor...');
    const createdDistributor = await sdk.getDistributor(distributorPDA);
    console.log('Verification successful!');
    console.log('Distributor V6 (Fixed V2) details:', {
      version: createdDistributor.version.toString(),
      root: Buffer.from(createdDistributor.root).toString('hex'),
      maxTotalClaim: createdDistributor.maxTotalClaim.toString(),
      totalClaimed: createdDistributor.totalAmountClaimed.toString(),
      admin: createdDistributor.admin.toString(),
      mint: createdDistributor.mint.toString(),
    });

    console.log('\n🎉 SUCCESS! Distributor V6 created with FIXED V2 (double hashing) merkle proofs!');
    console.log('📝 Key improvements over previous versions:');
    console.log('   ✅ Uses DOUBLE HASHING exactly like Rust program');
    console.log('   ✅ Step 1: hash(claimant + amount_unlocked + amount_locked)');
    console.log('   ✅ Step 2: hash(LEAF_PREFIX + step1_result)');
    console.log('   ✅ This should eliminate InvalidProof errors completely');
    console.log('\n📋 Next steps:');
    console.log('   1. Run the V6 claim script to test claiming your 1 USDC');
    console.log('   2. Verify that InvalidProof errors are finally resolved!');

  } catch (error) {
    console.error('\n❌ Error creating distributor V6:', error);
    throw error;
  }
}

// Run the script
if (require.main === module) {
  createDistributorV6Fixed()
    .then(() => process.exit(0))
    .catch((error) => {
      console.error('Script failed:', error);
      process.exit(1);
    });
}

export { createDistributorV6Fixed };