use aes_gcm::{Aes256Gcm, KeyInit};
use blstrs::{G1Projective, G2Projective, Scalar};
use elliptic_curve::{Field as _, Group as _};
use near_mpc_contract_interface::types::ProtocolContractState;
use near_mpc_contract_interface::types::{
    BitcoinExtractor, BitcoinRpcRequest, EDDSA_PAYLOAD_SIZE_LOWER_BOUND_BYTES,
    EDDSA_PAYLOAD_SIZE_UPPER_BOUND_BYTES, ForeignChainRpcRequest, ForeignTxPayloadVersion,
    VerifyForeignTransactionRequestArgs,
};
use rand::rngs::OsRng;
use std::collections::BTreeMap;
use std::net::{Ipv4Addr, SocketAddr};

use tokio::sync::{RwLock, watch};

use crate::config::{ParticipantsConfig, PersistentSecrets, SecretsConfig};
use crate::coordinator::Coordinator;
use crate::db::SecretDB;
use crate::indexer::IndexerAPI;
use crate::indexer::fake::{FakeIndexerManager, FakeReadSupportedForeignChain};
use crate::indexer::handler::{
    CKDArgs, CKDRequestFromChain, SignArgs, SignatureRequestFromChain,
    VerifyForeignTxRequestFromChain,
};
use crate::keyshare::{KeyStorageConfig, Keyshare};
use crate::migration_service::spawn_recovery_server_and_run_onboarding;
use crate::p2p::testing::{PortSeed, generate_test_p2p_configs};
use mpc_node_config::{
    CKDConfig, ConfigFile, ForeignChainsConfig, IndexerConfig, KeygenConfig, PresignatureConfig,
    SignatureConfig, SyncMode, TripleConfig,
};

use crate::primitives::ParticipantId;
use crate::tests::common::MockTransactionSender;
use crate::tracking::{self, AutoAbortTask, start_root_task};
use crate::web::recent_transactions::SharedRecentTransactions;
use crate::web::{start_web_server, static_web_data};
use assert_matches::assert_matches;
use mpc_primitives::domain::{Curve, Protocol};
use near_account_id::AccountId;
use near_indexer_primitives::CryptoHash;
use near_indexer_primitives::types::Finality;
use near_mpc_bounded_collections::BoundedVec;
use near_mpc_contract_interface::types::DomainConfig;
use near_mpc_contract_interface::types::Payload;
use near_time::Clock;
use rand::{Rng, RngCore};
use std::path::{Path, PathBuf};
use std::str::FromStr;
use std::sync::{Arc, OnceLock};
use tokio::time::timeout;

pub mod common;
pub(crate) mod dto_conversions;

mod asset_generation_signing_contention;
mod basic_cluster;
mod changing_participant_details;
mod faulty;
mod foreign_chain_configuration;
mod multidomain;
mod onboarding;
mod protocol_yielding;
mod resharing;

const DEFAULT_BLOCK_TIME: std::time::Duration = std::time::Duration::from_millis(300);
const DEFAULT_MAX_PROTOCOL_WAIT_TIME: std::time::Duration = std::time::Duration::from_secs(60);
const DEFAULT_MAX_SIGNATURE_WAIT_TIME: std::time::Duration = std::time::Duration::from_secs(60);

/// Data needed to start running a test node.
pub struct OneNodeTestConfig {
    clock: Clock,
    home_dir: PathBuf,
    pub config: ConfigFile,
    secrets: SecretsConfig,
    indexer: IndexerAPI<MockTransactionSender, FakeReadSupportedForeignChain>,
    _indexer_task: AutoAbortTask<()>,
    currently_running_job_name: Arc<std::sync::Mutex<String>>,
}

pub fn make_key_storage_config(
    home_dir: PathBuf,
    local_encryption_key: [u8; 16],
) -> KeyStorageConfig {
    KeyStorageConfig {
        home_dir,
        local_encryption_key,
        gcp: None,
    }
}

