use anchor_lang::prelude::*;
use anchor_lang::solana_program::program::invoke;
use anchor_lang::solana_program::instruction::Instruction;

pub mod state;
pub mod instructions;
pub mod error;
pub mod constants;

use state::*;
use instructions::*;
use error::HdGLError;
use constants::*;

declare_id!("HdGLMarketplace1111111111111111111111111111");

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

    /// Initializes the global state PDA.
    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(())
    }

    /// Free tick progression.
    pub fn free_tick(ctx: Context<FreeTick>) -> Result<()> {
        let counters = &mut ctx.accounts.counters;
        counters.free_counter += 1;
        emit!(FreeTickEvent {
            free_counter: counters.free_counter,
        });
        Ok(())
    }

    /// Premium submission with ZKP.
    pub fn premium_submit(
        ctx: Context<PremiumSubmit>,
        lattice_hash: [u8; 32],
        recursion_depth: u8,
        compute_units_used: u64,
        instance_id: u8,
        a: [[u64; 2]; 2],
        b: [[[u64; 2]; 2]; 2],
        c: [u64; 2],
        zkp_input: [u64; 4],
        commitment: [u64; 2],
    ) -> Result<()> {
        require!(instance_id < 8, HdGLError::InvalidInstanceId);
        require!(compute_units_used < 1e18 as u64, HdGLError::Overflow);

        // CPI to zk_verifier for ZKP validation
        let ix = Instruction {
            program_id: ctx.accounts.zk_verifier.key(),
            accounts: vec![
                AccountMeta::new_readonly(ctx.accounts.zk_verifier.key(), false),
                AccountMeta::new(ctx.accounts.global.key(), false),
            ],
            data: vec![
                // Serialize a, b, c, zkp_input, commitment
                (a[0][0] as u8), (a[0][1] as u8), (a[1][0] as u8), (a[1][1] as u8),
                (b[0][0][0] as u8), (b[0][0][1] as u8), (b[0][1][0] as u8), (b[0][1][1] as u8),
                (b[1][0][0] as u8), (b[1][0][1] as u8), (b[1][1][0] as u8), (b[1][1][1] as u8),
                (c[0] as u8), (c[1] as u8),
                (zkp_input[0] as u8), (zkp_input[1] as u8), (zkp_input[2] as u8), (zkp_input[3] as u8),
                (commitment[0] as u8), (commitment[1] as u8),
            ],
        };
        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(())
    }

    // Similar implementations for batch_premium_submit, epoch_submit, withdraw_eth, withdraw_chg, register_stealth_address

    fn _compute_reward(compute_units_used: u64) -> (u64, u64) {
        let base_eth = 1e15 as u64 * 1e6 as u64;
        let base_chg = 1000 as u64 * 1e6 as u64;
        let denominator = compute_units_used as u128 + 1;
        let reward_eth = (base_eth as u128 / denominator as u128) as u64;
        let reward_chg = (base_chg as u128 / denominator as u128) as u64;
        (reward_eth, reward_chg)
    }
}

#[derive(Accounts)]
pub struct Initialize<'info> {
    #[account(init, payer = owner, space = 8 + 32)]
    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)]
    pub counters: Account<'info, Counters>,
    pub global: Account<'info, GlobalState>,
}

#[derive(Accounts)]
pub struct PremiumSubmit<'info> {
    #[account(mut)]
    pub counters: Account<'info, Counters>,
    #[account(mut)]
    pub global: Account<'info, GlobalState>,
    #[account(mut)]
    pub contributor: Account<'info, ContributorState>,
    pub zk_verifier: Program<'info, ZKVerifier>,
    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,
}

// State structs
#[account]
pub struct GlobalState {
    pub owner: Pubkey,
    pub chg_token: Pubkey,
    pub zk_verifier: Pubkey,
    pub last_snapshot: [u8; 32],
    pub current_epoch: u64,
}

#[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,
}