//! Local network setup.

use commonware_codec::{Decode, Encode};
use commonware_cryptography::{
    bls12381::{
        dkg::ops,
        primitives::{
            group::Share,
            poly::{self, Public},
            variant::MinSig,
        },
    },
    ed25519::{PrivateKey, PublicKey},
    PrivateKeyExt, Signer,
};
use commonware_utils::{from_hex, hex, quorum};
use rand::{rngs::OsRng, seq::IteratorRandom};
use serde::{Deserialize, Serialize};
use std::{
    collections::HashMap,
    fs::{self, File},
    net::{IpAddr, Ipv4Addr, SocketAddr},
    path::{Path, PathBuf},
};
use tracing::{error, info};

/// A configuration for a validator participant generated by the [setup](run) procedure.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ParticipantConfig {
    /// The port for P2P to listen on.
    pub port: u16,
    /// The bootstrapper node identities.
    #[serde(with = "serde_peer_map")]
    pub bootstrappers: HashMap<PublicKey, SocketAddr>,
    /// The group polynomial, as a hex-encoded string.
    pub polynomial: Option<String>,
    /// The validator's ed25519 signing key.
    ///
    /// Used for P2P, and in the DKG bootstrap phase, for consensus message signing.
    #[serde(with = "serde_hex")]
    pub signing_key: PrivateKey,
    /// The validator's share, if active in the first epoch.
    #[serde(with = "serde_hex")]
    pub share: Option<Share>,
}

impl ParticipantConfig {
    /// Get the polynomial commitment for the DKG.
    pub fn polynomial(&self, threshold: u32) -> Option<Public<MinSig>> {
        self.polynomial.as_ref().map(|raw| {
            let bytes = from_hex(raw).expect("invalid hex string");
            Public::<MinSig>::decode_cfg(&mut bytes.as_slice(), &(threshold as usize)).unwrap()
        })
    }

    /// Update the participant config using the provided closure, then write it back to disk.
    pub fn update_and_write(mut self, path: &Path, f: impl FnOnce(&mut Self)) {
        f(&mut self);

        std::fs::write(
            path,
            serde_json::to_string_pretty(&self).expect("failed to serialize participant config"),
        )
        .unwrap();
    }
}

/// A list of all peers' public keys.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PeerConfig {
    /// The number of participants per epoch.
    pub num_participants_per_epoch: u32,
    /// All active peer public keys.
    #[serde(with = "serde_hex_vec")]
    pub active: Vec<PublicKey>,
    /// All inactive peer public keys.
    #[serde(with = "serde_hex_vec")]
    pub inactive: Vec<PublicKey>,
}

impl PeerConfig {
    pub fn threshold(&self) -> u32 {
        quorum(self.num_participants_per_epoch)
    }
}

/// Run the setup procedure, generating a number of random validator identities.
pub fn run(args: super::SetupArgs) {
    if args.datadir.exists() {
        error!("Data directory already exists; Remove it before setting up a new network");
        return;
    }

    fs::create_dir_all(&args.datadir).unwrap();

    let (polynomial, identities) = generate_identities(
        args.with_dkg,
        args.num_peers,
        args.num_participants_per_epoch,
    );
    let configs = generate_configs(&args, polynomial.as_ref(), &identities);

    let mprocs_validator_cmd = configs
        .into_iter()
        .fold(vec!["mprocs".to_string()], |mut acc, cfg| {
            acc.push(format!(
                r#""cargo run --bin commonware-reshare --release validator --cfg {} --peers {}""#,
                cfg.display(),
                args.datadir.join("peers.json").display()
            ));
            acc
        })
        .join(" ");

    if args.with_dkg {
        println!("\nThe network is configured to use DKG to distribute initial shares.");
        let mprocs_dkg_cmd = mprocs_validator_cmd.replace("validator", "dkg");
        println!("\nTo start the DKG process, run the following command:");
        println!("\n{mprocs_dkg_cmd}");
        println!("\nOnce the DKG process completes, exit the DKG processes and start the validators with the following command:");
    } else {
        println!("\nThe network is configured with a trusted threshold setup.");
        println!("\nTo start the validators, run the following command:");
    }

    println!("\n{mprocs_validator_cmd}");
}