pub async fn get_keyshares(
    home_dir: PathBuf,
    local_encryption_key: [u8; 16],
    keyset: &near_mpc_contract_interface::types::Keyset,
) -> anyhow::Result<Vec<Keyshare>> {
    let key_storage_config = make_key_storage_config(home_dir, local_encryption_key);
    let keystore = key_storage_config.create().await.unwrap();
    keystore.get_keyshares(keyset).await
}

impl OneNodeTestConfig {
    pub async fn run(self) -> anyhow::Result<()> {
        std::fs::create_dir_all(&self.home_dir)?;
        let my_account_id = self.config.my_near_account_id.clone();
        async move {
            let root_future = async move {
                let root_task_handle = tracking::current_task();
                let root_task = OnceLock::new();
                let _ = root_task.set(root_task_handle);
                let (debug_request_sender, _) = tokio::sync::broadcast::channel(10);

                let (_, dummy_protocol_state_receiver) =
                    watch::channel(ProtocolContractState::NotInitialized);
                let (_, dummy_migration_state_receiver) = watch::channel((0, BTreeMap::new()));
                // The fake indexer never records, so the buffer stays empty.
                // TODO(#3522): wire it into the fake indexer and assert on it.
                let web_server = start_web_server(
                    root_task.into(),
                    debug_request_sender.clone(),
                    self.config.web_ui,
                    static_web_data(&self.secrets, None),
                    dummy_protocol_state_receiver,
                    dummy_migration_state_receiver,
                    self.config.clone(),
                    SharedRecentTransactions::default(),
                )
                .await?;
                let _web_server = tracking::spawn_checked("web server", web_server);

                let secret_db = SecretDB::new(&self.home_dir, self.secrets.local_storage_aes_key)?;
                let key_storage_config =
                    make_key_storage_config(self.home_dir, self.secrets.local_storage_aes_key);
                let keystore = Arc::new(RwLock::new(
                    key_storage_config
                        .create()
                        .await
                        .expect("require keystore for integration tests"),
                ));

                spawn_recovery_server_and_run_onboarding(
                    self.config.migration_web_ui,
                    (&self.secrets).into(),
                    self.config.my_near_account_id.clone(),
                    keystore.clone(),
                    self.indexer.my_migration_info_receiver.clone(),
                    self.indexer.contract_state_receiver.clone(),
                    self.indexer.txn_sender.clone(),
                )
                .await
                .unwrap();

                let coordinator = Coordinator {
                    clock: self.clock,
                    config_file: self.config,
                    secrets: self.secrets,
                    secret_db,
                    keyshare_storage: RwLock::new(key_storage_config.create().await.unwrap())
                        .into(),
                    indexer: self.indexer,
                    currently_running_job_name: self.currently_running_job_name,
                    debug_request_sender,
                };
                coordinator.run().await
            };
            start_root_task(&format!("root for {}", my_account_id), root_future)
                .0
                .await?;
            Ok(())
        }
        .await
    }
}

/// Test fixture for integration tests, includes a fake indexer and a set of
/// nodes.
pub struct IntegrationTestSetup {
    pub indexer: FakeIndexerManager,
    pub configs: Vec<OneNodeTestConfig>,
    pub participants: ParticipantsConfig,
}

