import * as anchor from '@coral-xyz/anchor';
import { Connection, Keypair, PublicKey } from '@solana/web3.js';
import { getAssociatedTokenAddress } from '@solana/spl-token';
import {
  DriftClient,
  User,
  initialize,
  PositionDirection,
  convertToNumber,
  PRICE_PRECISION,
  QUOTE_PRECISION,
  BASE_PRECISION,
  Wallet,
  PerpMarkets,
  getMarketOrderParams,
  BulkAccountLoader,
  BN,
  calculateBidAskPrice,
} from '@drift-labs/sdk';

const main = async () => {
  // Step 1: Choose your environment
  const env = 'devnet';
  // const env = 'mainnet-beta'; // Uncomment for mainnet

  // Step 2: Initialize Drift SDK
  const sdkConfig = initialize({ env });
  console.log('Initialized Drift SDK for', env);

  // Step 3: Set up wallet and provider
  const rpcUrl = process.env.ANCHOR_PROVIDER_URL || 'https://api.devnet.solana.com';
  const connection = new Connection(rpcUrl, 'confirmed');
  
  // Load your wallet keypair
  const privateKey = JSON.parse(process.env.BOT_PRIVATE_KEY || '[]');
  const keypair = Keypair.fromSecretKey(Uint8Array.from(privateKey));
  const wallet = new Wallet(keypair);
  
  const provider = new anchor.AnchorProvider(
    connection,
    wallet,
    { preflightCommitment: 'confirmed', commitment: 'confirmed' }
  );

  // Step 4: Check SOL balance
  const lamportsBalance = await connection.getBalance(wallet.publicKey);
  console.log('Wallet:', wallet.publicKey.toString());
  console.log('SOL balance:', lamportsBalance / 10 ** 9);

  if (lamportsBalance < 0.1 * 10 ** 9) {
    throw new Error('Insufficient SOL balance. Need at least 0.1 SOL for fees.');
  }

  // Step 5: Initialize Drift Client
  const driftPublicKey = new PublicKey(sdkConfig.DRIFT_PROGRAM_ID);
  const bulkAccountLoader = new BulkAccountLoader(
    connection,
    'confirmed',
    1000
  );
  
  const driftClient = new DriftClient({
    connection,
    wallet,
    programID: driftPublicKey,
    accountSubscription: {
      type: 'polling',
      accountLoader: bulkAccountLoader,
    },
  });
  
  await driftClient.subscribe();
  console.log('✅ Connected to Drift Protocol');

  // Step 6: Set up User client
  const userPublicKey = await driftClient.getUserAccountPublicKey();
  const user = new User({
    driftClient,
    userAccountPublicKey: userPublicKey,
    accountSubscription: {
      type: 'polling',
      accountLoader: bulkAccountLoader,
    },
  });

  // Step 7: Check if user account exists
  const userAccountExists = await user.exists();

  if (!userAccountExists) {
    console.log('Creating Drift user account...');
    
    // Get USDC token address
    const usdcTokenAddress = await getAssociatedTokenAddress(
      new PublicKey(sdkConfig.USDC_MINT_ADDRESS),
      wallet.publicKey
    );

    // Deposit $10,000 USDC as collateral
    const depositAmount = new BN(10000).mul(QUOTE_PRECISION);
    
    console.log('Depositing', depositAmount.toString(), 'USDC (raw amount)...');
    
    await driftClient.initializeUserAccountAndDepositCollateral(
      depositAmount,
      usdcTokenAddress
    );
    
    console.log('✅ User account created and funded');
  } else {
    console.log('✅ User account already exists');
  }

  await user.subscribe();

  // Step 8: Get market information
  const solMarketInfo = PerpMarkets[env].find(
    (market) => market.baseAssetSymbol === 'SOL'
  );
  
  if (!solMarketInfo) {
    throw new Error('SOL-PERP market not found');
  }

  const marketIndex = solMarketInfo.marketIndex;
  console.log('\n📊 SOL-PERP Market (Index:', marketIndex + ')');

  // Step 9: Get current market price
  const perpMarketAccount = driftClient.getPerpMarketAccount(marketIndex);
  const oracleData = driftClient.getOracleDataForPerpMarket(marketIndex);
  
  const [bid, ask] = calculateBidAskPrice(
    perpMarketAccount.amm,
    oracleData
  );

  const formattedBidPrice = convertToNumber(bid, PRICE_PRECISION);
  const formattedAskPrice = convertToNumber(ask, PRICE_PRECISION);

  console.log('Current Bid:', '$' + formattedBidPrice.toFixed(2));
  console.log('Current Ask:', '$' + formattedAskPrice.toFixed(2));

  // Step 10: Place a market order
  console.log('\n🚀 Placing 1 SOL-PERP LONG order...');
  
  const orderParams = getMarketOrderParams({
    baseAssetAmount: new BN(1).mul(BASE_PRECISION), // 1 SOL
    direction: PositionDirection.LONG,
    marketIndex: marketIndex,
  });

  const txSig = await driftClient.placePerpOrder(orderParams);
  
  console.log('✅ Order placed successfully!');
  console.log('Transaction signature:', txSig);
  console.log('View on Solscan:', `https://solscan.io/tx/${txSig}?cluster=${env}`);

  // Step 11: Get position info
  await user.fetchAccounts();
  const position = user.getPerpPosition(marketIndex);
  
  if (position) {
    const baseAmount = convertToNumber(position.baseAssetAmount, BASE_PRECISION);
    const quoteAmount = convertToNumber(position.quoteAssetAmount, QUOTE_PRECISION);
    
    console.log('\n📈 Position Summary:');
    console.log('Base Asset Amount:', baseAmount.toFixed(4), 'SOL');
    console.log('Quote Asset Amount:', '$' + quoteAmount.toFixed(2));
  }

  console.log('\n✨ Trade complete!');
};

main()
  .catch(console.error)
  .finally(() => process.exit());