use super::consts::DEFAULT_MAX_TIMEOUT_TX_INCLUDED;
use super::shared_key_utils::{
    DomainKey, SharedSecretKey, derive_secret_key_ed25519, derive_secret_key_secp256k1,
    generate_random_app_public_key,
};
use digest::{Digest, FixedOutput};
use k256::{
    FieldBytes,
    elliptic_curve::{Field as _, Group as _},
};
use mpc_contract::{
    errors,
    primitives::{
        ckd::CKDRequest,
        signature::{SignatureRequest, YieldIndex},
    },
};
use near_account_id::AccountId;
use near_mpc_bounded_collections::BoundedVec;
use near_mpc_contract_interface::method_names::{
    GET_PENDING_CKD_REQUEST, GET_PENDING_REQUEST, REQUEST_APP_PRIVATE_KEY, RESPOND, RESPOND_CKD,
    SIGN,
};
use near_mpc_contract_interface::types::kdf::{derive_app_id, derive_tweak};
use near_mpc_contract_interface::types::{
    self as dtos, CKDAppPublicKey, CKDAppPublicKeyPV, CKDRequestArgs,
};
use near_mpc_contract_interface::types::{CKDResponse, DomainId, Payload};
use near_mpc_sdk::sign::{Ed25519Signature, K256Signature};
use near_mpc_sdk::sign::{SignRequestArgs, SignRequestBuilder, SignatureRequestResponse};
use near_workspaces::{
    Account, Contract, Worker, network::Sandbox, operations::TransactionStatus, types::NearToken,
};
use rand::{Rng, rngs::OsRng};
use rand_core::CryptoRngCore;
use serde::Serialize;
use sha2::Sha256;
use signature::DigestSigner;
use std::time::Duration;
use threshold_signatures::{
    KeygenOutput,
    confidential_key_derivation::{self as ckd, BLS12381SHA256, hash_app_id_with_pk},
    ecdsa as ts_ecdsa,
    frost::eddsa,
    frost_ed25519::{self},
};

#[derive(Debug)]
#[expect(clippy::large_enum_variant)]
pub enum DomainResponseTest {
    Sign(SignRequestTest),
    CKD(CKDRequestTest),
}

impl DomainResponseTest {
    pub fn new(
        rng: &mut impl CryptoRngCore,
        domain_key: &DomainKey,
        predecessor_id: &AccountId,
    ) -> Self {
        let domain_id = domain_key.domain_config.id;
        match &domain_key.domain_secret_key {
            SharedSecretKey::Secp256k1(sk) => DomainResponseTest::Sign(gen_secp_256k1_sign_test(
                rng,
                domain_id,
                predecessor_id,
                sk,
            )),
            SharedSecretKey::Ed25519(sk) => {
                DomainResponseTest::Sign(gen_ed25519_sign_test(rng, domain_id, predecessor_id, sk))
            }
            SharedSecretKey::Bls12381(sk) => {
                DomainResponseTest::CKD(CKDRequestTest::new(rng, domain_id, predecessor_id, sk))
            }
        }
    }

    pub async fn run(
        &self,
        account: &Account,
        contract: &Contract,
        attested_account: &Account,
    ) -> anyhow::Result<()> {
        let status = self
            .submit_request_ensure_included(account, contract)
            .await?;
        self.submit_response(contract, attested_account).await?;
        self.verify_execution_outcome(status).await?;
        Ok(())
    }

    pub async fn submit_response(
        &self,
        contract: &Contract,
        attested_account: &Account,
    ) -> anyhow::Result<()> {
        match &self {
            Self::Sign(inner) => {
                submit_signature_response(&inner.response, contract, attested_account).await
            }
            Self::CKD(inner) => {
                submit_ckd_response(&inner.response, contract, attested_account).await
            }
        }
    }

    pub async fn submit_request_ensure_included(
        &self,
        account: &Account,
        contract: &Contract,
    ) -> anyhow::Result<TransactionStatus> {
        match self {
            Self::Sign(inner) => {
                let status = submit_sign_request(account, &inner.args, contract).await?;
                await_request_in_contract_queue(contract, &inner.response.request, None).await?;
                Ok(status)
            }
            Self::CKD(inner) => {
                let status = submit_ckd_request(account, &inner.args, contract).await?;
                await_request_in_contract_queue(contract, &inner.response.request, None).await?;
                Ok(status)
            }
        }
    }

