use std::collections::BTreeMap;
use std::sync::{Arc, OnceLock};
use std::time::{Duration, Instant};

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

use crate::rtr::payload::{Aspa, Payload, RouteOrigin, RouterKey};

use super::ordering::{compare_payload_update_for_rtr, sort_payloads_for_rtr};

#[derive(Debug, Clone)]
pub struct DualTime {
    instant: Instant,
    utc: DateTime<Utc>,
}

impl DualTime {
    pub fn now() -> Self {
        Self {
            instant: Instant::now(),
            utc: Utc::now(),
        }
    }

    pub fn utc(&self) -> DateTime<Utc> {
        self.utc
    }

    pub fn elapsed(&self) -> Duration {
        self.instant.elapsed()
    }

    pub fn is_expired(&self, duration: Duration) -> bool {
        self.elapsed() >= duration
    }

    pub fn reset(&mut self) {
        self.instant = Instant::now();
        self.utc = Utc::now();
    }
}

impl Serialize for DualTime {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        self.utc.timestamp_millis().serialize(serializer)
    }
}

impl<'de> Deserialize<'de> for DualTime {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let millis = i64::deserialize(deserializer)?;
        let utc = DateTime::from_timestamp_millis(millis)
            .ok_or_else(|| serde::de::Error::custom("invalid timestamp"))?;

        Ok(Self {
            instant: Instant::now(),
            utc,
        })
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Snapshot {
    origins: Arc<Vec<RouteOrigin>>,
    router_keys: Arc<Vec<RouterKey>>,
    aspas: Arc<Vec<Aspa>>,
    created_at: DualTime,
    origins_hash: [u8; 32],
    router_keys_hash: [u8; 32],
    aspas_hash: [u8; 32],
    snapshot_hash: [u8; 32],
    #[serde(skip)]
    rtr_payloads_cache: OnceLock<Arc<Vec<Payload>>>,
}

impl Snapshot {
    pub fn new(
        origins: Vec<RouteOrigin>,
        router_keys: Vec<RouterKey>,
        aspas: Vec<Aspa>,
    ) -> Self {
        Self::from_shared_parts(
            Arc::new(sorted_dedup(origins)),
            Arc::new(sorted_dedup(router_keys)),
            Arc::new(normalize_aspas(aspas)),
        )
    }

    fn from_shared_parts(
        origins: Arc<Vec<RouteOrigin>>,
        router_keys: Arc<Vec<RouterKey>>,
        aspas: Arc<Vec<Aspa>>,
    ) -> Self {
        let mut snapshot = Snapshot {
            origins,
            router_keys,
            aspas,
            created_at: DualTime::now(),
            origins_hash: [0u8; 32],
            router_keys_hash: [0u8; 32],
            aspas_hash: [0u8; 32],
            snapshot_hash: [0u8; 32],
            rtr_payloads_cache: OnceLock::new(),
        };
        snapshot.recompute_hashes();
        let cached = build_snapshot_payloads_for_rtr(&snapshot);
        let _ = snapshot.rtr_payloads_cache.set(Arc::new(cached));
        snapshot
    }

    pub fn empty() -> Self {
        Self::new(Vec::new(), Vec::new(), Vec::new())
    }

    pub fn from_payloads(payloads: Vec<Payload>) -> Self {
        let mut builder = SnapshotBuilder::new();
        builder.extend(payloads);
        builder.finish()
    }

    pub fn builder() -> SnapshotBuilder {
        SnapshotBuilder::new()
    }

    pub fn project_for_version(&self, version: u8) -> Self {
        let router_keys = if version >= 1 {
            self.router_keys.clone()
        } else {
            Arc::new(Vec::new())
        };
        let aspas = if version >= 2 {
            self.aspas.clone()
        } else {
            Arc::new(Vec::new())
        };

        Self::from_shared_parts(self.origins.clone(), router_keys, aspas)
    }

    pub fn recompute_hashes(&mut self) {
        self.origins_hash = self.compute_origins_hash();
        self.router_keys_hash = self.compute_router_keys_hash();
        self.aspas_hash = self.compute_aspas_hash();
        self.snapshot_hash = self.compute_snapshot_hash();
    }

    fn compute_origins_hash(&self) -> [u8; 32] {
        Self::hash_ordered_iter(self.origins.iter())
    }

    fn compute_router_keys_hash(&self) -> [u8; 32] {
        Self::hash_ordered_iter(self.router_keys.iter())
    }

    fn compute_aspas_hash(&self) -> [u8; 32] {
        Self::hash_ordered_iter(self.aspas.iter())
    }

    fn compute_snapshot_hash(&self) -> [u8; 32] {
        let mut hasher = Sha256::new();
        hasher.update(b"snapshot:v1");
        hasher.update(self.origins_hash);
        hasher.update(self.router_keys_hash);
        hasher.update(self.aspas_hash);
        hasher.finalize().into()
    }

