use std::{
    collections::HashMap,
    sync::{atomic::AtomicU64, Mutex},
    thread::ThreadId,
    time::Instant,
};

mod perfetto {
    include!(concat!(env!("OUT_DIR"), "/perfetto.trace.rs"));
}

use perfetto::DebugAnnotation;
use prost::Message;
use tracing::{field::Visit, span, Dispatch, Level, Metadata, Subscriber};

pub fn install_perfetto_subscriber() {
    let subscriber = PerfettoSubscriber::new();
    let dispatcher = Dispatch::new(subscriber);
    tracing::dispatcher::set_global_default(dispatcher).unwrap();
}

pub struct PerfettoSubscriber {
    id_counter: AtomicU64,
    start_time: Instant,
    inner: Mutex<InnerPerfettoSubscriber>,
}

struct SpanInfo {
    debug_annotations: Vec<DebugAnnotation>,
    level: Level,
    name: String,
}

struct InnerPerfettoSubscriber {
    buffer: Vec<perfetto::TracePacket>,
    spans: HashMap<u64, SpanInfo>,
    threads: HashMap<ThreadId, u64>,
}

impl InnerPerfettoSubscriber {
    fn new(buffer: Vec<perfetto::TracePacket>) -> Self {
        Self {
            buffer,
            spans: HashMap::new(),
            threads: HashMap::new(),
        }
    }

    fn get_track_id(&mut self) -> u64 {
        let current = std::thread::current();
        let nextfree = self.threads.len() as u64;
        match self.threads.entry(current.id()) {
            std::collections::hash_map::Entry::Occupied(occupied_entry) => *occupied_entry.get(),
            std::collections::hash_map::Entry::Vacant(vacant_entry) => {
                vacant_entry.insert(nextfree);

                self.buffer.push(perfetto::TracePacket {
                    timestamp: None,
                    data: Some(perfetto::trace_packet::Data::TrackDescriptor(
                        perfetto::TrackDescriptor {
                            uuid: Some(nextfree),
                            parent_uuid: None,
                            static_or_dynamic_name: None,
                            process: None,
                            thread: Some(perfetto::ThreadDescriptor {
                                pid: Some(1234),
                                thread_name: current.name().map(|s| s.to_owned()),
                                tid: Some(nextfree as i32),
                            }),
                        },
                    )),
                    optional_trusted_packet_sequence_id: Some(
                    perfetto::trace_packet::OptionalTrustedPacketSequenceId::TrustedPacketSequenceId(
                        3903809,
                    ),
                ),
            });

                nextfree
            }
        }
    }
}

struct DebugAnnotationVisitor {
    data: Vec<DebugAnnotation>,
}

impl DebugAnnotationVisitor {
    fn new_with_metadata(metadata: &'static Metadata<'static>) -> Self {
        Self {
            data: vec![
                DebugAnnotation {
                    name_field: Some(perfetto::debug_annotation::NameField::Name(String::from(
                        "module_path",
                    ))),
                    value: metadata
                        .module_path()
                        .map(|m| perfetto::debug_annotation::Value::StringValue(m.to_owned())),
                },
                DebugAnnotation {
                    name_field: Some(perfetto::debug_annotation::NameField::Name(String::from(
                        "location",
                    ))),
                    value: Some(perfetto::debug_annotation::Value::StringValue(format!(
                        "{}: {}",
                        metadata.file().unwrap_or(""),
                        metadata.line().unwrap_or(0)
                    ))),
                },
            ],
        }
    }

    fn new() -> Self {
        Self { data: Vec::new() }
    }
}

