use anchor_lang::prelude::*;
use anchor_lang::solana_program::program::invoke;
use anchor_lang::solana_program::instruction::Instruction;
use bytemuck::{Pod, Zeroable};

declare_id!("HdGLMarketplace1111111111111111111111111111");

const MAX_COMPUTE_UNITS: u64 = 1_000_000_000_000; // 1e12

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

    pub fn initialize(ctx: Context<Initialize>, chg_token: Pubkey, zk_verifier: Pubkey) -> Result<()> {
        let global = &mut ctx.accounts.global;
        global.owner = ctx.accounts.owner.key();
        global.chg_token = chg_token;
        global.zk_verifier = zk_verifier;
        global.current_epoch = 1;
        Ok(())
    }

    pub fn free_tick(ctx: Context<FreeTick>) -> Result<()> {
        let counters = &mut ctx.accounts.counters;
        counters.free_counter = counters.free_counter.checked_add(1).ok_or(HdGLError::Overflow)?;
        emit!(FreeTickEvent {
            free_counter: counters.free_counter,
        });
        Ok(())
    }

    pub fn premium_submit(
        ctx: Context<PremiumSubmit>,
        lattice_hash: [u8; 32],
        recursion_depth: u8,
        compute_units_used: u64,
        instance_id: u8,
        a: [u64; 2],
        b: [[u64; 2]; 2],
        c: [u64; 2],
        zkp_input: [u64; 4],
        commitment: [u64; 2],
    ) -> Result<()> {
        require!(instance_id < 8, HdGLError::InvalidInstanceId);
        require!(compute_units_used < MAX_COMPUTE_UNITS, HdGLError::Overflow);
        require!(zkp_input[0] != 0, HdGLError::ZKPRequired);

        // CPI to zk_verifier with correct serialization
        let ix_data = ZKPData { a, b, c, zkp_input, commitment };
        let ix = Instruction {
            program_id: ctx.accounts.zk_verifier.key(),
            accounts: vec![
                AccountMeta::new_readonly(ctx.accounts.zk_verifier.key(), false),
                AccountMeta::new_readonly(ctx.accounts.global.key(), false),
            ],
            data: bytemuck::bytes_of(&ix_data).to_vec(),
        };
        invoke(&ix, &[
            ctx.accounts.zk_verifier.to_account_info(),
            ctx.accounts.global.to_account_info(),
        ])?;

        let counters = &mut ctx.accounts.counters;
        counters.counter_b = counters.counter_b.checked_add(1).ok_or(HdGLError::Overflow)?;

        let global = &mut ctx.accounts.global;
        global.last_snapshot = lattice_hash;

        let contributor = &mut ctx.accounts.contributor;
        contributor.last_lattice = lattice_hash;

        let (reward_eth, reward_chg) = _compute_reward(compute_units_used);
        contributor.balance_eth = contributor.balance_eth.checked_add(reward_eth).ok_or(HdGLError::Overflow)?;
        contributor.balance_chg = contributor.balance_chg.checked_add(reward_chg).ok_or(HdGLError::Overflow)?;

        emit!(PremiumSubmissionEvent {
            contributor: ctx.accounts.contributor.key(),
            block_number: Clock::get()?.unix_timestamp,
            lattice_hash,
            counter_b: counters.counter_b,
            reward_eth,
            reward_chg,
            recursion_depth,
            compute_units_used,
            instance_id,
        });

        Ok(())
    }

    pub fn batch_premium_submit(
        ctx: Context<BatchPremiumSubmit>,
        lattice_hashes: Vec<[u8; 32]>,
        recursion_depths: Vec<u8>,
        compute_units_used: Vec<u64>,
        instance_ids: Vec<u8>,
        a: Vec<[u64; 2]>,
        b: Vec<[[u64; 2]; 2]>,
        c: Vec<[u64; 2]>,
        zkp_inputs: Vec<[u64; 4]>,
        commitments: Vec<[u64; 2]>,
    ) -> Result<()> {
        require!(
            lattice_hashes.len() == recursion_depths.len() &&
            lattice_hashes.len() == compute_units_used.len() &&
            lattice_hashes.len() == instance_ids.len() &&
            lattice_hashes.len() == a.len() &&
            lattice_hashes.len() == b.len() &&
            lattice_hashes.len() == c.len() &&
            lattice_hashes.len() == zkp_inputs.len() &&
            lattice_hashes.len() == commitments.len(),
            HdGLError::ArrayLengthMismatch
        );
        require!(lattice_hashes.len() > 0 && lattice_hashes.len() <= 50, HdGLError::InvalidBatchSize);

        let counters = &mut ctx.accounts.counters;
        let global = &mut ctx.accounts.global;
        let contributor = &mut ctx.accounts.contributor;
        let mut total_reward_eth = 0;
        let mut total_reward_chg = 0;

        for i in 0..lattice_hashes.len() {
            require!(instance_ids[i] < 8, HdGLError::InvalidInstanceId);
            require!(compute_units_used[i] < MAX_COMPUTE_UNITS, HdGLError::Overflow);
            require!(zkp_inputs[i][0] != 0, HdGLError::ZKPRequired);

            let ix_data = ZKPData {
                a: a[i],
                b: b[i],
                c: c[i],
                zkp_input: zkp_inputs[i],
                commitment: commitments[i],
            };
            let ix = Instruction {
                program_id: ctx.accounts.zk_verifier.key(),
                accounts: vec![
                    AccountMeta::new_readonly(ctx.accounts.zk_verifier.key(), false),
                    AccountMeta::new_readonly(ctx.accounts.global.key(), false),
                ],
                data: bytemuck::bytes_of(&ix_data).to_vec(),
            };
            invoke(&ix, &[
                ctx.accounts.zk_verifier.to_account_info(),
                ctx.accounts.global.to_account_info(),
            ])?;

            require!(lattice_hashes[i] == zkp_inputs[i][0].to_le_bytes()[..32].try_into().unwrap(), HdGLError::ZKPHashMismatch);

            counters.counter_b = counters.counter_b.checked_add(1).ok_or(HdGLError::Overflow)?;
            global.last_snapshot = lattice_hashes[i];
            contributor.last_lattice = lattice_hashes[i];

            let (reward_eth, reward_chg) = _compute_reward(compute_units_used[i]);
            total_reward_eth = total_reward_eth.checked_add(reward_eth).ok_or(HdGLError::Overflow)?;
            total_reward_chg = total_reward_chg.checked_add(reward_chg).ok_or(HdGLError::Overflow)?;

            emit!(PremiumSubmissionEvent {
                contributor: ctx.accounts.contributor.key(),
                block_number: Clock::get()?.unix_timestamp,
                lattice_hash: lattice_hashes[i],
                counter_b: counters.counter_b,
                reward_eth,
                reward_chg,
                recursion_depth: recursion_depths[i],
                compute_units_used: compute_units_used[i],
                instance_id: instance_ids[i],
            });
        }

        contributor.balance_eth = contributor.balance_eth.checked_add(total_reward_eth).ok_or(HdGLError::Overflow)?;
        contributor.balance_chg = contributor.balance_chg.checked_add(total_reward_chg).ok_or(HdGLError::Overflow)?;

        Ok(())
    }

    pub fn epoch_submit(
        ctx: Context<EpochSubmit>,
        lattice_hash: [u8; 32],
        compute_units_total: u64,
        epoch: u64,
        a: [u64; 2],
        b: [[u64; 2]; 2],
        c: [u64; 2],
        zkp_input: [u64; 4],
        commitment: [u64; 2],
    ) -> Result<()> {
        let global = &mut ctx.accounts.global;
        require!(epoch == global.current_epoch, HdGLError::InvalidEpoch);
        require!(compute_units_total < MAX_COMPUTE_UNITS, HdGLError::Overflow);
        require!(zkp_input[0] != 0, HdGLError::ZKPRequired);

        let ix_data = ZKPData { a, b, c, zkp_input, commitment };
        let ix = Instruction {
            program_id: ctx.accounts.zk_verifier.key(),
            accounts: vec![
                AccountMeta::new_readonly(ctx.accounts.zk_verifier.key(), false),
                AccountMeta::new_readonly(ctx.accounts.global.key(), false),
            ],
            data: bytemuck::bytes_of(&ix_data).to_vec(),
        };
        invoke(&ix, &[
            ctx.accounts.zk_verifier.to_account_info(),
            ctx.accounts.global.to_account_info(),
        ])?;

        require!(lattice_hash == zkp_input[0].to_le_bytes()[..32].try_into().unwrap(), HdGLError::ZKPHashMismatch);

        let counters = &mut ctx.accounts.counters;
        counters.counter_b = counters.counter_b.checked_add(1).ok_or(HdGLError::Overflow)?;
        global.last_snapshot = lattice_hash;
        global.epoch_snapshots[epoch as usize] = lattice_hash;
        global.current_epoch = global.current_epoch.checked_add(1).ok_or(HdGLError::Overflow)?;

        let contributor = &mut ctx.accounts.contributor;
        let (reward_eth, reward_chg) = _compute_reward(compute_units_total);
        contributor.balance_eth = contributor.balance_eth.checked_add(reward_eth).ok_or(HdGLError::Overflow)?;
        contributor.balance_chg = contributor.balance_chg.checked_add(reward_chg).ok_or(HdGLError::Overflow)?;

        emit!(EpochSubmissionEvent {
            aggregator: ctx.accounts.contributor.key(),
            epoch,
            lattice_hash,
            compute_units_total,
            reward_eth,
            reward_chg,
        });

        Ok(())
    }

    pub fn register_stealth_address(ctx: Context<RegisterStealthAddress>, stealth_address: [u8; 32]) -> Result<()> {
        require!(stealth_address != [0u8; 32] && stealth_address[0] != 0, HdGLError::InvalidStealthAddress);
        let contributor = &mut ctx.accounts.contributor;
        contributor.stealth_address = stealth_address;
        Ok(())
    }

    pub fn withdraw_eth(ctx: Context<WithdrawEth>) -> Result<()> {
        let contributor = &mut ctx.accounts.contributor;
        let amount = contributor.balance_eth;
        let stealth_address = contributor.stealth_address;
        require!(amount > 0, HdGLError::NoBalance);
        require!(stealth_address != [0u8; 32], HdGLError::StealthAddressRequired);

        **ctx.accounts.contributor.to_account_info().lamports.borrow_mut() -= amount;
        **ctx.accounts.recipient.to_account_info().lamports.borrow_mut() += amount;

        contributor.balance_eth = 0;
        emit!(WithdrawalEthEvent {
            user: ctx.accounts.recipient.key(),
            stealth_address,
            amount,
        });
        Ok(())
    }

    pub fn withdraw_chg(ctx: Context<WithdrawChg>) -> Result<()> {
        let contributor = &mut ctx.accounts.contributor;
        let amount = contributor.balance_chg;
        let stealth_address = contributor.stealth_address;
        require!(amount > 0, HdGLError::NoBalance);
        require!(stealth_address != [0u8; 32], HdGLError::StealthAddressRequired);

        let cpi_accounts = anchor_spl::token::Transfer {
            from: ctx.accounts.contributor_token.to_account_info(),
            to: ctx.accounts.recipient_token.to_account_info(),
            authority: ctx.accounts.contributor.to_account_info(),
        };
        let cpi_program = ctx.accounts.chg_token.to_account_info();
        anchor_spl::token::transfer(CpiContext::new(cpi_program, cpi_accounts), amount)?;

        contributor.balance_chg = 0;
        emit!(WithdrawalChgEvent {
            user: ctx.accounts.recipient.key(),
            stealth_address,
            amount,
        });
        Ok(())
    }

    fn _compute_reward(compute_units_used: u64) -> (u64, u64) {
        require!(compute_units_used < MAX_COMPUTE_UNITS, HdGLError::Overflow);
        let base_eth = 1_000_000_000_000_000 * 1_000_000;
        let base_chg = 1_000 * 1_000_000;
        let denominator = compute_units_used as u128 + 1;
        let reward_eth = (base_eth as u128 / denominator).try_into().unwrap_or(0);
        let reward_chg = (base_chg as u128 / denominator).try_into().unwrap_or(0);
        (reward_eth, reward_chg)
    }
}

