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

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

// Load environment variables
config();

function getErrorMessage(error: any): string {
  if (error instanceof Error) {
    return error.message;
  }
  return String(error);
}

async function claimTokensV6Fixed() {
  console.log('🎯 Claiming USDC Tokens from Merkle Distributor V6 (FIXED V2 Double Hashing)\n');

  // Load distributor info
  let distributorInfo: any;
  try {
    const infoFile = fs.readFileSync('./distributor-v6-fixed-v2-info.json', 'utf-8');
    distributorInfo = JSON.parse(infoFile);
  } catch (error) {
    throw new Error('Could not load distributor-v6-fixed-v2-info.json. Make sure you have created the V6 distributor first.');
  }

  if (!distributorInfo.isFixedV2) {
    throw new Error('This claim script is for FIXED V2 distributors only. The loaded distributor info does not have the V2 flag.');
  }

  // 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 walletAddress = process.env.WALLET_ADDRESS || "CoJebSiqLWbXmSCmSSis8NSjva4ag93isJV7dxcz8x5q";

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

  console.log('📋 Configuration:');
  console.log(`RPC Endpoint: ${rpcEndpoint}`);
  console.log(`Wallet Address: ${walletAddress}`);
  console.log(`Distributor PDA: ${distributorInfo.distributorPDA}`);
  console.log(`Mint: ${distributorInfo.mint}`);
  console.log(`Merkle Root: ${distributorInfo.merkleRoot}`);
  console.log(`Version: ${distributorInfo.version} (V6 - FIXED V2 Double Hashing)`);

  // 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(`\n🔑 Connected wallet: ${wallet.publicKey.toString()}`);
  
  if (wallet.publicKey.toString() !== walletAddress) {
    throw new Error('Wallet address mismatch! Check your PRIVATE_KEY and WALLET_ADDRESS');
  }

  // Initialize SDK
  const sdk = new MerkleDistributor(provider);
  const distributorPDA = new PublicKey(distributorInfo.distributorPDA);
  const mint = new PublicKey(distributorInfo.mint);

  // Get distributor account
  console.log('\n🔍 Checking distributor status...');
  const distributor = await sdk.getDistributor(distributorPDA);
  console.log('Distributor found!');
  console.log(`  Version: ${distributor.version.toString()}`);
  console.log(`  Root: ${Buffer.from(distributor.root).toString('hex')}`);
  console.log(`  Max Total Claim: ${distributor.maxTotalClaim.toString()}`);
  console.log(`  Total Claimed: ${distributor.totalAmountClaimed.toString()}`);

  // Recreate the recipients list 
  const recipients: AirdropRecipient[] = [
    {
      address: new PublicKey("CoJebSiqLWbXmSCmSSis8NSjva4ag93isJV7dxcz8x5q"),
      unlockedAmount: 1000000, // 1 USDC
      lockedAmount: 0,
    }
  ];

  // Recreate the JITO merkle tree (double hashing)
  console.log('\n🌳 Recreating JITO (Double Hashing) Merkle Tree...');
  const { tree } = createJitoMerkleTree(recipients);
  const recreatedRoot = Buffer.from(tree.getRoot()).toString('hex');
  
  console.log(`Recreated Root: ${recreatedRoot}`);
  console.log(`Saved Root:     ${distributorInfo.merkleRoot}`);
  console.log(`Roots Match:    ${recreatedRoot === distributorInfo.merkleRoot}`);
  
  if (recreatedRoot !== distributorInfo.merkleRoot) {
    throw new Error('Merkle root mismatch! The recreated V2 tree does not match the saved info.');
  }
  console.log('✅ V2 merkle tree recreation successful!');

  // Generate JITO proof for the wallet
  console.log('\n🔐 Generating JITO (Double Hashing) Merkle Proof...');
  const { proof, index, recipient } = generateProofForRecipient(tree, recipients, wallet.publicKey);
  
  console.log(`Found recipient at index ${index}:`);
  console.log(`  Address: ${recipient.address.toString()}`);
  console.log(`  Unlocked: ${recipient.unlockedAmount / 1000000} USDC`);
  console.log(`  Locked: ${recipient.lockedAmount / 1000000} USDC`);
  console.log(`  Proof elements: ${proof.length}`);
  if (proof.length > 0) {
    console.log(`  Proof: ${proof.map(p => Buffer.from(p).toString('hex'))}`);
  } else {
    console.log(`  Proof: [] (empty - single recipient tree)`);
  }

  // Verify proof locally using FIXED V2 implementation
  const rawLeafData = tree.getRawLeafForRecipient(recipient);
  const isValidProof = tree.verifyProof(index, rawLeafData, proof.map(p => Buffer.from(p)));
  if (!isValidProof) {
    throw new Error('Generated FIXED V2 proof is invalid! There is still an issue with our V2 implementation.');
  }
  console.log('✅ FIXED V2 proof verified locally - should work with Rust program\'s double hashing!');

  // Show detailed comparison with Rust program expectations
  console.log('\n🔬 V2 Implementation Details:');
  console.log('This exactly matches the Rust program\'s two-step process:');
  console.log(`  Raw Data: ${rawLeafData.toString('hex')}`);
  
  const { createHash } = require('crypto');
  const step1Hash = createHash('sha256').update(rawLeafData).digest();
  const step2Hash = createHash('sha256').update(Buffer.concat([Buffer.from([0]), step1Hash])).digest();
  
  console.log(`  Step 1 (hashv raw): ${step1Hash.toString('hex')}`);
  console.log(`  Step 2 (+ prefix):  ${step2Hash.toString('hex')}`);
  console.log(`  Tree Root:          ${Buffer.from(tree.getRoot()).toString('hex')}`);
  console.log(`  Final Match:        ${step2Hash.toString('hex') === Buffer.from(tree.getRoot()).toString('hex')}`);

  // Check if already claimed
  const [claimStatusPDA] = PublicKey.findProgramAddressSync(
    [
      Buffer.from('ClaimStatus'),
      wallet.publicKey.toBuffer(),
      distributorPDA.toBuffer()
    ],
    sdk.programId
  );
  
  try {
    const claimStatus = await sdk.getClaimStatus(claimStatusPDA);
    console.log('\n⚠️  Tokens already claimed!');
    console.log(`  Claimed by: ${claimStatus.claimant.toString()}`);
    console.log(`  Unlocked amount: ${claimStatus.unlockedAmount.toString()}`);
    console.log(`  Locked amount: ${claimStatus.lockedAmount.toString()}`);
    return;
  } catch (error) {
    console.log('\n✅ No previous claim found - proceeding with V6 claim...');
  }

  // Get or create claimant's token account
  const claimantTokenAccount = await getAssociatedTokenAddress(mint, wallet.publicKey);
  
  console.log('\n💰 Token Account Setup:');
  console.log(`  Claimant Token Account: ${claimantTokenAccount.toString()}`);
  
  try {
    const accountInfo = await getAccount(connection, claimantTokenAccount);
    console.log(`  Current balance: ${accountInfo.amount.toString()}`);
  } catch (error) {
    console.log('  Account does not exist - will be created during claim');
  }

  try {
    console.log('\n🚀 Claiming tokens with FIXED V2 (Double Hashing) proof...');
    console.log('💡 This is the moment of truth - testing if double hashing resolves InvalidProof!');
    
    const signature = await sdk.claim({
      claimant: wallet.publicKey,
      distributor: distributorPDA,
      claimantTokenAccount: claimantTokenAccount,
      amountUnlocked: BigInt(recipient.unlockedAmount),
      amountLocked: BigInt(recipient.lockedAmount),
      proof: proof
    });

    console.log(`✅ CLAIM SUCCESSFUL! InvalidProof error has been RESOLVED!`);
    console.log(`Transaction signature: ${signature}`);
    console.log(`View on Solana Explorer: https://explorer.solana.com/tx/${signature}?cluster=mainnet-beta`);

    // Verify the claim
    console.log('\n🔍 Verifying successful claim...');
    const finalClaimStatus = await sdk.getClaimStatus(claimStatusPDA);
    console.log('Claim verification successful!');
    console.log(`  Claimant: ${finalClaimStatus.claimant.toString()}`);
    console.log(`  Unlocked Amount: ${finalClaimStatus.unlockedAmount.toString()}`);
    console.log(`  Locked Amount: ${finalClaimStatus.lockedAmount.toString()}`);

    // Check final token balance
    try {
      const finalAccount = await getAccount(connection, claimantTokenAccount);
      console.log(`  Final token balance: ${finalAccount.amount.toString()} (${Number(finalAccount.amount) / 1000000} USDC)`);
    } catch (error) {
      console.log('  Could not fetch final balance');
    }

    console.log('\n🎉🎉🎉 ULTIMATE SUCCESS! 🎉🎉🎉');
    console.log('🔧 The InvalidProof error has been COMPLETELY RESOLVED!');
    console.log('📝 Root cause was the double hashing approach:');
    console.log('   ✅ JavaScript V2 now matches Rust program exactly');
    console.log('   ✅ Step 1: hash(claimant + amount_unlocked + amount_locked)');
    console.log('   ✅ Step 2: hash(LEAF_PREFIX + step1_result)');
    console.log('   ✅ Merkle proof validation now works perfectly');
    console.log('   ✅ Token transfer completed successfully');
    console.log('\n🏆 The merkle distributor system is now fully functional!');

  } catch (error) {
    console.error('\n❌ Error claiming USDC tokens:', getErrorMessage(error));
    
    if (error instanceof Error) {
      if (error.message.includes('InvalidProof')) {
        console.log('\n💔 STILL getting InvalidProof error even with V2 double hashing!');
        console.log('   This suggests there may be even deeper issues to investigate.');
        console.log('   Possible remaining issues:');
        console.log('   - Different hash function (our SHA256 vs Solana\'s hashv)');
        console.log('   - Byte ordering or encoding differences');
        console.log('   - Program ID or account derivation mismatches');
      } else if (error.message.includes('ClaimExpired')) {
        console.log('\n💡 The claim window has expired.');
      } else if (error.message.includes('already claimed')) {
        console.log('\n💡 Tokens have already been claimed.');
      } else if (error.message.includes('insufficient funds')) {
        console.log('\n💡 Insufficient funds (check both SOL and token balances).');
      }
    }
    
    throw error;
  }
}

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

export { claimTokensV6Fixed };