    fn hash_ordered_iter<'a, T, I>(iter: I) -> [u8; 32]
    where
        T: Serialize + 'a,
        I: IntoIterator<Item = &'a T>,
    {
        let mut hasher = Sha256::new();
        hasher.update(b"set:v1");

        for item in iter {
            let encoded =
                serde_json::to_vec(item).expect("serialize snapshot item for hashing failed");
            let len = (encoded.len() as u32).to_be_bytes();
            hasher.update(len);
            hasher.update(encoded);
        }

        hasher.finalize().into()
    }

    pub fn diff(&self, new_snapshot: &Snapshot) -> (Vec<Payload>, Vec<Payload>) {
        let mut announced = Vec::new();
        let mut withdrawn = Vec::new();

        if !self.same_origins(new_snapshot) {
            diff_sorted(
                self.origins.as_slice(),
                new_snapshot.origins.as_slice(),
                &mut announced,
                &mut withdrawn,
                Payload::RouteOrigin,
            );
        }

        if !self.same_router_keys(new_snapshot) {
            diff_sorted(
                self.router_keys.as_slice(),
                new_snapshot.router_keys.as_slice(),
                &mut announced,
                &mut withdrawn,
                Payload::RouterKey,
            );
        }

        if !self.same_aspas(new_snapshot) {
            diff_aspas(
                &self.aspas,
                &new_snapshot.aspas,
                &mut announced,
                &mut withdrawn,
            );
        }

        (announced, withdrawn)
    }

    pub fn created_at(&self) -> DualTime {
        self.created_at.clone()
    }

    pub fn payloads(&self) -> Vec<Payload> {
        let mut v =
            Vec::with_capacity(self.origins.len() + self.router_keys.len() + self.aspas.len());

        v.extend(self.origins.iter().cloned().map(Payload::RouteOrigin));
        v.extend(self.router_keys.iter().cloned().map(Payload::RouterKey));
        v.extend(self.aspas.iter().cloned().map(Payload::Aspa));

        v
    }

    pub fn payloads_for_rtr(&self) -> Vec<Payload> {
        self.rtr_payloads_for_rtr_arc().as_ref().clone()
    }

    pub fn rtr_payloads_for_rtr_arc(&self) -> Arc<Vec<Payload>> {
        self.rtr_payloads_cache
            .get_or_init(|| Arc::new(build_snapshot_payloads_for_rtr(self)))
            .clone()
    }

    pub fn origins_hash(&self) -> [u8; 32] {
        self.origins_hash
    }

    pub fn router_keys_hash(&self) -> [u8; 32] {
        self.router_keys_hash
    }

    pub fn aspas_hash(&self) -> [u8; 32] {
        self.aspas_hash
    }

    pub fn snapshot_hash(&self) -> [u8; 32] {
        self.snapshot_hash
    }

    pub fn same_origins(&self, other: &Self) -> bool {
        self.origins_hash == other.origins_hash
    }

    pub fn same_router_keys(&self, other: &Self) -> bool {
        self.router_keys_hash == other.router_keys_hash
    }

    pub fn same_aspas(&self, other: &Self) -> bool {
        self.aspas_hash == other.aspas_hash
    }

    pub fn same_content(&self, other: &Self) -> bool {
        self.snapshot_hash == other.snapshot_hash
    }

    pub fn origins(&self) -> &[RouteOrigin] {
        self.origins.as_ref()
    }

    pub fn router_keys(&self) -> &[RouterKey] {
        self.router_keys.as_ref()
    }

    pub fn aspas(&self) -> &[Aspa] {
        self.aspas.as_ref()
    }

    pub fn is_empty(&self) -> bool {
        self.origins.is_empty() && self.router_keys.is_empty() && self.aspas.is_empty()
    }
}

#[derive(Debug, Default)]
pub struct SnapshotBuilder {
    origins: Vec<RouteOrigin>,
    router_keys: Vec<RouterKey>,
    aspas: Vec<Aspa>,
}

impl SnapshotBuilder {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn push(&mut self, payload: Payload) {
        match payload {
            Payload::RouteOrigin(o) => self.origins.push(o),
            Payload::RouterKey(k) => self.router_keys.push(k),
            Payload::Aspa(a) => self.aspas.push(a),
        }
    }

    pub fn extend<I>(&mut self, payloads: I)
    where
        I: IntoIterator<Item = Payload>,
    {
        for payload in payloads {
            self.push(payload);
        }
    }