#[derive(Accounts)]
pub struct Initialize<'info> {
    #[account(init, payer = owner, space = 8 + 32 + 32 + 32 + 32 + 8 + 8 * 1000, seeds = [b"global"], bump)]
    pub global: Account<'info, GlobalState>,
    #[account(mut)]
    pub owner: Signer<'info>,
    pub system_program: Program<'info, System>,
}

#[derive(Accounts)]
pub struct FreeTick<'info> {
    #[account(mut, seeds = [b"counters"], bump)]
    pub counters: Account<'info, Counters>,
    #[account(has_one = owner)]
    pub global: Account<'info, GlobalState>,
}

#[derive(Accounts)]
pub struct PremiumSubmit<'info> {
    #[account(mut, seeds = [b"counters"], bump)]
    pub counters: Account<'info, Counters>,
    #[account(mut, seeds = [b"global"], bump)]
    pub global: Account<'info, GlobalState>,
    #[account(mut, init_if_needed, payer = contributor, space = 8 + 32 + 8 + 8 + 32, seeds = [b"contributor", contributor.key().as_ref()], bump)]
    pub contributor: Account<'info, ContributorState>,
    pub zk_verifier: Program<'info, ZKVerifier>,
    pub chg_token: Program<'info, Token>,
    #[account(mut)]
    pub contributor: Signer<'info>,
    pub system_program: Program<'info, System>,
}