impl IntegrationTestSetup {
    /// Generates test node configs and a fake indexer; each config can then be used
    /// to start running the node.
    pub fn new(
        clock: Clock,
        temp_dir: &Path,
        participant_accounts: Vec<AccountId>,
        threshold: usize,
        txn_delay_blocks: u64,
        port_seed: PortSeed,
        block_time: std::time::Duration,
    ) -> IntegrationTestSetup {
        let p2p_configs =
            generate_test_p2p_configs(&participant_accounts, threshold, port_seed, None).unwrap();
        let participants = p2p_configs[0].0.participants.clone();
        let mut indexer_manager =
            FakeIndexerManager::new(clock.clone(), txn_delay_blocks, block_time);

        let mut configs = Vec::new();
        for (i, (_, p2p_key)) in p2p_configs.into_iter().enumerate() {
            let config = ConfigFile {
                cores: Some(4),
                separate_asset_generation_runtime: true,
                // Indexer config is just a dummy.
                indexer: IndexerConfig {
                    concurrency: 1.try_into().unwrap(),
                    finality: Finality::Final,
                    mpc_contract_id: "test".parse().unwrap(),
                    port_override: None,
                    sync_mode: SyncMode::Latest,
                    validate_genesis: false,
                },
                keygen: KeygenConfig { timeout_sec: 60 },
                my_near_account_id: participant_accounts[i].clone(),
                // Don't care since we use the fake indexer
                near_responder_account_id: AccountId::from_str("dont_care").unwrap(),
                presignature: PresignatureConfig {
                    concurrency: 1,
                    desired_presignatures_to_buffer: 5,
                    timeout_sec: 60,
                },
                signature: SignatureConfig { timeout_sec: 60 },
                ckd: CKDConfig { timeout_sec: 60 },
                foreign_chains: ForeignChainsConfig::default(),
                triple: TripleConfig {
                    concurrency: 1,
                    desired_triples_to_buffer: 10,
                    parallel_triple_generation_stagger_time_sec: 1,
                    timeout_sec: 60,
                },
                number_of_responder_keys: 0,
                web_ui: SocketAddr::new(Ipv4Addr::UNSPECIFIED.into(), port_seed.web_port(i)),
                migration_web_ui: SocketAddr::new(
                    Ipv4Addr::UNSPECIFIED.into(),
                    port_seed.migration_web_port(i),
                ),
                pprof_bind_address: SocketAddr::new(
                    Ipv4Addr::UNSPECIFIED.into(),
                    port_seed.pprof_web_port(i),
                ),
            };
            let secrets = SecretsConfig {
                persistent_secrets: PersistentSecrets {
                    p2p_private_key: p2p_key.clone(),
                    near_signer_key: ed25519_dalek::SigningKey::generate(&mut OsRng),
                    near_responder_keys: vec![ed25519_dalek::SigningKey::generate(&mut OsRng)],
                },
                local_storage_aes_key: rand::random(),
                backup_encryption_key: Aes256Gcm::generate_key(OsRng).into(),
            };
            let (indexer_api, task, currently_running_job_name) = indexer_manager.add_indexer_node(
                i.into(),
                participant_accounts[i].clone(),
                p2p_key.verifying_key(),
            );
            configs.push(OneNodeTestConfig {
                clock: clock.clone(),
                config,
                home_dir: temp_dir.join(format!("{}", i)),
                secrets,
                indexer: indexer_api,
                _indexer_task: task,
                currently_running_job_name,
            });
        }
        IntegrationTestSetup {
            indexer: indexer_manager,
            configs,
            participants,
        }
    }