    pub fn finish(self) -> Snapshot {
        Snapshot::new(self.origins, self.router_keys, normalize_aspas(self.aspas))
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Delta {
    serial: u32,
    announced: Vec<Payload>,
    withdrawn: Vec<Payload>,
    created_at: DualTime,
}

impl Delta {
    pub fn new(serial: u32, mut announced: Vec<Payload>, mut withdrawn: Vec<Payload>) -> Self {
        dedup_payloads(&mut announced);
        dedup_payloads(&mut withdrawn);

        sort_payloads_for_rtr(&mut announced, true);
        sort_payloads_for_rtr(&mut withdrawn, false);

        Delta {
            serial,
            announced,
            withdrawn,
            created_at: DualTime::now(),
        }
    }

    pub fn serial(&self) -> u32 {
        self.serial
    }

    pub fn announced(&self) -> &[Payload] {
        &self.announced
    }

    pub fn withdrawn(&self) -> &[Payload] {
        &self.withdrawn
    }

    pub fn created_at(&self) -> DualTime {
        self.created_at.clone()
    }

    pub fn is_empty(&self) -> bool {
        self.announced.is_empty() && self.withdrawn.is_empty()
    }

    pub fn payload_updates_for_rtr(&self) -> Vec<(bool, &Payload)> {
        build_payload_updates_for_rtr_refs(&self.announced, &self.withdrawn)
    }
}

fn build_payload_updates_for_rtr_refs<'a>(
    announced: &'a [Payload],
    withdrawn: &'a [Payload],
) -> Vec<(bool, &'a Payload)> {
    let mut updates = Vec::with_capacity(announced.len() + withdrawn.len());

    updates.extend(announced.iter().map(|p| (true, p)));
    updates.extend(withdrawn.iter().map(|p| (false, p)));

    updates.sort_by(|(a_upd, a_payload), (b_upd, b_payload)| {
        compare_payload_update_for_rtr(a_payload, *a_upd, b_payload, *b_upd)
    });

    updates
}

fn build_payload_updates_for_rtr(
    announced: &[Payload],
    withdrawn: &[Payload],
) -> Vec<(bool, Payload)> {
    let mut updates = Vec::with_capacity(announced.len() + withdrawn.len());

    updates.extend(announced.iter().cloned().map(|p| (true, p)));
    updates.extend(withdrawn.iter().cloned().map(|p| (false, p)));

    updates.sort_by(|(a_upd, a_payload), (b_upd, b_payload)| {
        compare_payload_update_for_rtr(a_payload, *a_upd, b_payload, *b_upd)
    });

    updates
}

fn build_snapshot_payloads_for_rtr(snapshot: &Snapshot) -> Vec<Payload> {
    let mut payloads = snapshot.payloads();
    sort_payloads_for_rtr(&mut payloads, true);
    payloads
}

fn dedup_payloads(payloads: &mut Vec<Payload>) {
    payloads.sort();
    payloads.dedup();
}

fn normalize_aspas<I>(aspas: I) -> Vec<Aspa>
where
    I: IntoIterator<Item = Aspa>,
{
    let mut by_customer = BTreeMap::<u32, Vec<_>>::new();

    for aspa in aspas {
        let providers = by_customer
            .entry(aspa.customer_asn().into_u32())
            .or_default();
        providers.extend(aspa.provider_asns().iter().copied());
    }

    let mut normalized = by_customer
        .into_iter()
        .map(|(customer_asn, providers)| {
            Aspa::new(customer_asn.into(), providers)
        })
        .collect::<Vec<_>>();
    normalized.sort();
    normalized
}

fn diff_aspas(
    current: &[Aspa],
    next: &[Aspa],
    announced: &mut Vec<Payload>,
    withdrawn: &mut Vec<Payload>,
) {
    let current = current
        .iter()
        .map(|aspa| (aspa.customer_asn().into_u32(), aspa))
        .collect::<BTreeMap<_, _>>();
    let next = next
        .iter()
        .map(|aspa| (aspa.customer_asn().into_u32(), aspa))
        .collect::<BTreeMap<_, _>>();

    let customers = current
        .keys()
        .chain(next.keys())
        .copied()
        .collect::<std::collections::BTreeSet<_>>();

    for customer in customers {
        match (current.get(&customer), next.get(&customer)) {
            (None, Some(new_aspa)) => announced.push(Payload::Aspa((*new_aspa).clone())),
            (Some(old_aspa), None) => withdrawn.push(Payload::Aspa((*old_aspa).clone())),
            (Some(old_aspa), Some(new_aspa)) if old_aspa != new_aspa => {
                announced.push(Payload::Aspa((*new_aspa).clone()));
            }
            _ => {}
        }
    }
}

fn sorted_dedup<T: Ord>(mut items: Vec<T>) -> Vec<T> {
    items.sort();
    items.dedup();
    items
}

fn diff_sorted<T, F>(
    current: &[T],
    next: &[T],
    announced: &mut Vec<Payload>,
    withdrawn: &mut Vec<Payload>,
    wrap: F,
) where
    T: Ord + Clone,
    F: Fn(T) -> Payload,
{
    let mut i = 0usize;
    let mut j = 0usize;
    while i < current.len() && j < next.len() {
        match current[i].cmp(&next[j]) {
            std::cmp::Ordering::Less => {
                withdrawn.push(wrap(current[i].clone()));
                i += 1;
            }
            std::cmp::Ordering::Greater => {
                announced.push(wrap(next[j].clone()));
                j += 1;
            }
            std::cmp::Ordering::Equal => {
                i += 1;
                j += 1;
            }
        }
    }
    while i < current.len() {
        withdrawn.push(wrap(current[i].clone()));
        i += 1;
    }
    while j < next.len() {
        announced.push(wrap(next[j].clone()));
        j += 1;
    }
}
