use borsh::{BorshDeserialize, BorshSerialize};
use curve25519_dalek::{constants::ED25519_BASEPOINT_POINT, EdwardsPoint};
use k256::{
    elliptic_curve::{
        group::GroupEncoding, sec1::FromEncodedPoint, sec1::ToEncodedPoint, subtle::CtOption,
        CurveArithmetic, PrimeField,
    },
    EncodedPoint, Secp256k1,
};
use near_sdk::{ext_contract, near, AccountId, CurveType, PublicKey};
use serde::{Deserialize, Serialize};
use sha3::{Digest, Sha3_256};
use std::{fmt::Display, panic, str::FromStr};
use tiny_keccak::{Hasher, Keccak};

use crate::chainsig::{chain_to_domain, chain_to_domain_id, PublicKeyAddress};
pub mod evm;
pub mod near;
pub mod svm;

#[near(serializers=[json,borsh])]
pub enum Domain {
    SEC1,
    ED25519,
}

#[near(serializers = [json, borsh])]
#[derive(Clone)]
pub enum ChainEnvironment {
    SVM,
    EVM,
    NearWasm,
}

impl ChainEnvironment {
    pub fn from_string(chain: String) -> ChainEnvironment {
        if chain == "SVM" {
            return ChainEnvironment::SVM;
        } else if chain == "EVM" {
            return ChainEnvironment::EVM;
        } else if chain == "NearWasm" {
            return ChainEnvironment::NearWasm;
        }

        panic!("Unsupported chain value is passed in")
    }
}

#[near(serializers = [json, borsh])]
pub enum NearNetwork {
    Mainnet,
    Testnet,
}

#[near(serializers = [json, borsh])]
pub enum ChainSigPayloadV2 {
    Ecdsa(String),
    Eddsa(String),
}

#[near(serializers = [json, borsh])]
pub struct ChainSigSignRequest {
    pub path: String,
    pub domain_id: u8,
    pub payload_v2: ChainSigPayloadV2,
}

#[near(serializers = [json, borsh])]
pub struct BigR {
    pub affine_point: String,
}

#[near(serializers = [json, borsh])]
pub struct Scalar {
    pub scalar: String,
}

#[near(serializers = [json, borsh])]
#[serde(tag = "scheme")]
pub enum ChainSigResponse {
    #[serde(rename = "Ed25519")]
    Ed25519 { signature: Vec<u8> },

    #[serde(rename = "Secp256k1")]
    Secp256k1 {
        big_r: BigR,
        s: Scalar,
        recovery_id: u8,
    },
}

#[allow(unused)]
#[ext_contract(chainsig_contract)]
pub trait ChainSigContract {
    fn sign(&self, request: ChainSigSignRequest) -> ChainSigResponse;
}

pub fn get_mpc_contract_id_by_network(near_network: &NearNetwork) -> AccountId {
    match near_network {
        NearNetwork::Mainnet => AccountId::from_str("v1.signer").unwrap(),
        NearNetwork::Testnet => AccountId::from_str("v1.signer-prod.testnet").unwrap(),
    }
}

pub fn get_mpc_contract_by_network(
    near_network: &NearNetwork,
) -> chainsig_contract::ChainSigContractExt {
    chainsig_contract::ext(get_mpc_contract_id_by_network(near_network))
}

#[derive(Debug, Clone, Eq, Ord, PartialEq, PartialOrd)]
#[near(serializers=[borsh, json])]
pub struct Tweak([u8; 32]);

impl Tweak {
    pub fn as_bytes(&self) -> [u8; 32] {
        self.0
    }

    pub fn new(bytes: [u8; 32]) -> Self {
        Self(bytes)
    }
}

pub trait ChainSig {
    // Should return an array of hashes to be signed with chain sig
    fn prepare_transaction_for_signing(encoded_tx: String) -> Vec<String>;
}

// Constant prefix that ensures tweak derivation values are used specifically for
// near-mpc-recovery with key derivation protocol vX.Y.Z.
const TWEAK_DERIVATION_PREFIX: &str = "near-mpc-recovery v0.1.0 epsilon derivation:";