    pub fn get_config(&self, node_index: usize) -> Option<&OneNodeTestConfig> {
        self.configs.get(node_index)
    }
}
/// Request a signature from the indexer and wait for the response.
/// Returns the time taken to receive the response, or None if timed out.
pub async fn request_signature_and_await_response(
    indexer: &mut FakeIndexerManager,
    user: &str,
    domain: &DomainConfig,
    timeout_sec: std::time::Duration,
) -> Option<std::time::Duration> {
    let payload = match domain.protocol {
        Protocol::CaitSith | Protocol::DamgardEtAl => {
            let mut payload = [0; 32];
            rand::thread_rng().fill_bytes(payload.as_mut());

            Payload::Ecdsa(payload.into())
        }
        Protocol::Frost => {
            let len = rand::thread_rng().gen_range(
                EDDSA_PAYLOAD_SIZE_LOWER_BOUND_BYTES..EDDSA_PAYLOAD_SIZE_UPPER_BOUND_BYTES,
            );
            let mut payload = vec![0; len];
            rand::thread_rng().fill_bytes(payload.as_mut());

            let bounded_payload: BoundedVec<
                u8,
                EDDSA_PAYLOAD_SIZE_LOWER_BOUND_BYTES,
                EDDSA_PAYLOAD_SIZE_UPPER_BOUND_BYTES,
            > = payload.try_into().unwrap();

            Payload::Eddsa(bounded_payload)
        }
        Protocol::ConfidentialKeyDerivation => unreachable!(),
    };
    let request = SignatureRequestFromChain {
        signature_id: CryptoHash(rand::random()),
        receipt_id: CryptoHash(rand::random()),
        predecessor_id: user.parse().unwrap(),
        request: SignArgs {
            domain_id: domain.id,
            path: "m/44'/60'/0'/0/0".to_string(),
            payload,
        },
    };
    tracing::info!(
        "Sending signature request from user {}, payload {:?}",
        user,
        request.request.payload
    );
    indexer.request_signature(request.clone());
    let start_time = std::time::Instant::now();
    loop {
        match timeout(timeout_sec, indexer.next_response()).await {
            Ok(signature) => {
                if signature.request.payload != request.request.payload {
                    // This can legitimately happen when multiple nodes submit responses
                    // for the same signature request. In tests this can happen if the
                    // secondary leader thinks the primary leader is offline when in fact
                    // the network just has not yet been fully established.
                    tracing::info!(
                        "Received signature is not for the signature request we sent (user {})
                         Expected {:?}, actual {:?}",
                        user,
                        request.request.payload,
                        signature.request.payload
                    );
                    continue;
                }
                if signature.request.domain_id != domain.id {
                    tracing::info!(
                        "Received signature is not for the domain we requested (user {})
                         Expected {:?}, actual {:?}",
                        user,
                        domain.id,
                        signature.request.domain_id
                    );
                    continue;
                }
                tracing::info!("Got signature response for user {}", user);
                return Some(start_time.elapsed());
            }
            Err(_) => {
                tracing::info!("Timed out waiting for signature response for user {}", user);
                return None;
            }
        }
    }
}

/// Request a ckd from the indexer and wait for the response.
/// Returns the time taken to receive the response, or None if timed out.
pub async fn request_ckd_and_await_response(
    indexer: &mut FakeIndexerManager,
    user: &str,
    domain: &DomainConfig,
    timeout_sec: std::time::Duration,
) -> Option<std::time::Duration> {
    let app_public_key = near_mpc_contract_interface::types::CKDAppPublicKey::AppPublicKey(
        "bls12381g1:6KtVVcAAGacrjNGePN8bp3KV6fYGrw1rFsyc7cVJCqR16Zc2ZFg3HX3hSZxSfv1oH6"
            .parse()
            .unwrap(),
    );
    do_request_ckd_and_await_response(indexer, user, domain, timeout_sec, app_public_key).await
}

/// Request a ckd with public verifiability from the indexer and wait for the response.
/// Returns the time taken to receive the response, or None if timed out.
pub async fn request_ckd_pv_and_await_response(
    indexer: &mut FakeIndexerManager,
    user: &str,
    domain: &DomainConfig,
    timeout_sec: std::time::Duration,
) -> Option<std::time::Duration> {
    let app_sk = Scalar::random(&mut OsRng);
    let pk1 = G1Projective::generator() * app_sk;
    let pk2 = G2Projective::generator() * app_sk;
    let app_public_key = near_mpc_contract_interface::types::CKDAppPublicKey::AppPublicKeyPV(
        near_mpc_contract_interface::types::CKDAppPublicKeyPV {
            pk1: (&pk1).into(),
            pk2: (&pk2).into(),
        },
    );
    do_request_ckd_and_await_response(indexer, user, domain, timeout_sec, app_public_key).await
}

