use anchor_lang::prelude::*;
use anchor_spl::token::{self, Burn, Mint, MintTo, Token, TokenAccount};
use anchor_spl::associated_token::AssociatedToken;

declare_id!("8hcuEUBcuVBYyD53QPFQrRSWJjQy3UVURiqJpcuvgrTf");

#[program]
pub mod lottery_simple {
    use super::*;

    pub fn initialize(ctx: Context<Initialize>) -> Result<()> {
        let state = &mut ctx.accounts.state;
        state.bump = ctx.bumps.state;
        state.vault_bump = ctx.bumps.vault;
        state.authority = ctx.accounts.authority.key();
        state.mint = ctx.accounts.mint.key();
        state.vault = ctx.accounts.vault.key();
        state.total_supply = 0;
        state.dev_fee_bps = 300; // 3%
        state.dividend_fee_bps = 200; // 2%
        state.token_price_initial = 100_000; // 0.0001 SOL
        state.token_price_increment = 10_000; // 0.00001 SOL per token
        state.profit_per_share = 0;
        msg!("PoWH initialized with mint: {}", ctx.accounts.mint.key());
        Ok(())
    }

    pub fn register_user(ctx: Context<RegisterUser>) -> Result<()> {
        let user = &mut ctx.accounts.user;
        user.owner = ctx.accounts.owner.key();
        user.token_account = ctx.accounts.token_account.key();
        user.dividends_withdrawn = 0;
        msg!("User registered: {}", ctx.accounts.owner.key());
        Ok(())
    }

    pub fn buy_tokens(ctx: Context<BuyTokens>, sol_amount: u64) -> Result<()> {
        // Capture needed values before mutable borrow
        let state_bump = ctx.accounts.state.bump;
        let buyer_key = ctx.accounts.buyer.key();
        let state_account_info = ctx.accounts.state.to_account_info();
        
        let state = &mut ctx.accounts.state;
        
        // Calculate fees
        let dev_fee = (sol_amount as u128 * state.dev_fee_bps as u128) / 10_000;
        let dividend_fee = (sol_amount as u128 * state.dividend_fee_bps as u128) / 10_000;
        let net_purchase = sol_amount as u128 - dev_fee - dividend_fee;
        
        // Transfer dev fee from buyer to dev wallet
        if dev_fee > 0 {
            let transfer_ix = anchor_lang::solana_program::system_instruction::transfer(
                &ctx.accounts.buyer.key(),
                &ctx.accounts.dev_wallet.key(),
                dev_fee as u64,
            );
            
            anchor_lang::solana_program::program::invoke(
                &transfer_ix,
                &[
                    ctx.accounts.buyer.to_account_info(),
                    ctx.accounts.dev_wallet.to_account_info(),
                    ctx.accounts.system_program.to_account_info(),
                ],
            )?;
        }
        
        // Transfer net purchase + dividend fee to vault (vault holds everything except dev fee)
        let vault_amount = net_purchase + dividend_fee;
        if vault_amount > 0 {
            let transfer_ix = anchor_lang::solana_program::system_instruction::transfer(
                &ctx.accounts.buyer.key(),
                &ctx.accounts.vault.key(),
                vault_amount as u64,
            );
            
            anchor_lang::solana_program::program::invoke(
                &transfer_ix,
                &[
                    ctx.accounts.buyer.to_account_info(),
                    ctx.accounts.vault.to_account_info(),
                    ctx.accounts.system_program.to_account_info(),
                ],
            )?;
        }
        
        // Distribute dividends to existing holders
        if state.total_supply > 0 && dividend_fee > 0 {
            let dividend_per_token = (dividend_fee * PRECISION) / state.total_supply as u128;
            state.profit_per_share += dividend_per_token;
        }
        
        // Calculate tokens to mint using bonding curve
        let tokens_to_mint = calculate_tokens_for_sol(net_purchase as u64, state.total_supply, state)?;
        
        // Mint tokens to buyer
        let seeds = &[b"state".as_ref(), &[state_bump]];
        let signer_seeds = &[&seeds[..]];
        
        let cpi_ctx = CpiContext::new_with_signer(
            ctx.accounts.token_program.to_account_info(),
            MintTo {
                mint: ctx.accounts.mint.to_account_info(),
                to: ctx.accounts.buyer_token_account.to_account_info(),
                authority: state_account_info,
            },
            signer_seeds,
        );
        token::mint_to(cpi_ctx, tokens_to_mint)?;
        
        state.total_supply += tokens_to_mint;
        
        emit!(TradeEvent {
            user: buyer_key,
            tokens: tokens_to_mint,
            sol_amount,
            is_buy: true,
        });
        
        msg!("Bought {} tokens for {} SOL", tokens_to_mint, sol_amount);
        Ok(())
    }
    