/// Generate shares, ed25519 private keys, and a commitment for a given number of participants.
#[allow(clippy::type_complexity)]
fn generate_identities(
    is_dkg: bool,
    num_peers: u32,
    num_participants_per_epoch: u32,
) -> (Option<Public<MinSig>>, Vec<(PrivateKey, Option<Share>)>) {
    // Generate consensus key
    let threshold = quorum(num_participants_per_epoch);
    let (polynomial, shares) = if is_dkg {
        (None, Vec::new())
    } else {
        let (polynomial, shares) = ops::generate_shares::<_, MinSig>(
            &mut OsRng,
            None,
            num_participants_per_epoch,
            threshold,
        );
        info!(identity = ?poly::public::<MinSig>(&polynomial), "generated network key");
        (Some(polynomial), shares)
    };

    // Assign shares to peers (those not participating in the first epoch get no share.)
    let mut shares = shares.into_iter().map(Some).collect::<Vec<_>>();
    shares.resize(num_peers as usize, None);

    // Generate p2p private keys
    let mut peer_signers = (0..num_peers)
        .map(|_| PrivateKey::from_rng(&mut OsRng))
        .collect::<Vec<_>>();
    peer_signers.sort_by_key(|signer| signer.public_key());

    let identities = peer_signers.into_iter().zip(shares).collect::<Vec<_>>();

    info!(num_peers, threshold, "generated participant identities");
    (polynomial, identities)
}

/// Generates all [ParticipantConfig] files from the provided identities.
fn generate_configs(
    args: &super::SetupArgs,
    polynomial: Option<&Public<MinSig>>,
    identities: &[(PrivateKey, Option<Share>)],
) -> Vec<PathBuf> {
    let bootstrappers = identities
        .iter()
        .enumerate()
        .choose_multiple(&mut OsRng, args.num_bootstrappers)
        .into_iter()
        .map(|(i, (signer, _))| {
            (
                signer.public_key(),
                SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), args.base_port + i as u16),
            )
        })
        .collect::<HashMap<_, _>>();

    let mut configs = Vec::with_capacity(identities.len());
    for (index, (signer, share)) in identities.iter().enumerate() {
        let config_path = args.datadir.join(format!("participant-{index}.json"));
        let participant_config = ParticipantConfig {
            port: args.base_port + index as u16,
            bootstrappers: bootstrappers.clone(),
            polynomial: polynomial.map(|polynomial| hex(polynomial.encode().as_ref())),
            signing_key: signer.clone(),
            share: share.clone(),
        };
        let config_file = File::create(&config_path).unwrap();
        serde_json::to_writer_pretty(config_file, &participant_config).unwrap();

        configs.push(config_path);
    }
    info!("wrote participant configurations");

    let peers = if args.with_dkg {
        PeerConfig {
            num_participants_per_epoch: args.num_participants_per_epoch,
            active: identities
                .iter()
                .map(|(signer, _)| signer.public_key())
                .take(args.num_participants_per_epoch as usize)
                .collect(),
            inactive: identities
                .iter()
                .map(|(signer, _)| signer.public_key())
                .skip(args.num_participants_per_epoch as usize)
                .take((args.num_peers - args.num_participants_per_epoch) as usize)
                .collect(),
        }
    } else {
        PeerConfig {
            num_participants_per_epoch: args.num_participants_per_epoch,
            active: identities
                .iter()
                .filter(|(_, share)| share.is_some())
                .map(|(signer, _)| signer.public_key())
                .collect(),
            inactive: identities
                .iter()
                .filter(|(_, share)| share.is_none())
                .map(|(signer, _)| signer.public_key())
                .collect(),
        }
    };
    let peers_file = File::create(args.datadir.join("peers.json")).unwrap();
    serde_json::to_writer_pretty(peers_file, &peers).unwrap();
    info!("wrote peers map");

    configs
}

// --- `serde` helpers ---

mod serde_hex {
    use super::*;
    use commonware_codec::DecodeExt;
    use commonware_utils::from_hex;
    use serde::{Deserializer, Serializer};
    use serde_json::Value;