#[derive(Accounts)]
pub struct BatchPremiumSubmit<'info> {
    #[account(mut, seeds = [b"counters"], bump)]
    pub counters: Account<'info, Counters>,
    #[account(mut, seeds = [b"global"], bump)]
    pub global: Account<'info, GlobalState>,
    #[account(mut, init_if_needed, payer = contributor, space = 8 + 32 + 8 + 8 + 32, seeds = [b"contributor", contributor.key().as_ref()], bump)]
    pub contributor: Account<'info, ContributorState>,
    pub zk_verifier: Program<'info, ZKVerifier>,
    pub chg_token: Program<'info, Token>,
    #[account(mut)]
    pub contributor: Signer<'info>,
    pub system_program: Program<'info, System>,
}

#[derive(Accounts)]
pub struct EpochSubmit<'info> {
    #[account(mut, seeds = [b"counters"], bump)]
    pub counters: Account<'info, Counters>,
    #[account(mut, seeds = [b"global"], bump)]
    pub global: Account<'info, GlobalState>,
    #[account(mut, init_if_needed, payer = contributor, space = 8 + 32 + 8 + 8 + 32, seeds = [b"contributor", contributor.key().as_ref()], bump)]
    pub contributor: Account<'info, ContributorState>,
    pub zk_verifier: Program<'info, ZKVerifier>,
    pub chg_token: Program<'info, Token>,
    #[account(mut)]
    pub contributor: Signer<'info>,
    pub system_program: Program<'info, System>,
}