pub fn derive_tweak(predecessor_id: &AccountId, path: &str) -> Tweak {
    // ',' is ACCOUNT_DATA_SEPARATOR from nearcore that indicate the end
    // of the accound id in the trie key. We reuse the same constant to
    // indicate the end of the account id in derivation path.
    // Do not reuse this hash function on anything that isn't an account
    // ID or it'll be vunerable to Hash Melleability/extention attacks.
    let derivation_path = format!("{TWEAK_DERIVATION_PREFIX}{},{}", predecessor_id, path);
    let mut hasher = Sha3_256::new();
    hasher.update(derivation_path);
    let hash: [u8; 32] = hasher.finalize().into();
    Tweak::new(hash)
}

/// Each domain corresponds to a specific root key in a specific signature scheme. There may be
/// multiple domains per signature scheme. The domain ID uniquely identifies a domain.
#[near(serializers=[borsh, json])]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct DomainId(pub u64);

impl Default for DomainId {
    fn default() -> Self {
        Self::legacy_ecdsa_id()
    }
}

impl DomainId {
    /// Returns the DomainId of the single ECDSA key present in the contract before V2.
    pub fn legacy_ecdsa_id() -> Self {
        Self(0)
    }
}

impl Display for DomainId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.0.fmt(f)
    }
}

#[cfg_attr(
    all(feature = "abi", not(target_arch = "wasm32")),
    derive(::near_sdk::schemars::JsonSchema),
    derive(::borsh::BorshSchema)
)]
#[derive(
    Debug,
    PartialEq,
    Serialize,
    Deserialize,
    Eq,
    Clone,
    Copy,
    derive_more::From,
    derive_more::AsRef,
    derive_more::Deref,
)]
pub struct SerializableEdwardsPoint(
    #[cfg_attr(
        all(feature = "abi", not(target_arch = "wasm32")),
        schemars(with = "[u8; 32]"),
        borsh(schema(with_funcs(
            declaration = "<[u8; 32] as ::borsh::BorshSchema>::declaration",
            definitions = "<[u8; 32] as ::borsh::BorshSchema>::add_definitions_recursively"
        ),))
    )]
    EdwardsPoint,
);

impl GroupEncoding for SerializableEdwardsPoint {
    type Repr = [u8; 32];

    fn from_bytes(bytes: &Self::Repr) -> CtOption<Self> {
        EdwardsPoint::from_bytes(bytes).map(Into::into)
    }

    fn from_bytes_unchecked(bytes: &Self::Repr) -> CtOption<Self> {
        Self::from_bytes(bytes)
    }

    fn to_bytes(&self) -> Self::Repr {
        self.compress().to_bytes()
    }
}

impl BorshSerialize for SerializableEdwardsPoint {
    fn serialize<W: std::io::prelude::Write>(&self, writer: &mut W) -> std::io::Result<()> {
        let bytes = self.0.to_bytes();
        BorshSerialize::serialize(&bytes, writer)
    }
}

impl BorshDeserialize for SerializableEdwardsPoint {
    fn deserialize_reader<R: std::io::prelude::Read>(reader: &mut R) -> std::io::Result<Self> {
        let bytes: [u8; 32] = BorshDeserialize::deserialize_reader(reader)?;

        SerializableEdwardsPoint::from_bytes(&bytes)
            .into_option()
            .ok_or(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "The provided bytes is not a valid edwards point.",
            ))
    }
}

#[cfg_attr(
    all(feature = "abi", not(target_arch = "wasm32")),
    derive(::near_sdk::schemars::JsonSchema),
    derive(::borsh::BorshSchema)
)]
#[derive(Debug, PartialEq, Eq, Clone, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub enum PublicKeyExtended {
    Secp256k1 {
        near_public_key: near_sdk::PublicKey,
    },
    // Invariant: `edwards_point` is always the decompressed representation of `near_public_key_compressed`.
    Ed25519 {
        /// Serialized compressed Edwards-y point.
        near_public_key_compressed: near_sdk::PublicKey,
        /// Decompressed Edwards point used for curve arithmetic operations.
        edwards_point: SerializableEdwardsPoint,
    },
}