    pub fn serialize<T, S>(value: &T, serializer: S) -> Result<S::Ok, S::Error>
    where
        T: Encode,
        S: Serializer,
    {
        hex(&value.encode()).serialize(serializer)
    }

    pub fn deserialize<'de, T, D>(deserializer: D) -> Result<T, D::Error>
    where
        T: DecodeExt<()>,
        D: Deserializer<'de>,
    {
        if let Value::String(s) = Value::deserialize(deserializer)? {
            let bytes = from_hex(&s).ok_or(serde::de::Error::custom(
                "failed to deserialize: invalid hex string",
            ))?;
            T::decode(&mut bytes.as_slice())
                .map_err(|_| serde::de::Error::custom("failed to decode bytes"))
        } else {
            Err(serde::de::Error::custom(
                "failed to deserialize: expected a hex string",
            ))
        }
    }
}

mod serde_hex_vec {
    use super::*;
    use commonware_codec::DecodeExt;
    use commonware_utils::from_hex;
    use core::fmt;
    use serde::{
        de::{SeqAccess, Visitor},
        Deserializer, Serializer,
    };

    pub fn serialize<T, S>(value: &[T], serializer: S) -> Result<S::Ok, S::Error>
    where
        T: Encode,
        S: Serializer,
    {
        // Serialize each element as a hex string
        serializer.collect_seq(value.iter().map(|v| hex(&v.encode())))
    }

    pub fn deserialize<'de, T, D>(deserializer: D) -> Result<Vec<T>, D::Error>
    where
        T: DecodeExt<()>,
        D: Deserializer<'de>,
    {
        struct HexVecVisitor<T>(std::marker::PhantomData<T>);

        impl<'de, T> Visitor<'de> for HexVecVisitor<T>
        where
            T: DecodeExt<()>,
        {
            type Value = Vec<T>;

            fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
                f.write_str("a sequence of hex-encoded values")
            }

            fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
            where
                A: SeqAccess<'de>,
            {
                let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(0));
                while let Some(s) = seq.next_element::<String>()? {
                    let bytes = from_hex(&s).ok_or(serde::de::Error::custom(
                        "failed to convert from hex to bytes",
                    ))?;

                    out.push(T::decode(&mut bytes.as_ref()).map_err(serde::de::Error::custom)?);
                }
                Ok(out)
            }
        }

        deserializer.deserialize_seq(HexVecVisitor(std::marker::PhantomData))
    }
}

mod serde_peer_map {
    use super::*;
    use commonware_codec::DecodeExt;
    use commonware_utils::from_hex;
    use serde::{Deserializer, Serializer};
    use serde_json::Value;

    pub fn serialize<S: Serializer>(
        value: &HashMap<PublicKey, SocketAddr>,
        serializer: S,
    ) -> Result<S::Ok, S::Error> {
        let hex_map = value
            .iter()
            .map(|(k, v)| (hex(&k.encode()), *v))
            .collect::<HashMap<_, _>>();
        hex_map.serialize(serializer)
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<HashMap<PublicKey, SocketAddr>, D::Error>
    where
        D: Deserializer<'de>,
    {
        if let Value::Object(map) = Value::deserialize(deserializer)? {
            let mut result = HashMap::new();
            for (k, v) in map {
                let pk_bytes = from_hex(&k).ok_or(serde::de::Error::custom(
                    "failed to deserialize: invalid hex string for public key",
                ))?;
                let pk = PublicKey::decode(&mut pk_bytes.as_slice())
                    .map_err(|_| serde::de::Error::custom("failed to decode public key bytes"))?;

                let Value::String(addr_str) = v else {
                    return Err(serde::de::Error::custom(
                        "failed to deserialize: expected a string for socket address",
                    ));
                };
                let socket_addr = addr_str.parse::<SocketAddr>().map_err(|_| {
                    serde::de::Error::custom("failed to parse socket address from string")
                })?;

                result.insert(pk, socket_addr);
            }
            Ok(result)
        } else {
            Err(serde::de::Error::custom(
                "failed to deserialize: expected a map",
            ))
        }
    }
}