    pub async fn verify_execution_outcome(&self, status: TransactionStatus) -> anyhow::Result<()> {
        match self {
            Self::Sign(inner) => inner.verify_execution_outcome(status).await,
            Self::CKD(inner) => inner.verify_execution_outcome(status).await,
        }
    }
}

#[derive(Debug)]
pub struct SignRequestTest {
    pub response: SignResponseArgs,
    pub args: SignRequestArgs,
}

impl SignRequestTest {
    pub fn request_json_args(&self) -> serde_json::Value {
        serde_json::json!({
            "request": self.args,
        })
    }

    pub fn payload(&self) -> &Payload {
        &self.response.request.payload
    }

    pub fn path(&self) -> &str {
        &self.args.path
    }

    pub async fn verify_execution_outcome(&self, status: TransactionStatus) -> anyhow::Result<()> {
        let execution = status.await?;
        dbg!(&execution);
        let execution = execution.into_result()?;
        let returned_resp: SignatureRequestResponse = execution.json()?;
        assert_eq!(
            returned_resp, self.response.response,
            "Returned signature request does not match"
        );
        Ok(())
    }

    pub async fn submit_response(
        &self,
        contract: &Contract,
        attested_account: &Account,
    ) -> anyhow::Result<()> {
        submit_signature_response(&self.response, contract, attested_account).await
    }
}

#[derive(Debug, Serialize)]
pub struct SignResponseArgs {
    pub request: SignatureRequest,
    pub response: SignatureRequestResponse,
}

impl SignResponseArgs {
    pub fn json_args(&self) -> serde_json::Value {
        serde_json::to_value(self).unwrap()
    }
}

pub async fn verify_timeout(status: TransactionStatus) -> anyhow::Result<()> {
    let execution = status.await?;
    dbg!(&execution);
    assert!(execution.is_failure());
    let err = execution
        .into_result()
        .expect_err("expect execution failure");
    assert!(
        err.to_string()
            .contains(&errors::RequestError::Timeout.to_string())
    );
    Ok(())
}

#[derive(Debug)]
pub struct CKDRequestTest {
    pub response: CKDResponseArgs,
    pub args: CKDRequestArgs,
}

fn gen_ckd_derivation_path(rng: &mut impl CryptoRngCore) -> String {
    let empty: bool = rng.r#gen();
    if empty {
        "".to_string()
    } else {
        rng.r#gen::<usize>().to_string()
    }
}

impl CKDRequestTest {
    pub fn new(
        rng: &mut impl CryptoRngCore,
        domain_id: DomainId,
        predecessor_id: &AccountId,
        sk: &ckd::KeygenOutput,
    ) -> CKDRequestTest {
        let derivation_path = gen_ckd_derivation_path(rng);
        let app_public_key = generate_random_app_public_key(rng);
        let (request, response) = create_response_ckd(
            predecessor_id,
            &app_public_key,
            &domain_id,
            sk,
            &derivation_path,
        );
        let args = CKDRequestArgs {
            derivation_path,
            app_public_key: CKDAppPublicKey::AppPublicKey(app_public_key),
            domain_id,
        };

        CKDRequestTest {
            response: CKDResponseArgs { request, response },
            args,
        }
    }
    pub fn new_pv(
        rng: &mut impl CryptoRngCore,
        domain_id: DomainId,
        predecessor_id: &AccountId,
        sk: &ckd::KeygenOutput,
    ) -> CKDRequestTest {
        let derivation_path = gen_ckd_derivation_path(rng);
        let (request, response) =
            create_response_ckd_pv(rng, predecessor_id, &domain_id, sk, &derivation_path);
        let args = CKDRequestArgs {
            derivation_path,
            app_public_key: request.app_public_key.clone(),
            domain_id,
        };

        CKDRequestTest {
            response: CKDResponseArgs { request, response },
            args,
        }
    }

    pub fn request_json_args(&self) -> serde_json::Value {
        serde_json::json!({
            "request": self.args,
        })
    }