#[derive(Accounts)]
pub struct RegisterStealthAddress<'info> {
    #[account(mut, init_if_needed, payer = contributor, space = 8 + 32 + 8 + 8 + 32, seeds = [b"contributor", contributor.key().as_ref()], bump)]
    pub contributor: Account<'info, ContributorState>,
    #[account(mut)]
    pub contributor: Signer<'info>,
    pub system_program: Program<'info, System>,
}

#[derive(Accounts)]
pub struct WithdrawEth<'info> {
    #[account(mut, seeds = [b"global"], bump)]
    pub global: Account<'info, GlobalState>,
    #[account(mut, seeds = [b"contributor", contributor.key().as_ref()], bump)]
    pub contributor: Account<'info, ContributorState>,
    #[account(mut)]
    pub recipient: AccountInfo<'info>,
    #[account(mut)]
    pub contributor: Signer<'info>,
    pub system_program: Program<'info, System>,
}

#[derive(Accounts)]
pub struct WithdrawChg<'info> {
    #[account(mut, seeds = [b"global"], bump)]
    pub global: Account<'info, GlobalState>,
    #[account(mut, seeds = [b"contributor", contributor.key().as_ref()], bump)]
    pub contributor: Account<'info, ContributorState>,
    #[account(mut)]
    pub contributor_token: Account<'info, TokenAccount>,
    #[account(mut)]
    pub recipient_token: Account<'info, TokenAccount>,
    #[account(mut)]
    pub contributor: Signer<'info>,
    pub chg_token: Program<'info, Token>,
    pub system_program: Program<'info, System>,
}