    pub fn sell_tokens(ctx: Context<SellTokens>, tokens_to_sell: u64) -> Result<()> {
        let state = &mut ctx.accounts.state;
        require!(tokens_to_sell > 0, ErrorCode::InvalidAmount);
        require!(ctx.accounts.seller_token_account.amount >= tokens_to_sell, ErrorCode::InsufficientFunds);
        
        // Calculate SOL to return using bonding curve
        let calculated_sol = calculate_sol_for_tokens(tokens_to_sell, state.total_supply, state)?;
        
        // Safety check: limit to vault balance to prevent errors
        let vault_balance = ctx.accounts.vault.to_account_info().lamports();
        let sol_to_return = if calculated_sol > vault_balance {
            msg!("Calculated return {} exceeds vault balance {}, limiting to vault balance", calculated_sol, vault_balance);
            // Leave enough for rent exemption - be more conservative
            vault_balance.saturating_sub(5000) // Leave 5000 lamports for rent exemption
        } else {
            calculated_sol
        };
        
        // Burn tokens from seller
        let cpi_ctx = CpiContext::new(
            ctx.accounts.token_program.to_account_info(),
            Burn {
                mint: ctx.accounts.mint.to_account_info(),
                from: ctx.accounts.seller_token_account.to_account_info(),
                authority: ctx.accounts.seller.to_account_info(),
            },
        );
        token::burn(cpi_ctx, tokens_to_sell)?;
        
        state.total_supply -= tokens_to_sell;
        
        // Transfer SOL from vault to seller (direct lamport manipulation)
        // This is the correct approach for program-owned PDAs as per Solana documentation
        if sol_to_return > 0 {
            // Direct lamport manipulation - standard practice for program-owned accounts
            **ctx.accounts.vault.to_account_info().try_borrow_mut_lamports()? -= sol_to_return;
            **ctx.accounts.seller.to_account_info().try_borrow_mut_lamports()? += sol_to_return;
            
            msg!("Transferred {} lamports from vault to seller", sol_to_return);
        }
        
        emit!(TradeEvent {
            user: ctx.accounts.seller.key(),
            tokens: tokens_to_sell,
            sol_amount: sol_to_return,
            is_buy: false,
        });
        
        msg!("Sold {} tokens for {} SOL", tokens_to_sell, sol_to_return);
        Ok(())
    }
    
    pub fn withdraw_dividends(ctx: Context<WithdrawDividends>) -> Result<()> {
        let state = &ctx.accounts.state;
        let user = &mut ctx.accounts.user;
        let token_balance = ctx.accounts.token_account.amount as u128;
        
        // Calculate available dividends
        let total_dividends = (state.profit_per_share * token_balance) / PRECISION;
        let available_dividends = total_dividends - user.dividends_withdrawn;
        
        require!(available_dividends > 0, ErrorCode::NoDividends);
        
        user.dividends_withdrawn = total_dividends;
        
        // Transfer dividends from vault to user (direct lamport manipulation)
        // This is the correct approach for program-owned PDAs
        **ctx.accounts.vault.to_account_info().try_borrow_mut_lamports()? -= available_dividends as u64;
        **ctx.accounts.user_authority.to_account_info().try_borrow_mut_lamports()? += available_dividends as u64;
        
        emit!(WithdrawalEvent {
            user: ctx.accounts.user_authority.key(),
            amount: available_dividends as u64,
        });
        
        msg!("Withdrew {} lamports in dividends", available_dividends);
        Ok(())
    }
}

// Realistic bonding curve based on actual SOL in vault
fn calculate_tokens_for_sol(sol: u64, supply: u64, _state: &PowhState) -> Result<u64> {
    // Linear bonding curve: price increases with supply
    // Base price: 0.0001 SOL per token (100,000 lamports per token)
    // Price increases by 0.00001 SOL per billion tokens in supply
    let base_price_lamports = 100_000u128; // 0.0001 SOL in lamports
    let price_increment = 10_000u128; // 0.00001 SOL in lamports
    let supply_factor = (supply as u128) / 1_000_000_000; // Per billion tokens
    
    let current_price = base_price_lamports + (price_increment * supply_factor);
    let tokens = (sol as u128 * 1_000_000_000) / current_price; // Convert lamports to tokens
    
    msg!("Buy: {} lamports -> {} tokens at price {} lamports per token", sol, tokens, current_price);
    Ok(tokens as u64)
}

fn calculate_sol_for_tokens(tokens: u64, supply: u64, _state: &PowhState) -> Result<u64> {
    // Sell at 90% of current buy price to create spread
    let base_price_lamports = 100_000u128;
    let price_increment = 10_000u128;
    let supply_factor = (supply as u128) / 1_000_000_000;
    
    let current_price = base_price_lamports + (price_increment * supply_factor);
    let sell_price = (current_price * 90) / 100; // 90% of buy price
    let sol_lamports = (tokens as u128 * sell_price) / 1_000_000_000;
    
    msg!("Sell: {} tokens -> {} lamports at price {} lamports per token", tokens, sol_lamports, sell_price);
    Ok(sol_lamports as u64)
}

const PRECISION: u128 = 1_000_000_000_000_000_000; // 18 decimal places for dividend calculations