    async fn verify_execution_outcome(&self, status: TransactionStatus) -> anyhow::Result<()> {
        let execution = status.await?;
        dbg!(&execution);
        let execution = execution.into_result()?;
        let returned_resp: CKDResponse = execution.json()?;
        assert_eq!(
            returned_resp, self.response.response,
            "Returned ckd request does not match"
        );
        Ok(())
    }
}

#[derive(Debug, Serialize)]
pub struct CKDResponseArgs {
    pub request: CKDRequest,
    pub response: CKDResponse,
}

async fn submit_request(
    account: &Account,
    contract: &Contract,
    method: &str,
    args: impl Serialize,
) -> anyhow::Result<TransactionStatus> {
    let status = account
        .call(contract.id(), method)
        .args_json(serde_json::json!({ "request": args }))
        .deposit(NearToken::from_yoctonear(1))
        .max_gas()
        .transact_async()
        .await?;
    dbg!(&status);
    Ok(status)
}

async fn submit_sign_request(
    account: &Account,
    request: &SignRequestArgs,
    contract: &Contract,
) -> anyhow::Result<TransactionStatus> {
    submit_request(account, contract, SIGN, request).await
}

async fn submit_ckd_request(
    account: &Account,
    request: &CKDRequestArgs,
    contract: &Contract,
) -> anyhow::Result<TransactionStatus> {
    submit_request(account, contract, REQUEST_APP_PRIVATE_KEY, request).await
}

async fn submit_response(
    contract: &Contract,
    attested_account: &Account,
    method: &str,
    args: impl Serialize,
) -> anyhow::Result<()> {
    let respond = attested_account
        .call(contract.id(), method)
        .args_json(args)
        .max_gas()
        .transact()
        .await?;
    dbg!(&respond);
    respond.into_result()?;
    Ok(())
}

pub async fn submit_signature_response(
    response: &SignResponseArgs,
    contract: &Contract,
    attested_account: &Account,
) -> anyhow::Result<()> {
    submit_response(contract, attested_account, RESPOND, response).await
}

pub async fn submit_ckd_response(
    response: &CKDResponseArgs,
    contract: &Contract,
    attested_account: &Account,
) -> anyhow::Result<()> {
    submit_response(contract, attested_account, RESPOND_CKD, response).await
}

pub async fn submit_ckd_response_measure_gas(
    response: &CKDResponseArgs,
    contract: &Contract,
    attested_account: &Account,
) -> anyhow::Result<near_workspaces::types::Gas> {
    let respond = attested_account
        .call(contract.id(), RESPOND_CKD)
        .args_json(response)
        .max_gas()
        .transact()
        .await?;
    let gas = respond.total_gas_burnt;
    respond.into_result()?;
    Ok(gas)
}

trait ContractQueueRequest: serde::Serialize + Sync {
    async fn is_in_queue(&self, contract: &Contract) -> Option<YieldIndex>;
}

impl ContractQueueRequest for CKDRequest {
    async fn is_in_queue(&self, contract: &Contract) -> Option<YieldIndex> {
        contract
            .view(GET_PENDING_CKD_REQUEST)
            .args_json(serde_json::json!({ "request": self }))
            .await
            .unwrap()
            .json()
            .unwrap()
    }
}

impl ContractQueueRequest for SignatureRequest {
    async fn is_in_queue(&self, contract: &Contract) -> Option<YieldIndex> {
        contract
            .view(GET_PENDING_REQUEST)
            .args_json(serde_json::json!({ "request": self }))
            .await
            .unwrap()
            .json()
            .unwrap()
    }
}

async fn await_request_in_contract_queue<T: ContractQueueRequest>(
    contract: &Contract,
    request: &T,
    max_timeout: Option<Duration>,
) -> anyhow::Result<()> {
    let timeout = max_timeout.unwrap_or(DEFAULT_MAX_TIMEOUT_TX_INCLUDED);
    let start = std::time::Instant::now();

    loop {
        if request.is_in_queue(contract).await.is_some() {
            return Ok(());
        }

        if start.elapsed() >= timeout {
            anyhow::bail!("timed out waiting for request to appear in queue");
        }

        tokio::time::sleep(Duration::from_millis(50)).await;
    }
}