pub fn public_key_extended(domain_id: DomainId, near_network: &NearNetwork) -> PublicKeyExtended {
    if domain_id == DomainId(0) {
        let secp: PublicKey = match near_network {
            NearNetwork::Mainnet => {
                "secp256k1:3tFRbMqmoa6AAALMrEFAYCEoHcqKxeW38YptwowBVBtXK1vo36HDbUWuR6EZmoK4JcH6HDkNMGGqP1ouV7VZUWya".parse().unwrap()
            },
            NearNetwork::Testnet => {
                "secp256k1:4NfTiv3UsGahebgTaHyD9vF8KYKMBnfd6kh94mK6xv8fGBiJB8TBtFMP5WWXz6B89Ac1fbpzPwAvoyQebemHFwx3".parse().unwrap()
            },
        };
        PublicKeyExtended::Secp256k1 {
            near_public_key: (secp),
        }
    } else if domain_id == DomainId(1) {
        let (ed, public_key_bytes): (PublicKey, [u8; 32]) = match near_network {
            NearNetwork::Mainnet => (
                "ed25519:G9hwngxWNKdmqMCmU1Yt6LPhFpayJeKFxyAV1HqMNLtF"
                    .parse()
                    .unwrap(),
                [
                    225, 26, 63, 32, 184, 108, 201, 78, 91, 117, 55, 216, 102, 209, 196, 106, 227,
                    228, 26, 70, 161, 208, 143, 122, 234, 22, 150, 164, 165, 0, 191, 208,
                ],
            ),
            NearNetwork::Testnet => (
                "ed25519:6vSEtQxrQj6txUMh33WC4ERyCWmNMRTdufDWAaDY3Un2"
                    .parse()
                    .unwrap(),
                [
                    87, 250, 193, 125, 62, 246, 7, 222, 13, 52, 179, 58, 18, 195, 58, 130, 12, 149,
                    243, 102, 186, 172, 40, 227, 40, 252, 15, 7, 157, 181, 179, 127,
                ],
            ),
        };

        let edwards_point = SerializableEdwardsPoint::from_bytes(&public_key_bytes)
            .into_option()
            .unwrap();

        PublicKeyExtended::Ed25519 {
            near_public_key_compressed: (ed),
            edwards_point,
        }
    } else {
        panic!("Invalid domain")
    }
}

pub mod k256_types {
    use k256::{elliptic_curve::CurveArithmetic, Secp256k1};

    pub type PublicKey = <Secp256k1 as CurveArithmetic>::AffinePoint;
}

pub fn near_public_key_to_affine_point(pk: near_sdk::PublicKey) -> k256_types::PublicKey {
    assert_eq!(
        pk.curve_type(),
        near_sdk::CurveType::SECP256K1,
        "Expected a key on the SECP256K1 curve"
    );

    let mut bytes = pk.into_bytes();
    bytes[0] = 0x04;
    let point = EncodedPoint::from_bytes(bytes).unwrap();
    k256_types::PublicKey::from_encoded_point(&point).unwrap()
}

#[derive(Debug, Clone)]
pub struct TweakNotOnCurve;

pub fn derive_key_secp256k1(
    public_key: &k256_types::PublicKey,
    tweak: &Tweak,
) -> Result<k256_types::PublicKey, TweakNotOnCurve> {
    let tweak = k256::Scalar::from_repr(tweak.as_bytes().into())
        .into_option()
        .ok_or(TweakNotOnCurve)?;

    Ok(
        (<Secp256k1 as CurveArithmetic>::ProjectivePoint::GENERATOR * tweak + public_key)
            .to_affine(),
    )
}

pub fn naj_to_uncompressed_public_key_sec1(naj_public_key: &String) -> String {
    let parts: Vec<&str> = naj_public_key.split(':').collect();
    let key_part = parts[1];

    let decoded_key = bs58::decode(key_part)
        .into_vec()
        .expect("Invalid base58 key");

    format!("04{}", hex::encode(decoded_key))
}

pub fn derive_public_key_edwards_point_ed25519(
    public_key_edwards_point: &curve25519_dalek::EdwardsPoint,
    tweak: &Tweak,
) -> curve25519_dalek::EdwardsPoint {
    let tweak = curve25519_dalek::Scalar::from_bytes_mod_order(tweak.as_bytes());
    public_key_edwards_point + ED25519_BASEPOINT_POINT * tweak
}