#[derive(Accounts)]
pub struct Initialize<'info> {
    #[account(
        init,
        payer = authority,
        space = 8 + PowhState::INIT_SPACE,
        seeds = [b"state"],
        bump
    )]
    pub state: Account<'info, PowhState>,
    
    /// CHECK: This is the vault PDA for holding SOL
    #[account(
        init,
        payer = authority,
        space = 0,
        seeds = [b"vault"],
        bump
    )]
    pub vault: UncheckedAccount<'info>,
    
    #[account(mut)]
    pub mint: Account<'info, Mint>,
    
    #[account(mut)]
    pub authority: Signer<'info>,
    
    pub system_program: Program<'info, System>,
    pub token_program: Program<'info, Token>,
}

#[derive(Accounts)]
pub struct RegisterUser<'info> {
    #[account(
        seeds = [b"state"],
        bump = state.bump
    )]
    pub state: Account<'info, PowhState>,
    
    #[account(
        init,
        payer = owner,
        space = 8 + User::INIT_SPACE,
        seeds = [b"user", owner.key().as_ref()],
        bump
    )]
    pub user: Account<'info, User>,
    
    #[account(
        init_if_needed,
        payer = owner,
        associated_token::mint = mint,
        associated_token::authority = owner,
    )]
    pub token_account: Account<'info, TokenAccount>,
    
    pub mint: Account<'info, Mint>,
    
    #[account(mut)]
    pub owner: Signer<'info>,
    
    pub system_program: Program<'info, System>,
    pub token_program: Program<'info, Token>,
    pub associated_token_program: Program<'info, AssociatedToken>,
}

#[derive(Accounts)]
pub struct BuyTokens<'info> {
    #[account(
        mut,
        seeds = [b"state"],
        bump = state.bump
    )]
    pub state: Account<'info, PowhState>,
    
    /// CHECK: This is the vault PDA
    #[account(
        mut,
        seeds = [b"vault"],
        bump = state.vault_bump
    )]
    pub vault: UncheckedAccount<'info>,
    
    #[account(mut)]
    pub mint: Account<'info, Mint>,
    
    #[account(
        mut,
        token::mint = mint,
        token::authority = buyer,
    )]
    pub buyer_token_account: Account<'info, TokenAccount>,
    
    /// CHECK: Dev wallet for fees
    #[account(mut)]
    pub dev_wallet: UncheckedAccount<'info>,
    
    #[account(mut)]
    pub buyer: Signer<'info>,
    
    pub system_program: Program<'info, System>,
    pub token_program: Program<'info, Token>,
}

#[derive(Accounts)]
pub struct SellTokens<'info> {
    #[account(
        mut,
        seeds = [b"state"],
        bump = state.bump
    )]
    pub state: Account<'info, PowhState>,
    
    /// CHECK: This is the vault PDA
    #[account(
        mut,
        seeds = [b"vault"],
        bump = state.vault_bump
    )]
    pub vault: UncheckedAccount<'info>,
    
    #[account(mut)]
    pub mint: Account<'info, Mint>,
    
    #[account(
        mut,
        token::mint = mint,
        token::authority = seller,
    )]
    pub seller_token_account: Account<'info, TokenAccount>,
    
    #[account(mut)]
    pub seller: Signer<'info>,
    
    pub system_program: Program<'info, System>,
    pub token_program: Program<'info, Token>,
}

#[derive(Accounts)]
pub struct WithdrawDividends<'info> {
    #[account(
        seeds = [b"state"],
        bump = state.bump
    )]
    pub state: Account<'info, PowhState>,
    
    /// CHECK: This is the vault PDA
    #[account(
        mut,
        seeds = [b"vault"],
        bump = state.vault_bump
    )]
    pub vault: UncheckedAccount<'info>,
    
    #[account(
        mut,
        seeds = [b"user", user_authority.key().as_ref()],
        bump
    )]
    pub user: Account<'info, User>,
    
    #[account(
        token::mint = mint,
        token::authority = user_authority,
    )]
    pub token_account: Account<'info, TokenAccount>,
    
    pub mint: Account<'info, Mint>,
    
    #[account(mut)]
    pub user_authority: Signer<'info>,
    
    pub system_program: Program<'info, System>,
}

#[account]
#[derive(InitSpace)]
pub struct PowhState {
    pub bump: u8,
    pub vault_bump: u8,
    pub authority: Pubkey,
    pub mint: Pubkey,
    pub vault: Pubkey,
    pub total_supply: u64,
    pub dev_fee_bps: u16,
    pub dividend_fee_bps: u16,
    pub token_price_initial: u64,
    pub token_price_increment: u64,
    pub profit_per_share: u128,
}

#[account]
#[derive(InitSpace)]
pub struct User {
    pub owner: Pubkey,
    pub token_account: Pubkey,
    pub dividends_withdrawn: u128,
}

#[event]
pub struct TradeEvent {
    pub user: Pubkey,
    pub tokens: u64,
    pub sol_amount: u64,
    pub is_buy: bool,
}

#[event]
pub struct WithdrawalEvent {
    pub user: Pubkey,
    pub amount: u64,
}

#[error_code]
pub enum ErrorCode {
    #[msg("Insufficient funds")]
    InsufficientFunds,
    #[msg("Invalid amount")]
    InvalidAmount,
    #[msg("No dividends available")]
    NoDividends,
}