/// Derives a confidential key following https://github.com/near/threshold-signatures/blob/main/docs/confidential_key_derivation.md
pub fn create_response_ckd(
    account_id: &AccountId,
    app_public_key: &dtos::Bls12381G1PublicKey,
    domain_id: &DomainId,
    key_package: &KeygenOutput<BLS12381SHA256>,
    derivation_path: &str,
) -> (CKDRequest, CKDResponse) {
    let request = CKDRequest::new(
        CKDAppPublicKey::AppPublicKey(app_public_key.clone()),
        *domain_id,
        account_id,
        derivation_path,
    );

    let app_id = derive_app_id(account_id, derivation_path);
    let app_pk: ckd::ElementG1 = app_public_key
        .try_into()
        .expect("invalid BLS12-381 G1 point");
    let msk = key_package.private_share.to_scalar();

    let big_s = hash_app_id_with_pk(&key_package.public_key, app_id.as_ref()) * msk;
    let y = ckd::Scalar::random(OsRng);
    let big_y = ckd::ElementG1::generator() * y;
    let big_c = big_s + app_pk * y;

    let response = CKDResponse {
        big_y: (&big_y).into(),
        big_c: (&big_c).into(),
    };
    (request, response)
}

fn create_response_ckd_pv(
    rng: &mut impl CryptoRngCore,
    account_id: &AccountId,
    domain_id: &DomainId,
    key_package: &KeygenOutput<BLS12381SHA256>,
    derivation_path: &str,
) -> (CKDRequest, CKDResponse) {
    let app_scalar = ckd::Scalar::random(&mut *rng);
    let app_pk1 = ckd::ElementG1::generator() * app_scalar;
    let app_pk2 = ckd::ElementG2::generator() * app_scalar;
    let app_public_key = CKDAppPublicKey::AppPublicKeyPV(CKDAppPublicKeyPV {
        pk1: dtos::Bls12381G1PublicKey::from(&app_pk1),
        pk2: dtos::Bls12381G2PublicKey::from(&app_pk2),
    });

    let request = CKDRequest::new(app_public_key, *domain_id, account_id, derivation_path);

    let app_id = derive_app_id(account_id, derivation_path);
    let msk = key_package.private_share.to_scalar();

    let big_s = hash_app_id_with_pk(&key_package.public_key, app_id.as_ref()) * msk;
    let y = ckd::Scalar::random(rng);
    let big_y = ckd::ElementG1::generator() * y;
    let big_c = big_s + app_pk1 * y;

    let response = CKDResponse {
        big_y: (&big_y).into(),
        big_c: (&big_c).into(),
    };
    (request, response)
}

pub async fn make_and_submit_requests(
    keys: &[DomainKey],
    contract: &Contract,
    worker: &Worker<Sandbox>,
    rng: &mut impl CryptoRngCore,
) -> (Vec<PendingSignRequest>, Vec<PendingCKDRequest>) {
    let mut pending_sign_requests = vec![];
    let mut pending_ckd_requests = vec![];

    const NUM_TESTS: usize = 2;

    let alice = worker.dev_create_account().await.unwrap();
    let alice_id = alice.id();

    for key in keys {
        for _ in 0..NUM_TESTS {
            match DomainResponseTest::new(rng, key, alice_id) {
                DomainResponseTest::Sign(inner) => {
                    let transaction = submit_sign_request(&alice, &inner.args, contract)
                        .await
                        .unwrap();
                    pending_sign_requests.push(PendingSignRequest {
                        transaction,
                        response: inner.response,
                    });
                }
                DomainResponseTest::CKD(inner) => {
                    let transaction = submit_ckd_request(&alice, &inner.args, contract)
                        .await
                        .unwrap();
                    pending_ckd_requests.push(PendingCKDRequest {
                        transaction,
                        ckd_response: inner.response,
                    });
                }
            }
        }
    }
    (pending_sign_requests, pending_ckd_requests)
}