pub fn derived_public_key(
    path: String,
    predecessor: AccountId,
    domain_id: DomainId,
    near_network: &NearNetwork,
) -> PublicKey {
    let tweak = derive_tweak(&predecessor, &path);

    let domain = domain_id;

    let public_key = public_key_extended(domain, near_network);

    let derived_public_key = match public_key {
        PublicKeyExtended::Secp256k1 { near_public_key } => {
            let derived_public_key =
                derive_key_secp256k1(&near_public_key_to_affine_point(near_public_key), &tweak)
                    .unwrap();

            let encoded_point = derived_public_key.to_encoded_point(false);
            let slice: &[u8] = &encoded_point.as_bytes()[1..65];
            PublicKey::from_parts(CurveType::SECP256K1, slice.to_vec())
        }
        PublicKeyExtended::Ed25519 { edwards_point, .. } => {
            let derived_public_key_edwards_point =
                derive_public_key_edwards_point_ed25519(&edwards_point, &tweak);

            let encoded_point: [u8; 32] = derived_public_key_edwards_point.compress().to_bytes();

            PublicKey::from_parts(CurveType::ED25519, encoded_point.into())
        }
    };

    derived_public_key.unwrap()
}

pub fn derived_public_key_with_address(
    path: &String,
    chain: &ChainEnvironment,
    near_network: &NearNetwork,
    predecessor: &AccountId,
) -> PublicKeyAddress {
    let domain_id = chain_to_domain_id(chain);

    let public_key_result = derived_public_key(
        path.to_owned(),
        predecessor.clone(),
        DomainId(domain_id.into()),
        near_network,
    );

    let mut public_key: String = String::from(&public_key_result);

    let domain = chain_to_domain(chain);

    public_key = match domain {
        Domain::SEC1 => naj_to_uncompressed_public_key_sec1(&public_key),
        Domain::ED25519 => public_key,
    };

    let public_key_address = match chain {
        ChainEnvironment::SVM => {
            // strip out ed25519:
            let base58_key = &public_key[8..];

            PublicKeyAddress {
                public_key: base58_key.to_string(),
                address: base58_key.to_string(),
            }
        }
        ChainEnvironment::EVM => {
            let public_key_no_prefix = &public_key[2..];

            let public_key_bytes =
                hex::decode(public_key_no_prefix).expect("Invalid hex in public key");

            let mut hasher = Keccak::v256();
            hasher.update(&public_key_bytes);
            let mut hash = [0u8; 32];
            hasher.finalize(&mut hash);

            let address_bytes = &hash[12..];

            PublicKeyAddress {
                public_key,
                address: format!("0x{}", hex::encode(address_bytes)),
            }
        }
        ChainEnvironment::NearWasm => {
            // strip out ed25519:
            let ori_public_key = public_key.clone();
            let base58_key = &public_key[8..];

            let bytes = bs58::decode(base58_key).into_vec().expect("Invalid base58");

            let hex_public_key = hex::encode(bytes);

            PublicKeyAddress {
                public_key: ori_public_key,
                address: hex_public_key.to_string(),
            }
        }
    };

    public_key_address
}

#[cfg(all(test, not(target_arch = "wasm32")))]
mod test_derive_public_key {

    use crate::helpers::chainsig::{derived_public_key, DomainId, NearNetwork};

    #[test]
    fn test_ed25519_derive() {
        let key = derived_public_key(
            "2".to_owned(),
            "alice.near".parse().unwrap(),
            DomainId(1),
            &NearNetwork::Mainnet,
        );
        let s: String = String::from(&key);
        assert_eq!(s, "ed25519:9TCzxc6DXNPxYp2RQKHfGWrbxG6NLgLxDsFTyHKCWRHY");
    }

    #[test]
    fn test_secp256k1_derive() {
        let key = derived_public_key(
            "1".to_owned(),
            "alice.near".parse().unwrap(),
            DomainId(0),
            &NearNetwork::Mainnet,
        );
        let s: String = String::from(&key);
        assert_eq!(s, "secp256k1:2PK3TZFzPq8KktvFuvuXe121nrbH57ASSF3cttJZFjG74X7RD1ubxQSMEsfyYDmwmpSM7rnPMmJ7oVWKfQQB3gQY");
    }
}