async fn do_request_ckd_and_await_response(
    indexer: &mut FakeIndexerManager,
    user: &str,
    domain: &DomainConfig,
    timeout_sec: std::time::Duration,
    app_public_key: near_mpc_contract_interface::types::CKDAppPublicKey,
) -> Option<std::time::Duration> {
    assert_matches!(
        Curve::from(domain.protocol),
        Curve::Bls12381,
        "`request_ckd_and_await_response` must be called with a compatible domain",
    );
    let request = CKDRequestFromChain {
        ckd_id: CryptoHash(rand::random()),
        receipt_id: CryptoHash(rand::random()),
        request: CKDArgs {
            app_public_key,
            domain_id: domain.id,
            app_id: [1u8; 32].into(),
        },
    };
    tracing::info!(
        "Sending ckd request from user {}, app_id {:?}, app_public_key {:?}",
        user,
        request.request.app_id,
        request.request.app_public_key,
    );
    indexer.request_ckd(request.clone());
    let start_time = std::time::Instant::now();
    loop {
        match timeout(timeout_sec, indexer.next_response_ckd()).await {
            Ok(ckd_response_args) => {
                if ckd_response_args.request.app_id != request.request.app_id {
                    // This can legitimately happen when multiple nodes submit responses
                    // for the same ckd request. In tests this can happen if the
                    // secondary leader thinks the primary leader is offline when in fact
                    // the network just has not yet been fully established.
                    tracing::info!(
                        "Received ckd is not for the ckd request we sent (user {})
                         Expected {:?}, actual {:?}",
                        user,
                        request.request.app_id,
                        ckd_response_args.request.app_id
                    );
                    continue;
                }
                if ckd_response_args.request.domain_id != domain.id {
                    tracing::info!(
                        "Received ckd is not for the domain we requested (user {})
                         Expected {:?}, actual {:?}",
                        user,
                        domain.id,
                        ckd_response_args.request.domain_id
                    );
                    continue;
                }
                if ckd_response_args.request.app_public_key != request.request.app_public_key {
                    tracing::info!(
                        "Received ckd is not for the app_public_key we requested (user {})
                         Expected {:?}, actual {:?}",
                        user,
                        request.request.app_public_key,
                        ckd_response_args.request.app_public_key
                    );
                    continue;
                }
                tracing::info!("Got response ckd for user {}", user);
                return Some(start_time.elapsed());
            }
            Err(_) => {
                tracing::info!("Timed out waiting for ckd response for user {}", user);
                return None;
            }
        }
    }
}

/// Request a verify foreign tx from the indexer and wait for the response.
/// Returns the time taken to receive the response, or None if timed out.
// TODO: remove this when tests are added for this functionality
#[expect(unused)]
pub async fn request_verify_foreign_tx_and_await_response(
    indexer: &mut FakeIndexerManager,
    user: &str,
    domain: &DomainConfig,
    timeout_sec: std::time::Duration,
) -> Option<std::time::Duration> {
    assert_matches!(
        Curve::from(domain.protocol),
        Curve::Secp256k1,
        "`request_verify_foreign_tx_and_await_response` must be called with a compatible domain",
    );
    let request = VerifyForeignTxRequestFromChain {
        verify_foreign_tx_id: CryptoHash(rand::random()),
        receipt_id: CryptoHash(rand::random()),
        request: VerifyForeignTransactionRequestArgs {
            request: ForeignChainRpcRequest::Bitcoin(BitcoinRpcRequest {
                tx_id: [42u8; 32].into(),
                confirmations: 2.into(),
                extractors: vec![BitcoinExtractor::BlockHash],
            }),
            domain_id: domain.id.0.into(),
            payload_version: ForeignTxPayloadVersion::V1,
        },
    };
    tracing::info!(
        "Sending verify foreign tx request from user {}, request {:?}",
        user,
        request.request,
    );
    indexer.request_verify_foreign_tx(request.clone());
    let start_time = std::time::Instant::now();
    loop {
        match timeout(timeout_sec, indexer.next_response_verify_foreign_tx()).await {
            Ok(verify_foreign_tx_response_args) => {
                if verify_foreign_tx_response_args.request.request != request.request.request {
                    // This can legitimately happen when multiple nodes submit responses
                    // for the same verify foreign tx request. In tests this can happen if the
                    // secondary leader thinks the primary leader is offline when in fact
                    // the network just has not yet been fully established.
                    tracing::info!(
                        "Received verify foreign tx is not for the verify foreign tx request we sent (user {})
                         Expected {:?}, actual {:?}",
                        user,
                        request.request.request,
                        verify_foreign_tx_response_args.request.request
                    );
                    continue;
                }
                if verify_foreign_tx_response_args.request.domain_id.0 != domain.id.0 {
                    tracing::info!(
                        "Received verify foreign tx is not for the domain we requested (user {})
                         Expected {:?}, actual {:?}",
                        user,
                        domain.id,
                        verify_foreign_tx_response_args.request.domain_id
                    );
                    continue;
                }
                tracing::info!("Got response verify foreign tx for user {}", user);
                return Some(start_time.elapsed());
            }
            Err(_) => {
                tracing::info!(
                    "Timed out waiting for verify foreign tx response for user {}",
                    user
                );
                return None;
            }
        }
    }
}