pub fn create_response_secp256k1(
    domain_id: DomainId,
    predecessor_id: &AccountId,
    msg: &str,
    path: &str,
    signing_key: &ts_ecdsa::KeygenOutput,
) -> (Payload, SignatureRequest, SignatureRequestResponse) {
    let (digest, payload) = process_message(msg);
    let tweak = derive_tweak(predecessor_id, path);
    let derived_sk = derive_secret_key_secp256k1(signing_key, &tweak);
    let signing_key =
        k256::ecdsa::SigningKey::from_bytes(&derived_sk.private_share.to_scalar().into()).unwrap();

    let (sig, recovery_id) = signing_key.try_sign_digest(digest).unwrap();

    let signature = K256Signature::from_ecdsa_recoverable(&sig, recovery_id);
    let respond_req = SignatureRequest::new(domain_id, payload.clone(), predecessor_id, path);

    (
        payload,
        respond_req,
        SignatureRequestResponse::Secp256k1(signature),
    )
}

pub fn create_response_ed25519(
    domain_id: DomainId,
    predecessor_id: &AccountId,
    msg: &str,
    path: &str,
    signing_key: &eddsa::KeygenOutput,
) -> (Payload, SignatureRequest, SignatureRequestResponse) {
    let tweak = derive_tweak(predecessor_id, path);
    let derived_signing_key = derive_secret_key_ed25519(signing_key, &tweak);

    let payload: [u8; 32] = {
        let mut hasher = Sha256::new();
        hasher.update(msg);
        hasher.clone().finalize().into()
    };

    let derived_signing_key =
        frost_ed25519::SigningKey::from_scalar(derived_signing_key.private_share.to_scalar())
            .unwrap();

    let signature: [u8; 64] = derived_signing_key
        .sign(OsRng, &payload)
        .serialize()
        .unwrap()
        .try_into()
        .unwrap();

    let bytes = BoundedVec::from(payload);
    let payload = Payload::Eddsa(bytes);

    let respond_req = SignatureRequest::new(domain_id, payload.clone(), predecessor_id, path);

    let signature_response = SignatureRequestResponse::Ed25519 {
        signature: Ed25519Signature::from(signature),
    };

    (payload, respond_req, signature_response)
}

/// Process the message, creating the same hash with type of [`Digest`] and [`Payload`]
fn process_message(msg: &str) -> (impl Digest + use<>, Payload) {
    let msg = msg.as_bytes();
    let digest = <k256::Secp256k1 as ecdsa::hazmat::DigestPrimitive>::Digest::new_with_prefix(msg);
    let bytes: FieldBytes = digest.clone().finalize_fixed();

    let payload_hash = Payload::from_legacy_ecdsa(bytes.into());
    (digest, payload_hash)
}

pub struct PendingSignRequest {
    pub transaction: TransactionStatus,
    pub response: SignResponseArgs,
}

pub struct PendingCKDRequest {
    pub transaction: TransactionStatus,
    pub ckd_response: CKDResponseArgs,
}

fn gen_ed25519_sign_test(
    rng: &mut impl Rng,
    domain_id: DomainId,
    predecessor_id: &AccountId,
    sk: &eddsa::KeygenOutput,
) -> SignRequestTest {
    let msg: String = rng.r#gen::<usize>().to_string();
    let path: String = rng.r#gen::<usize>().to_string();
    let (payload, request, response) =
        create_response_ed25519(domain_id, predecessor_id, &msg, &path, sk);
    let args = SignRequestBuilder::new()
        .with_path(path)
        .with_payload(near_mpc_sdk::sign::Payload::Eddsa(
            payload.as_eddsa().unwrap().to_vec().try_into().unwrap(),
        ))
        .with_domain_id(domain_id.0)
        .build();
    SignRequestTest {
        response: SignResponseArgs { request, response },
        args,
    }
}

pub fn gen_secp_256k1_sign_test(
    rng: &mut impl Rng,
    domain_id: DomainId,
    predecessor_id: &AccountId,
    sk: &ts_ecdsa::KeygenOutput,
) -> SignRequestTest {
    let msg: String = rng.r#gen::<usize>().to_string();
    let path: String = rng.r#gen::<usize>().to_string();
    let (payload, request, response) =
        create_response_secp256k1(domain_id, predecessor_id, &msg, &path, sk);

    let payload_bytes: [u8; 32] = *payload.as_ecdsa().unwrap();

    let args = SignRequestBuilder::new()
        .with_path(path)
        .with_payload(near_mpc_sdk::sign::Payload::Ecdsa(payload_bytes.into()))
        .with_domain_id(domain_id.0)
        .build();
    SignRequestTest {
        response: SignResponseArgs { request, response },
        args,
    }
}