#[error_code]
pub enum HdGLError {
    #[msg("Invalid instance ID")]
    InvalidInstanceId,
    #[msg("Arithmetic overflow")]
    Overflow,
    #[msg("Array length mismatch")]
    ArrayLengthMismatch,
    #[msg("Invalid batch size")]
    InvalidBatchSize,
    #[msg("ZKP required")]
    ZKPRequired,
    #[msg("ZKP hash mismatch")]
    ZKPHashMismatch,
    #[msg("Invalid epoch")]
    InvalidEpoch,
    #[msg("No balance to withdraw")]
    NoBalance,
    #[msg("Stealth address required")]
    StealthAddressRequired,
    #[msg("Invalid stealth address format")]
    InvalidStealthAddress,
}

#[account]
pub struct GlobalState {
    pub owner: Pubkey,
    pub chg_token: Pubkey,
    pub zk_verifier: Pubkey,
    pub last_snapshot: [u8; 32],
    pub current_epoch: u64,
    pub epoch_snapshots: [[u8; 32]; 1000], // Fixed size for simplicity
}

#[account]
pub struct Counters {
    pub counter_b: u64,
    pub free_counter: u64,
}

#[account]
pub struct ContributorState {
    pub last_lattice: [u8; 32],
    pub balance_eth: u64,
    pub balance_chg: u64,
    pub stealth_address: [u8; 32],
}

#[event]
pub struct FreeTickEvent {
    pub free_counter: u64,
}

#[event]
pub struct PremiumSubmissionEvent {
    pub contributor: Pubkey,
    pub block_number: i64,
    pub lattice_hash: [u8; 32],
    pub counter_b: u64,
    pub reward_eth: u64,
    pub reward_chg: u64,
    pub recursion_depth: u8,
    pub compute_units_used: u64,
    pub instance_id: u8,
}

#[event]
pub struct EpochSubmissionEvent {
    pub aggregator: Pubkey,
    pub epoch: u64,
    pub lattice_hash: [u8; 32],
    pub compute_units_total: u64,
    pub reward_eth: u64,
    pub reward_chg: u64,
}

#[event]
pub struct WithdrawalEthEvent {
    pub user: Pubkey,
    pub stealth_address: [u8; 32],
    pub amount: u64,
}

#[event]
pub struct WithdrawalChgEvent {
    pub user: Pubkey,
    pub stealth_address: [u8; 32],
    pub amount: u64,
}

#[derive(Copy, Clone, Pod, Zeroable)]
#[repr(C)]
struct ZKPData {
    a: [u64; 2],
    b: [[u64; 2]; 2],
    c: [u64; 2],
    zkp_input: [u64; 4],
    commitment: [u64; 2],
}

#[program]
pub mod zk_verifier {
    use super::*;
    pub fn verify_proof(_ctx: Context<VerifyProof>, _a: [u64; 2], _b: [[u64; 2]; 2], _c: [u64; 2], _zkp_input: [u64; 4], _commitment: [u64; 2]) -> Result<()> {
        Ok(()) // Placeholder for actual ZKP verification
    }
}

#[derive(Accounts)]
pub struct VerifyProof<'info> {
    pub global: Account<'info, GlobalState>,
}