impl Visit for DebugAnnotationVisitor {
    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
        self.data.push(DebugAnnotation {
            name_field: Some(perfetto::debug_annotation::NameField::Name(
                field.name().to_owned(),
            )),
            value: Some(perfetto::debug_annotation::Value::StringValue(format!(
                "{value:?}"
            ))),
        });
    }

    fn record_i64(&mut self, field: &tracing::field::Field, value: i64) {
        self.data.push(DebugAnnotation {
            name_field: Some(perfetto::debug_annotation::NameField::Name(
                field.name().to_owned(),
            )),
            value: Some(perfetto::debug_annotation::Value::IntValue(value)),
        });
    }

    fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
        self.data.push(DebugAnnotation {
            name_field: Some(perfetto::debug_annotation::NameField::Name(
                field.name().to_owned(),
            )),
            value: Some(perfetto::debug_annotation::Value::UintValue(value)),
        });
    }

    fn record_f64(&mut self, field: &tracing::field::Field, value: f64) {
        self.data.push(DebugAnnotation {
            name_field: Some(perfetto::debug_annotation::NameField::Name(
                field.name().to_owned(),
            )),
            value: Some(perfetto::debug_annotation::Value::DoubleValue(value)),
        });
    }

    fn record_bool(&mut self, field: &tracing::field::Field, value: bool) {
        self.data.push(DebugAnnotation {
            name_field: Some(perfetto::debug_annotation::NameField::Name(
                field.name().to_owned(),
            )),
            value: Some(perfetto::debug_annotation::Value::BoolValue(value)),
        });
    }

    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
        self.data.push(DebugAnnotation {
            name_field: Some(perfetto::debug_annotation::NameField::Name(
                field.name().to_owned(),
            )),
            value: Some(perfetto::debug_annotation::Value::StringValue(
                value.to_owned(),
            )),
        });
    }
}

impl PerfettoSubscriber {
    pub fn new() -> Self {
        let mut packets = Vec::new();
        packets.push(perfetto::TracePacket {
            timestamp: None,
            data: Some(perfetto::trace_packet::Data::TrackDescriptor(
                perfetto::TrackDescriptor {
                    uuid: Some(894893984),
                    parent_uuid: None,
                    static_or_dynamic_name: None,
                    process: Some(perfetto::ProcessDescriptor {
                        pid: Some(1234),
                        process_name: Some(String::from("process")),
                    }),
                    thread: None,
                },
            )),
            optional_trusted_packet_sequence_id: Some(
                perfetto::trace_packet::OptionalTrustedPacketSequenceId::TrustedPacketSequenceId(
                    3903809,
                ),
            ),
        });
        Self {
            id_counter: AtomicU64::new(1),
            inner: Mutex::new(InnerPerfettoSubscriber::new(packets)),
            start_time: Instant::now(),
        }
    }

    fn flush(&self) {
        let data = self.inner.lock().unwrap();

        let trace = perfetto::Trace {
            packet: data.buffer.clone(),
        };
        let buf = trace.encode_to_vec();

        let _ = std::fs::write("out.trace", buf);
    }
}

impl Subscriber for PerfettoSubscriber {
    fn enabled(&self, metadata: &tracing::Metadata<'_>) -> bool {
        true
    }

    fn new_span(&self, span: &span::Attributes<'_>) -> span::Id {
        // println!("new_span {:?}", span);
        let id = self
            .id_counter
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);

        let mut debug_annotations = DebugAnnotationVisitor::new_with_metadata(span.metadata());
        span.record(&mut debug_annotations);

