//! Local network setup.
use commonware_codec::{Decode, Encode};
use commonware_cryptography::{
    bls12381::{
        dkg::{deal, Output},
        primitives::{group::Share, variant::MinSig},
    },
    ed25519::{PrivateKey, PublicKey},
    Signer,
};
use commonware_math::algebra::Random;
use commonware_utils::{
    from_hex, hex,
    ordered::{Map, Set},
    quorum, TryCollect, NZU32,
};
use rand::{
    rngs::{OsRng, StdRng},
    seq::IteratorRandom,
    SeedableRng,
};
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 output (containing public information) as a hex string.
    pub output: 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 output(&self, max_participants_per_round: u32) -> Option<Output<MinSig, PublicKey>> {
        self.output.as_ref().map(|raw| {
            let bytes = from_hex(raw).expect("invalid hex string");
            Output::<MinSig, PublicKey>::decode_cfg(
                &mut bytes.as_slice(),
                &NZU32!(max_participants_per_round),
            )
            .expect("failed to decode polynomial")
        })
    }

    /// 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"),
        )
        .expect("failed to write participant config");
    }
}

/// A list of all peers' public keys.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PeerConfig<P: commonware_cryptography::PublicKey = PublicKey> {
    /// The number of participants per round.
    ///
    /// This is a vec which cycles through different numbers in each round.
    ///
    /// This MUST be non-empty.
    ///
    /// E.g. `vec![3, 4]` will start with 3 participants, then use 4, then use
    /// 3, etc.
    pub num_participants_per_round: Vec<u32>,
    /// All active peer public keys.
    #[serde(with = "serde_hex_ordered")]
    pub participants: Set<P>,
}

impl<P: commonware_cryptography::PublicKey> PeerConfig<P> {
    /// Returns the maximum number of participants per round.
    pub fn max_participants_per_round(&self) -> u32 {
        self.num_participants_per_round
            .iter()
            .copied()
            .max()
            .expect("num_participants_per_round must not be empty")
    }

    /// Returns the number of participants in the given round.
    pub fn num_participants_in_round(&self, round: u64) -> u32 {
        self.num_participants_per_round
            [(round % self.num_participants_per_round.len() as u64) as usize]
    }

    /// Pick the dealers for a particular round.
    ///
    /// The first round will use the first [`Self::num_participants_in_round`] players
    /// as the dealers.
    ///
    /// Subsequent rounds will their corresponding [`Self::num_participants_in_round`], and
    /// so on.
    pub fn dealers(&self, round: u64) -> Set<P> {
        let p_iter = self.participants.iter().cloned();
        let to_choose = self.num_participants_in_round(round) as usize;
        if round == 0 {
            return p_iter.take(to_choose).try_collect().unwrap();
        }
        let mut rng = StdRng::seed_from_u64(round);
        p_iter
            .choose_multiple(&mut rng, to_choose)
            .into_iter()
            .try_collect()
            .unwrap()
    }
}

/// 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).expect("failed to create data directory");

    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<Output<MinSig, PublicKey>>,
    Vec<(PrivateKey, Option<Share>)>,
) {
    // Generate p2p private keys
    let peer_signers = (0..num_peers)
        .map(|_| PrivateKey::random(&mut OsRng))
        .collect::<Vec<_>>();

    // Generate consensus key
    let threshold = quorum(num_participants_per_epoch);
    let all_participants: Set<PublicKey> = peer_signers
        .iter()
        .map(|s| s.public_key())
        .try_collect()
        .unwrap();
    let (output, shares) = if is_dkg {
        (None, Map::default())
    } else {
        let (output, shares) = deal(
            OsRng,
            Default::default(),
            all_participants
                .iter()
                .take(num_participants_per_epoch as usize)
                .cloned()
                .try_collect()
                .unwrap(),
        )
        .expect("deal failed: should have sufficient players");
        (Some(output), shares)
    };
    info!(num_peers, threshold, "generated participant identities");
    let identities = peer_signers
        .into_iter()
        .map(|s| {
            let share = shares.get_value(&s.public_key()).cloned();
            (s, share)
        })
        .collect::<Vec<_>>();
    (output, identities)
}

/// Generates all [ParticipantConfig] files from the provided identities.
fn generate_configs(
    args: &super::SetupArgs,
    output: Option<&Output<MinSig, PublicKey>>,
    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(),
            output: output.map(|o| hex(o.encode().as_ref())),
            signing_key: signer.clone(),
            share: share.clone(),
        };
        let config_file =
            File::create(&config_path).expect("failed to create participant config file");
        serde_json::to_writer_pretty(config_file, &participant_config)
            .expect("failed to serialize participant config");

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

    let peers = PeerConfig {
        num_participants_per_round: vec![args.num_participants_per_epoch],
        participants: identities
            .iter()
            .map(|(signer, _)| signer.public_key())
            .try_collect()
            .unwrap(),
    };
    let peers_file =
        File::create(args.datadir.join("peers.json")).expect("failed to create peers config file");
    serde_json::to_writer_pretty(peers_file, &peers).expect("failed to serialize peers config");
    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_ordered {
    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: &Set<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<Set<T>, D::Error>
    where
        T: Ord + DecodeExt<()>,
        D: Deserializer<'de>,
    {
        struct HexVecVisitor<T>(std::marker::PhantomData<T>);

        impl<'de, T> Visitor<'de> for HexVecVisitor<T>
        where
            T: Ord + DecodeExt<()>,
        {
            type Value = Set<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)?);
                }
                Set::try_from(out).map_err(|e| serde::de::Error::custom(format!("{e:?}")))
            }
        }

        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",
            ))
        }
    }
}