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

use crate::audit::{
    AuditDownloadEvent, AuditDownloadKind, AuditDownloadKindStats, AuditDownloadObjectsStat,
    AuditDownloadStats,
};

#[derive(Clone, Debug, Default)]
pub struct DownloadLogHandle {
    inner: Arc<Mutex<Vec<AuditDownloadEvent>>>,
}

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

    pub fn record_event(&self, event: AuditDownloadEvent) {
        self.inner.lock().expect("download log lock").push(event);
    }

    pub fn snapshot_events(&self) -> Vec<AuditDownloadEvent> {
        self.inner.lock().expect("download log lock").clone()
    }

    pub fn stats_from_events(events: &[AuditDownloadEvent]) -> AuditDownloadStats {
        let mut out = AuditDownloadStats {
            events_total: events.len() as u64,
            by_kind: BTreeMap::new(),
        };
        for e in events {
            let kind_key = match e.kind {
                AuditDownloadKind::RrdpNotification => "rrdp_notification",
                AuditDownloadKind::RrdpSnapshot => "rrdp_snapshot",
                AuditDownloadKind::RrdpDelta => "rrdp_delta",
                AuditDownloadKind::Rsync => "rsync",
            }
            .to_string();

            let st = out
                .by_kind
                .entry(kind_key)
                .or_insert_with(|| AuditDownloadKindStats {
                    ok_total: 0,
                    fail_total: 0,
                    duration_ms_total: 0,
                    bytes_total: None,
                    objects_count_total: None,
                    objects_bytes_total: None,
                });
            if e.success {
                st.ok_total = st.ok_total.saturating_add(1);
            } else {
                st.fail_total = st.fail_total.saturating_add(1);
            }
            st.duration_ms_total = st.duration_ms_total.saturating_add(e.duration_ms);
            if let Some(b) = e.bytes {
                st.bytes_total = Some(st.bytes_total.unwrap_or(0).saturating_add(b));
            }
            if let Some(objects) = &e.objects {
                st.objects_count_total = Some(
                    st.objects_count_total
                        .unwrap_or(0)
                        .saturating_add(objects.objects_count),
                );
                st.objects_bytes_total = Some(
                    st.objects_bytes_total
                        .unwrap_or(0)
                        .saturating_add(objects.objects_bytes_total),
                );
            }
        }
        out
    }

    pub fn stats(&self) -> AuditDownloadStats {
        let events = self.snapshot_events();
        Self::stats_from_events(&events)
    }

    pub fn span_download<'a>(
        &'a self,
        kind: AuditDownloadKind,
        uri: &'a str,
    ) -> DownloadSpanGuard<'a> {
        DownloadSpanGuard {
            handle: self,
            kind,
            uri,
            start_instant: Instant::now(),
            started_at: time::OffsetDateTime::now_utc(),
            bytes: None,
            objects: None,
            error: None,
            success: None,
        }
    }
}

pub struct DownloadSpanGuard<'a> {
    handle: &'a DownloadLogHandle,
    kind: AuditDownloadKind,
    uri: &'a str,
    start_instant: Instant,
    started_at: time::OffsetDateTime,
    bytes: Option<u64>,
    objects: Option<AuditDownloadObjectsStat>,
    error: Option<String>,
    success: Option<bool>,
}

impl DownloadSpanGuard<'_> {
    pub fn set_bytes(&mut self, bytes: u64) {
        self.bytes = Some(bytes);
    }

    pub fn set_objects(&mut self, objects_count: u64, objects_bytes_total: u64) {
        self.objects = Some(AuditDownloadObjectsStat {
            objects_count,
            objects_bytes_total,
        });
    }

    pub fn set_ok(&mut self) {
        self.success = Some(true);
    }

    pub fn set_err(&mut self, msg: impl Into<String>) {
        self.success = Some(false);
        self.error = Some(msg.into());
    }
}

impl Drop for DownloadSpanGuard<'_> {
    fn drop(&mut self) {
        use time::format_description::well_known::Rfc3339;
        let finished_at = time::OffsetDateTime::now_utc();
        let dur = self.start_instant.elapsed();
        let duration_ms = duration_to_ms(dur);
        let started_at_rfc3339_utc = self
            .started_at
            .to_offset(time::UtcOffset::UTC)
            .format(&Rfc3339)
            .unwrap_or_else(|_| "<format-error>".to_string());
        let finished_at_rfc3339_utc = finished_at
            .to_offset(time::UtcOffset::UTC)
            .format(&Rfc3339)
            .unwrap_or_else(|_| "<format-error>".to_string());
        let success = self.success.unwrap_or(false);
        let event = AuditDownloadEvent {
            kind: self.kind.clone(),
            uri: self.uri.to_string(),
            started_at_rfc3339_utc,
            finished_at_rfc3339_utc,
            duration_ms,
            success,
            error: if success { None } else { self.error.clone() },
            bytes: self.bytes,
            objects: self.objects.clone(),
        };
        self.handle.record_event(event);
    }
}

fn duration_to_ms(d: Duration) -> u64 {
    let ms = d.as_millis();
    ms.min(u128::from(u64::MAX)) as u64
}