        let mut inner = self.inner.lock().unwrap();
        inner.spans.insert(
            id,
            SpanInfo {
                debug_annotations: debug_annotations.data,
                level: *span.metadata().level(),
                name: span.metadata().name().to_owned(),
            },
        );
        span::Id::from_u64(id)
    }

    fn record(&self, span: &span::Id, values: &span::Record<'_>) {
        // println!("record {:?} {:?}", span, values);
        let mut debug_annotations = DebugAnnotationVisitor::new();
        values.record(&mut debug_annotations);

        let mut inner = self.inner.lock().unwrap();
        inner
            .spans
            .get_mut(&span.into_u64())
            .unwrap()
            .debug_annotations
            .append(&mut debug_annotations.data);
    }

    fn record_follows_from(&self, span: &span::Id, follows: &span::Id) {
        // println!("record_follows_from {:?} {:?}", span, follows);
    }

    fn event(&self, event: &tracing::Event<'_>) {
        // println!("event {:?}", event);
        {
            let time = self.start_time.elapsed().as_nanos() as u64;
            let mut inner = self.inner.lock().unwrap();

            let mut debug_annotations = DebugAnnotationVisitor::new_with_metadata(event.metadata());
            event.record(&mut debug_annotations);

            let track_id = inner.get_track_id();
            inner.buffer.push(perfetto::TracePacket {
            timestamp: Some(time),
            data: Some(perfetto::trace_packet::Data::TrackEvent(
                perfetto::TrackEvent {
                    track_uuid: Some(track_id),
                    categories: vec![event.metadata().level().as_str().to_owned()],
                    category_iids: vec![],
                    name_field: Some(perfetto::track_event::NameField::Name(String::from(event.metadata().name()))),
                    r#type: Some(perfetto::track_event::Type::Instant as i32),
            debug_annotations: debug_annotations.data,
                },
            )),
            optional_trusted_packet_sequence_id: Some(
                perfetto::trace_packet::OptionalTrustedPacketSequenceId::TrustedPacketSequenceId(
                    3903809,
                ),
            ),
        });
        }
        self.flush();
    }

    fn enter(&self, span: &span::Id) {
        // println!("enter {:?}", span);
        {
            let time = self.start_time.elapsed().as_nanos() as u64;
            let mut inner = self.inner.lock().unwrap();

            let debug_annotations = inner.spans[&span.into_u64()].debug_annotations.clone();
            let level = inner.spans[&span.into_u64()].level;
            let name = inner.spans[&span.into_u64()].name.clone();

            let track_id = inner.get_track_id();
            inner.buffer.push(perfetto::TracePacket {
            timestamp: Some(time),
            data: Some(perfetto::trace_packet::Data::TrackEvent(
                perfetto::TrackEvent {
                    track_uuid: Some(track_id),
                    categories: vec![level.as_str().to_owned()],
                    category_iids: vec![],
                    name_field: Some(perfetto::track_event::NameField::Name(name)),
                    r#type: Some(perfetto::track_event::Type::SliceBegin as i32),
                    debug_annotations,
                },
            )),
            optional_trusted_packet_sequence_id: Some(
                perfetto::trace_packet::OptionalTrustedPacketSequenceId::TrustedPacketSequenceId(
                    3903809,
                ),
            ),
        });
        }
        self.flush();
    }

    fn exit(&self, span: &span::Id) {
        // println!("exit {:?}", span);
        {
            let time = self.start_time.elapsed().as_nanos() as u64;
            let mut inner = self.inner.lock().unwrap();

            let debug_annotations = inner.spans[&span.into_u64()].debug_annotations.clone();
            let level = inner.spans[&span.into_u64()].level;
            let name = inner.spans[&span.into_u64()].name.clone();

            let track_id = inner.get_track_id();
            inner.buffer.push(perfetto::TracePacket {
            timestamp: Some(time),
            data: Some(perfetto::trace_packet::Data::TrackEvent(
                perfetto::TrackEvent {
                    track_uuid: Some(track_id),
                    categories: vec![level.as_str().to_owned()],
                    category_iids: vec![],
                    name_field: Some(perfetto::track_event::NameField::Name(name)),
                    r#type: Some(perfetto::track_event::Type::SliceEnd as i32),
                    debug_annotations,
                },
            )),
            optional_trusted_packet_sequence_id: Some(
                perfetto::trace_packet::OptionalTrustedPacketSequenceId::TrustedPacketSequenceId(
                    3903809,
                ),
            ),
        });
        }
        self.flush();
    }
}