pub fn into_participant_ids(
    participants: &[threshold_signatures::participants::Participant],
) -> Vec<ParticipantId> {
    participants.iter().map(|p| (*p).into()).collect()
}

#[test]
fn test_build_info_metric() {
    // Test that the build info metric can be initialized without panicking
    crate::metrics::init_build_info_metric();

    // Verify that the built crate information is available
    let version = crate::built_info::PKG_VERSION;
    let build_time = crate::built_info::BUILT_TIME_UTC;
    let commit = crate::built_info::GIT_COMMIT_HASH_SHORT.unwrap_or("unknown");
    let rustc_version = crate::built_info::RUSTC_VERSION;

    // Verify that the version information is not "unknown"
    assert_ne!(version, "unknown", "PKG_VERSION should be set");
    assert_ne!(build_time, "unknown", "BUILT_TIME_UTC should be set");
    assert_ne!(commit, "unknown", "GIT_COMMIT_HASH_SHORT should be set");
    assert_ne!(rustc_version, "unknown", "RUSTC_VERSION should be set");

    // Verify that the version string contains all the expected information
    let version_string = &*crate::MPC_VERSION_STRING;
    assert!(version_string.contains(version));
    assert!(version_string.contains(build_time));
    assert!(version_string.contains(commit));
    assert!(version_string.contains(rustc_version));
}

#[test]
fn test_build_info_metric_initialization() {
    // Test that the build info metric can be initialized without panicking
    // This verifies that the compile-time constants are accessible
    crate::metrics::init_build_info_metric();

    // If we get here without panicking, the test passes
    // The metric should now be available in Prometheus with the correct values
}

#[test]
fn test_build_info_metric_values() {
    // Test that the build info metric has the correct values
    crate::metrics::init_build_info_metric();

    // Get the metric value directly
    let metric = &crate::metrics::MPC_BUILD_INFO;
    let version = crate::built_info::PKG_VERSION;
    let build_time = crate::built_info::BUILT_TIME_UTC;
    let commit = crate::built_info::GIT_COMMIT_HASH_SHORT.unwrap_or("unknown");
    let rustc_version = crate::built_info::RUSTC_VERSION;

    // Check that the metric exists with the correct labels
    let gauge = metric.with_label_values(&[version, build_time, commit, rustc_version]);
    let value = gauge.get();

    println!("Metric value: {}", value);
    println!(
        "Expected labels: version={}, build_time={}, commit={}, rustc_version={}",
        version, build_time, commit, rustc_version
    );

    // The value should be 1
    assert_eq!(value, 1);
}