use crate::primary_market::balance_prover::run_setup::{
    check_keys_exist, run_trusted_setup_automated, run_trusted_setup_interactive, validate_keys,
};
use crate::primary_market::balance_prover::trusted_setup::{
    load_proving_key_from_disk, load_verifying_key_from_disk,
};
use crate::secondary_market::matching_engine::OrderPool;
use ark_bn254::Bn254;
use ark_groth16::{ProvingKey, VerifyingKey};
use solana_client::rpc_client::RpcClient;
use solana_sdk::{commitment_config::CommitmentConfig, pubkey::Pubkey};
use std::str::FromStr;
use std::sync::Arc;

#[derive(Clone)]
pub struct AppState {
    pub rpc_client: Arc<RpcClient>,
    pub program_id: Pubkey,
    pub order_pool: OrderPool,
    // ZK-SNARK keys (optional - None if keys not available)
    pub proving_key: Option<Arc<ProvingKey<Bn254>>>,
    pub verifying_key: Option<Arc<VerifyingKey<Bn254>>>,
}

impl AppState {
    /// Creates new AppState with ZK setup handling
    pub fn new(
        rpc_url: &str,
        program_id_str: &str,
        skip_zk_setup: bool,
        force_zk_setup: bool,
    ) -> Result<Self, String> {
        let rpc_client =
            RpcClient::new_with_commitment(rpc_url.to_string(), CommitmentConfig::finalized());

        let program_id =
            Pubkey::from_str(program_id_str).map_err(|e| format!("Invalid program ID: {}", e))?;

        // Initialize the order pool (in-memory orderbook)
        let order_pool = OrderPool::new();

        // Handle ZK-SNARK setup
        let (proving_key, verifying_key) = if skip_zk_setup {
            println!("ZK setup skipped (SKIP_ZK_SETUP=true)");
            if !check_keys_exist() {
                println!("WARNING: ZK keys do not exist! Primary market features will not work.");
                println!(
                    "Run without SKIP_ZK_SETUP=true or set FORCE_ZK_SETUP=true to generate keys."
                );
            }
            (None, None)
        } else {
            Self::setup_zk_keys(force_zk_setup)?
        };

        Ok(AppState {
            rpc_client: Arc::new(rpc_client),
            program_id,
            order_pool,
            proving_key,
            verifying_key,
        })
    }

    /// Handles ZK-SNARK setup and key loading
    fn setup_zk_keys(
        force_setup: bool,
    ) -> Result<
        (
            Option<Arc<ProvingKey<Bn254>>>,
            Option<Arc<VerifyingKey<Bn254>>>,
        ),
        String,
    > {
        // Check if running in interactive mode (terminal)
        let is_interactive = atty::is(atty::Stream::Stdin);

        // Run trusted setup if needed
        if is_interactive {
            // Interactive mode: prompt user
            match run_trusted_setup_interactive() {
                Ok(true) => println!("✓ Trusted setup completed"),
                Ok(false) => println!("✓ Using existing keys"),
                Err(e) => {
                    println!("Trusted setup failed: {}", e);
                    println!("Primary market features will not be available");
                    return Ok((None, None));
                }
            }
        } else {
            // Non-interactive mode: automated
            match run_trusted_setup_automated(force_setup) {
                Ok(true) => println!("✓ Trusted setup completed"),
                Ok(false) => println!("✓ Using existing keys"),
                Err(e) => {
                    println!("Trusted setup failed: {}", e);
                    println!("Primary market features will not be available");
                    return Ok((None, None));
                }
            }
        }

        // Validate keys exist
        if !check_keys_exist() {
            println!("WARNING: ZK keys do not exist after setup!");
            return Ok((None, None));
        }

        // Validate keys can be loaded
        if let Err(e) = validate_keys() {
            println!("Key validation failed: {}", e);
            println!("Primary market features will not be available");
            return Ok((None, None));
        }

        // Load keys into memory
        println!("Loading ZK-SNARK keys into memory...");

        let pk = load_proving_key_from_disk()
            .map_err(|e| format!("Failed to load proving key: {}", e))?;

        let vk = load_verifying_key_from_disk()
            .map_err(|e| format!("Failed to load verifying key: {}", e))?;

        println!("✓ ZK-SNARK keys loaded into AppState");
        println!("  Proving key: {} bytes", std::mem::size_of_val(&pk));
        println!("  Verifying key: {} bytes", std::mem::size_of_val(&vk));

        Ok((Some(Arc::new(pk)), Some(Arc::new(vk))))
    }

    /// Checks if ZK keys are available
    pub fn has_zk_keys(&self) -> bool {
        self.proving_key.is_some() && self.verifying_key.is_some()
    }

    /// Gets proving key reference (returns error if not available)
    pub fn get_proving_key(&self) -> Result<&ProvingKey<Bn254>, String> {
        self.proving_key
            .as_ref()
            .map(|arc| arc.as_ref())
            .ok_or_else(|| {
                "Proving key not available. ZK setup may have failed or been skipped.".to_string()
            })
    }

    /// Gets verifying key reference (returns error if not available)
    pub fn get_verifying_key(&self) -> Result<&VerifyingKey<Bn254>, String> {
        self.verifying_key
            .as_ref()
            .map(|arc| arc.as_ref())
            .ok_or_else(|| {
                "Verifying key not available. ZK setup may have failed or been skipped.".to_string()
            })
    }
}