use tracing::{Subscriber, Event};
use tracing::metadata::Metadata;
use tracing::span::{Attributes, Id, Record as TracingRecord};
use std::collections::{HashMap, hash_map::Entry};
use std::sync::Mutex;
use std::time::{Instant, Duration};
use rand::Rng;

struct SpanData {
    name: &'static str,
    enter_time: Option<Instant>,
}
struct SpanTimerSubscriber {
    spans: Mutex<HashMap<Id, SpanData>>,
}
impl SpanTimerSubscriber {
    pub fn new() -> Self {
        SpanTimerSubscriber {
            spans: Mutex::new(HashMap::new())
        }
    }
}
impl tracing::Subscriber for SpanTimerSubscriber {
    fn enabled(&self, metadata: &tracing::metadata::Metadata<'_>) -> bool {
        true
    }

    fn new_span(&self, span: &Attributes<'_>) -> Id {
        let span_data = SpanData {
            name: span.metadata().name(),
            enter_time: None,
        };
        let mut rng = rand::thread_rng();
        // will panic if you get 0 or a repeat, unlucky!
        let id = Id::from_u64(rng.gen::<u64>());

        let mut mutex_guard = self.spans.lock().unwrap();
        mutex_guard.insert(id.clone(), span_data);
        id
    }

    fn record(&self, _span: &Id, _values: &TracingRecord<'_>) {}

    fn record_follows_from(&self, _span: &Id, _follows: &Id) {}

    fn event(&self, _event: &Event<'_>) {}

    fn enter(&self, span: &Id) {
        let mut hash_map = self.spans.lock().unwrap();
        match hash_map.entry(span.clone()) {
            Entry::Occupied(mut entry) => {
                let entry = entry.get_mut();
                if entry.enter_time.is_some() {
                    panic!("Entered twice");
                }
                entry.enter_time = Some(Instant::now());
            }
            _ => panic!("Expected occupied entry for {:?}", span),
        }
    }

    fn exit(&self, span: &Id) {
        let mut hash_map = self.spans.lock().unwrap();
        match hash_map.entry(span.clone()) {
            Entry::Occupied(mut entry) => {
                let entry = entry.get_mut();
                match entry.enter_time {
                    None => panic!("Exited twice"),
                    Some(enter_time) => {
                        let exit_time = Instant::now();
                        let duration = exit_time.duration_since(enter_time);
                        if duration > Duration::from_millis(500) {
                            println!("SPAN: '{}', DURATION: '{:?}'", entry.name, duration);
                        }
                    }
                }
                entry.enter_time = None;
            }
            _ => panic!("Expected occupied exit for {:?}", span),
        }
    }
}