use crate::function_trace::*;
use color_eyre::eyre::{Result, WrapErr};
use fnv::{FnvHashMap, FnvHashSet};
use indexmap::set::IndexSet;
use serde::Serialize;
use serde_repr::Serialize_repr;
use std::borrow::Cow;
use std::collections::HashMap;
use std::fs::File;
use std::io::BufWriter;
use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{Duration, SystemTime};

/// Number of ms since `meta.start_duration`.
type FirefoxTime = f64;

// TODO: Should we be creating these?
#[derive(Serialize, Debug)]
enum Unused {}

type StringIndex = usize;

/// Represents an opaque promise that a thread can be registered with the given id, including some
/// information internally necessary for doing so.
pub struct FirefoxProfileThreadId {
    /// A unique id representing this thread.  It has no correlation to Unix thread ids.
    pub tid: usize,

    /// The initial system time for all tracing in this profile.
    /// This allows threads to emit offset times based on a single master time for the profile.
    start_duration: Duration,
}

/// The top-level structure representing a profile that can be consumed by the Firefox profiler.
/// This will be exported mostly as-is to JSON at the end of a tracing session.
///
/// Based off of <https://github.com/firefox-devtools/profiler/blob/main/src/types/profile.js>.
#[derive(Serialize, Debug)]
pub struct FirefoxProfile {
    meta: Metadata,
    counters: Vec<Unused>,
    threads: Vec<FirefoxThread>,
    libs: Vec<Unused>,

    ///////////////////////////////////////////////////////////////////////////
    // Internal
    ///////////////////////////////////////////////////////////////////////////
    // The number of registered threads.
    #[serde(skip_serializing)]
    thread_count: usize,
}

/// Represent a function call in a way that can either be cheaply created (for [`HashMap::get`]) or
/// allocated on the heap (for [`HashMap::insert`]).
#[derive(Hash, Debug, PartialEq, Eq)]
enum FunctionCall<'a> {
    Native {
        function_name: Cow<'a, str>,
        module: Cow<'a, str>,
    },
    Python {
        function_name: Cow<'a, str>,
        module: Cow<'a, str>,
        line: u32,
    },
}

///////////////////////////////////////////////////////////////////////////////
// Event Categories
///////////////////////////////////////////////////////////////////////////////
// https://github.com/firefox-devtools/profiler/blob/78070314c3f0d5f4d51d9dd1e44d11a5e242d921/src/utils/colors.js
#[derive(Serialize, Debug)]
#[serde(rename_all = "lowercase")]
#[allow(dead_code)]
enum Color {
    Transparent,
    Purple,
    Green,
    Orange,
    Yellow,
    LightBlue,
    Grey,
    Blue,
    Brown,
}

#[derive(Serialize, Debug)]
struct Category {
    name: &'static str,
    color: Color,
    subcategories: [&'static str; 1],
}
#[derive(Serialize_repr, Debug, Clone, Copy)]
#[repr(usize)]
enum CategoryIndex {
    Logs,
    Python,
    Native,
    Exceptions,
    Imports,
    GarbageCollection,
    Length,
}

const CATEGORIES: [Category; CategoryIndex::Length as usize] = [
    Category {
        name: "Logs",
        color: Color::Grey,
        subcategories: ["Unused"],
    },
    Category {
        name: "Python",
        color: Color::Orange,
        subcategories: ["Code"],
    },
    Category {
        name: "Native",
        color: Color::LightBlue,
        subcategories: ["Code"],
    },
    Category {
        name: "Exceptions",
        color: Color::Brown,
        subcategories: ["Unused"],
    },
    Category {
        name: "Imports",
        color: Color::Grey,
        subcategories: ["Unused"],
    },
    Category {
        name: "GarbageCollection",
        color: Color::Purple,
        subcategories: ["Unused"],
    },
];

///////////////////////////////////////////////////////////////////////////////
// Tables
///////////////////////////////////////////////////////////////////////////////
// There are many table structures where the profile is expected to output struct-of-lists rather
// than list-of-structs in order to save on memory allocation overhead in JS.
// Rather than dealing with all of these manually, this macro allows us to quickly define the
// structs as well as useful helper methods.
macro_rules! struct_table {
    (struct $struct_name: ident {
        $($element: ident: $ty: ty),+
    } $($additional_field:ident: $additional_type: ty),*) => {
        paste::item! {
            #[derive(Serialize, Debug, Copy, Clone, PartialEq, Eq, Hash, Default)]
            struct [<$struct_name Index>] (usize);

            #[derive(Serialize, Debug, Default)]
            #[serde(rename_all = "camelCase")]
            struct [<$struct_name Table>] {
                $($element: Vec<$ty>,)*
                $($additional_field: $additional_type,)*
                length: usize
            }

            struct $struct_name {
                $($element: $ty),*,
            }

            impl [<$struct_name Table>] {
                #[allow(dead_code)]
                fn new($([<$additional_field _arg>]: $additional_type,)*) -> Self {
                    Self {
                        $($element: Vec::new()),*,
                        $($additional_field: [<$additional_field _arg>],)*
                        length: 0
                    }
                }

                #[allow(dead_code)]
                fn add(&mut self, arg: $struct_name) -> [<$struct_name Index>] {
                    let length = self.length;
                    $(self.$element.push(arg.$element);)*
                    self.length += 1;

                    [<$struct_name Index>](length)
                }

                // Sometimes we want to skip serializing if the table is empty.
                #[allow(dead_code)]
                fn is_empty(&self) -> bool {
                    self.length == 0
                }

                $(
                // Support typesafe-ish accessors for elements.
                #[allow(dead_code)]
                fn [<get_ $element>](&self, index: [<$struct_name Index>]) -> &$ty {
                    &self.$element[index.0]
                })*
            }
        }
    }
}

struct_table! {
    struct Sample {
        event_delay: Option<u32>,
        stack: StackIndex,
        time: FirefoxTime,
        weight: Option<FirefoxTime>
        // TODO: We can draw CPU usage graphs with this, but we don't currently have enough
        // information to do that.
        // threadCPUDelta: f32
    }
    weight_type: &'static str
}

struct_table! {
    struct Stack {
        frame: FrameIndex,
        category: CategoryIndex,
        subcategory: u32,
        prefix: Option<StackIndex>
    }
}
struct_table! {
    struct Frame {
        address: Option<u32>, // TODO: What's the proper type here?
        inlineDepth: u32,
        category: CategoryIndex,
        subcategory: u32,
        func: FunctionIndex,
        nativeSymbol: Option<Unused>,
        innerWindowID: Option<Unused>,
        implementation: Option<Unused>,
        line: Option<u32>,
        column: Option<u32>
    }
}
struct_table! {
    struct Function {
        isJS: bool,
        name: StringIndex,
        resource: i32,
        relevantForJS: bool,
        file_name: StringIndex, // Or module-name for native functions
        line_number: Option<u32>,
        column_number: Option<u32>
    }
}
struct_table! {
    struct NativeAllocation {
        time: FirefoxTime,
        weight: isize,  // This is actually bytes
        stack: Option<StackIndex>,
        memory_address: usize,
        thread_id: usize
    }
    weight_type: &'static str
}
struct_table! {
    struct NativeSymbol {
        libIndex: Option<Unused>,
        address: Option<u32>,
        name: StringIndex,
        functionSize: Option<u32>
    }
}
struct_table! {
    struct Marker {
        data: serde_json::Value,
        name: StringIndex,
        startTime: FirefoxTime,
        endTime: Option<FirefoxTime>,
        phase: u64, // TODO: This seems interesting - how can we use it?
        category: CategoryIndex
    }
}
struct_table! {
    struct Resource {
        lib: Unused,
        name: Unused,
        host: Unused,
        r#type: Unused
    }
}

///////////////////////////////////////////////////////////////////////////////
// Markers
///////////////////////////////////////////////////////////////////////////////
// Based off of https://github.com/firefox-devtools/profiler/blob/main/src/types/markers.js
#[derive(Serialize, Debug)]
#[serde(rename_all = "camelCase")]
struct MarkerSchema {
    name: &'static str,
    tooltip_label: &'static str,
    table_label: &'static str,
    chart_label: &'static str,
    display: &'static [&'static str],
    data: &'static [MarkerDataSchema],
}

#[derive(Serialize, Debug)]
struct MarkerDataSchema {
    key: &'static str,
    label: &'static str,
    format: &'static str,
    searchable: bool,
}

const MARKER_DISPLAYS_ALL: &[&str] = &[
    "marker-chart",
    "marker-table",
    "timeline-overview",
    "stack-chart",
];

const MARKER_SCHEMAS_SUPPORTED: &[MarkerSchema] = &[
    MarkerSchema {
        name: "Log",
        tooltip_label: "Log event from {marker.data.origin}",
        chart_label: "{marker.data.origin}",
        table_label: "{marker.data.origin}: {marker.data.value}",
        display: MARKER_DISPLAYS_ALL,
        data: &[
            MarkerDataSchema {
                key: "origin",
                label: "Log Origin",
                format: "string",
                searchable: true,
            },
            MarkerDataSchema {
                key: "value",
                label: "Value",
                format: "string",
                searchable: true,
            },
        ],
    },
    MarkerSchema {
        name: "Import",
        tooltip_label: "Imported {marker.data.module}",
        chart_label: "{marker.data.module}",
        table_label: "imported {marker.data.module}",
        display: MARKER_DISPLAYS_ALL,
        data: &[MarkerDataSchema {
            key: "module",
            label: "Imported Module",
            format: "string",
            searchable: true,
        }],
    },
    MarkerSchema {
        name: "Exception",
        tooltip_label: "Exception {marker.data.exception} in {marker.data.module}",
        chart_label: "{marker.data.exception}",
        table_label: "Exception {marker.data.exception}: {marker.data.value}",
        display: MARKER_DISPLAYS_ALL,
        data: &[
            MarkerDataSchema {
                key: "exception",
                label: "Exception Type",
                format: "string",
                searchable: true,
            },
            MarkerDataSchema {
                key: "value",
                label: "Exception Data",
                format: "string",
                searchable: true,
            },
            MarkerDataSchema {
                key: "module",
                label: "Module",
                format: "string",
                searchable: true,
            },
        ],
    },
];

///////////////////////////////////////////////////////////////////////////////
// Top-Level Structures
///////////////////////////////////////////////////////////////////////////////
#[derive(Serialize, Debug)]
#[serde(rename_all = "camelCase")]
struct Metadata {
    ///////////////////////////////////////////////////////////////////////////
    // External API
    ///////////////////////////////////////////////////////////////////////////
    interval: f32,
    /// Number of ms since Unix epoch time
    start_time: f64,
    /// Should be 0
    process_type: u32,
    categories: &'static [Category],
    /// Should be 0
    stackwalk: u32,
    debug: bool,
    version: u32,
    preprocessed_profile_version: u32,
    symbolicated: bool,

    product: String,
    #[serde(rename = "appBuildID")]
    app_build_id: String,
    abi: String,
    platform: String,
    misc: String,

    #[serde(rename = "physicalCPUs")]
    physical_cpus: usize,
    #[serde(rename = "logicalCPUs")]
    logical_cpus: usize,

    marker_schema: &'static [MarkerSchema],

    ///////////////////////////////////////////////////////////////////////////
    // Internal API
    ///////////////////////////////////////////////////////////////////////////
    /// The duration corresponding to `self.start_time`.
    #[serde(skip_serializing)]
    start_duration: Duration,
}

/// The `FirefoxProfile` representation of a thread, with additional internal augmentations to
/// store data necessary for efficiently parsing the thread's trace log.
///
/// NOTE: This is designed to be usable entirely standalone from a `FirefoxProfile` during parsing.
#[derive(Serialize, Debug, Default)]
#[serde(rename_all = "camelCase")]
pub struct FirefoxThread {
    ///////////////////////////////////////////////////////////////////////////
    // External API
    ///////////////////////////////////////////////////////////////////////////
    process_type: &'static str,
    process_startup_time: FirefoxTime, // Always 0.0
    process_shutdown_time: FirefoxTime,
    register_time: FirefoxTime,
    unregister_time: Option<FirefoxTime>, // Always None
    paused_ranges: Vec<Unused>,
    name: &'static str,
    is_main_thread: bool,
    process_name: String,
    is_js_tracer: bool,
    pid: String,
    tid: usize,

    samples: SampleTable,
    markers: MarkerTable,
    stack_table: StackTable,
    frame_table: FrameTable,
    func_table: FunctionTable,
    string_array: IndexSet<String>,

    // If an empty allocation table is emitted, the profiler silently fails.
    #[serde(skip_serializing_if = "NativeAllocationTable::is_empty")]
    native_allocations: NativeAllocationTable,

    ///////////////////////////////////////////////////////////////////////////
    // External (Unused)
    ///////////////////////////////////////////////////////////////////////////
    resource_table: ResourceTable,
    native_symbols: NativeSymbolTable,

    ///////////////////////////////////////////////////////////////////////////
    // Internal
    ///////////////////////////////////////////////////////////////////////////
    /// A stack of sample indices, allowing calls/returns to know what their
    /// parent was.
    #[serde(skip_serializing)]
    calls: Vec<SampleIndex>,

    /// A mapping of function UIDs to the de-duplicated store in `func_table`.
    #[serde(skip_serializing)]
    functions: HashMap<FunctionCall<'static>, FunctionIndex>,

    /// The starting time for this thread, which all other times are supposed to be relative to.
    /// This is a copy of meta.start_duration.
    #[serde(skip_serializing)]
    start_duration: Duration,

    /// Track allocation address -> bytes.
    // TODO: This should probably be COW when we do multiprocess support,
    // and probably can't live on a Thread once we do multithread support.
    #[serde(skip_serializing)]
    allocations: FnvHashMap<usize, usize>,

    /// Track the set of Stacks (information about a (function, caller) pair) that exist,
    /// preventing us from emitting duplicates.
    #[serde(skip_serializing)]
    existing_stacks: FnvHashMap<(FunctionIndex, Option<StackIndex>), StackIndex>,
}

impl FirefoxProfile {
    /// Create a new `FirefoxProfile` representing a given trace.
    #[must_use]
    pub fn new(info: TraceInitialization) -> FirefoxProfile {
        FirefoxProfile {
            meta: Metadata {
                // Mark the time we received this message as the time the program started.
                // Additionally, the process being traced sent us a Duration that we should record.
                start_time: SystemTime::now()
                    .duration_since(SystemTime::UNIX_EPOCH)
                    .map_or(0.0, |n| n.as_secs_f64() * 1000.0),
                start_duration: info.time,
                interval: 0.000_001, // 1ns
                process_type: 0,
                categories: &CATEGORIES,
                stackwalk: 0,

                // Latest versions as of Aug 2023.
                // Follow the migration code at
                // https://github.com/firefox-devtools/profiler/blob/main/src/profile-logic/processed-profile-versioning.js
                // to update these.
                version: 27,
                preprocessed_profile_version: 47,
                debug: false,
                symbolicated: true,
                physical_cpus: num_cpus::get_physical(),
                logical_cpus: num_cpus::get(),

                product: info.program_name,
                app_build_id: info.program_version,
                abi: info.lang_version,
                platform: info.platform,
                misc: "<git revision>".to_string(),

                marker_schema: MARKER_SCHEMAS_SUPPORTED,
            },
            counters: Vec::new(),
            threads: Vec::new(),
            libs: Vec::new(),

            thread_count: 0,
        }
    }

    /// Register a new thread, allowing us to parse traces for it without needing ownership of the
    /// `FirefoxProfile` object.
    pub fn register_thread(&mut self) -> FirefoxProfileThreadId {
        self.thread_count += 1;

        FirefoxProfileThreadId {
            tid: self.thread_count,
            start_duration: self.meta.start_duration,
        }
    }

    /// Finalize this thread, attaching it to the profile.
    pub fn finalize_thread(&mut self, thread: FirefoxThread) {
        // Add this thread to the profile.
        self.threads.push(thread);
    }

    /// Export the recorded profile to the given output directory.
    pub fn export(&mut self, output_dir: PathBuf) -> Result<()> {
        // All threads have exited and we're the last one recording.
        // Sort threads by the time they started, to ensure earlier threads show higher in the
        // profile.
        self.threads.sort_by(|x, y| {
            x.register_time
                .partial_cmp(&y.register_time)
                .expect("Registration time is reasonable")
        });

        // Determine when the last thread exited, in case we need to update other threads to point
        // to that.
        let end_time = self
            .threads
            .iter()
            .map(|x| x.process_shutdown_time)
            .fold(0.0, f64::max);

        // The set of pids we've seen so far.  This tells us whether a thread that is registering
        // should be the main thread for its pid.
        // NOTE: This assumes pids aren't reused, as otherwise top-level processes will be attached as
        // threads under a dead process.
        let mut registered_pids = FnvHashSet::default();

        // Fixup some metadata since we have better ordering information now.
        for thread in &mut self.threads {
            thread.name = if registered_pids.contains(&thread.pid) {
                "thread"
            } else {
                thread.is_main_thread = true;
                "GeckoMain"
            };

            registered_pids.insert(thread.pid.clone());

            if thread.functions.is_empty() {
                // If we didn't log any calls, attach a long call to ensure the user will know this
                // occurred.
                // TODO: This should display as full width in the Stack Chart, but doesn't for some
                // reason.
                let message = "[WARNING] FunctionTrace didn't log any data for this thread.  This may be because the thread crashed or otherwise exited before flushing its trace buffer.";
                thread.emit_call(
                    message,
                    thread.start_duration + Duration::from_secs_f64(thread.register_time / 1000.0),
                    "FunctionTrace",
                    None,
                );
                thread.emit_return(message, end_time);
                thread.process_shutdown_time = end_time;
            }
        }

        // TODO: If we find empty threads with a parent process that has data, we should delete the
        // threads and add a marker about them to the parent process.

        // We gzip profiles to save some disk space when transferring files.  However, this can
        // take a significant amount of time, which may not be worth it, so it can be disabled with
        // this environment variable.
        let compressed = std::env::var("FUNCTIONTRACE_COMPRESSION")
            .map(|x| !matches!(x.as_str(), "false" | "no" | "disable"))
            .unwrap_or(true);

        // Pick a reasonably unique filename in the output directory to write the profile to.
        let output_path = {
            let now = chrono::Local::now();

            output_dir.join(format!(
                "functiontrace.{}.json{}",
                now.format("%F-%T.%3f"),
                if compressed { ".gz" } else { "" }
            ))
        };

        // Emit the current profile, rendering a spinner while waiting for it to compress.
        {
            use spinoff::spinners::{Arrow3, SpinnerFrames};

            // Track the number of bytes we've written.
            let bytes = AtomicUsize::new(0);

            let mut spinner = spinoff::Spinner::new(
                Arrow3,
                "Exporting FunctionTrace profile...",
                spinoff::Color::Blue,
            );
            let gzipped_bytes = std::thread::scope(|s| {
                // Spawn a background thread to keep the spinner updated, then kick off the JSON
                // encoding.
                s.spawn(|| {
                    let spin_interval = {
                        let spinner_frame = SpinnerFrames::from(Arrow3);

                        Duration::from_millis(
                            spinner_frame.frames.len() as u64 * spinner_frame.interval as u64,
                        )
                    };

                    loop {
                        std::thread::sleep(spin_interval);
                        let written = bytes.load(Ordering::Relaxed);

                        if written == usize::MAX {
                            break;
                        }

                        spinner.update_text(format!(
                            "Exporting FunctionTrace profile... {}",
                            bytesize::to_string(written as u64, false)
                        ));
                    }
                });

                // Note: Ordering on these streams matters - if we put the progress tracker in
                // front of [`BufWriter`], we'll end up getting notified every few bytes.
                let progress_writer = BufWriter::new(progress_streams::ProgressWriter::new(
                    File::create(&output_path)?,
                    |progress: usize| {
                        bytes.fetch_add(progress, Ordering::Relaxed);
                    },
                ));

                if compressed {
                    let gzip = flate2::write::GzEncoder::new(
                        progress_writer,
                        // Level 4 appears to be significantly faster than the default (6) on this
                        // format, while also only using a few percent more disk space.
                        flate2::Compression::new(4),
                    );

                    serde_json::to_writer(gzip, &self)?;
                } else {
                    serde_json::to_writer(progress_writer, &self)?;
                }

                // Notfiy the spinner thread that we're done by using [`usize::MAX`] as a flag
                // value.
                Ok::<_, std::io::Error>(bytes.swap(usize::MAX, Ordering::SeqCst))
            })?;

            spinner.success(&format!(
                "Exported FunctionTrace profile to {} ({})",
                output_path.display(),
                bytesize::to_string(gzipped_bytes as u64, false)
            ));
        };

        // Link `functiontrace.latest.json.gz` to this file.  We need to remove any existing symlink
        // before doing so.
        let symlink = output_dir.join(format!(
            "functiontrace.latest.json{}",
            if compressed { ".gz" } else { "" }
        ));
        let _ = std::fs::remove_file(&symlink);
        std::os::unix::fs::symlink(&output_path, &symlink)
            .wrap_err("Symlinking functiontrace.latest.json.gz failed")?;

        Ok(())
    }
}

impl FirefoxThread {
    /// Create a new `FirefoxThread`, given some information about it and a
    /// `FirefoxProfileThreadId` granting permission to do so.
    #[must_use]
    pub fn new(registration: ThreadRegistration, thread: &FirefoxProfileThreadId) -> FirefoxThread {
        FirefoxThread {
            process_type: "default",
            process_startup_time: 0.0,
            process_shutdown_time: 0.0, // Updated on each event
            register_time: (registration.time - thread.start_duration).as_secs_f64() * 1000.0,
            start_duration: thread.start_duration,
            unregister_time: None,
            name: "<filled in at export>",
            process_name: registration.program_name,
            is_js_tracer: true,
            pid: registration.pid.to_string(),
            tid: thread.tid,
            samples: SampleTable::new("tracing-ms"),
            native_allocations: NativeAllocationTable::new("bytes"),
            ..Default::default()
        }
    }

    /// Times come in as seconds but need to be converted to milliseconds and adjusted to an offset
    /// from `FirefoxThread::start_time`.
    fn offset_time(&self, time: Duration) -> FirefoxTime {
        (time - self.start_duration).as_secs_f64() * 1000.0
    }

    /// We use an `IndexSet` to store our strings, but this makes for slightly more complicated code
    /// if we're trying to avoid unnecessary allocations.  This is a small wrapper for those.
    fn unique_string(&mut self, string: &str) -> StringIndex {
        match self.string_array.get_full(string) {
            Some((index, _)) => index,
            None => self.string_array.insert_full(string.to_string()).0,
        }
    }

    fn emit_call(
        &mut self,
        name: &str,
        time: Duration,
        file_or_module: &str,
        line_number: Option<u32>,
    ) -> FirefoxTime {
        // Create a UID for this function to avoid creating multiple entries in the output for it.
        let (func_uid, category) = match line_number {
            // Only native functions don't have known line numbers.
            None => (
                FunctionCall::Native {
                    function_name: std::borrow::Cow::Borrowed(name),
                    module: std::borrow::Cow::Borrowed(file_or_module),
                },
                CategoryIndex::Native,
            ),
            Some(line) => (
                FunctionCall::Python {
                    function_name: std::borrow::Cow::Borrowed(name),
                    module: std::borrow::Cow::Borrowed(file_or_module),
                    line,
                },
                CategoryIndex::Python,
            ),
        };

        // Ensure we have a FunctionTable and FrameTable entry for this function.
        let function_id = match self.functions.get(&func_uid) {
            Some(&x) => x,
            None => {
                // We haven't seen this function before.  Reserve an index in the
                // FunctionTable for it.
                let funcname_id = self.unique_string(name);
                let filename_id = self.unique_string(file_or_module);

                // Create a function and frame entry for this function.
                // There's currently a one-to-one mapping of these.
                let function_id = self.func_table.add(Function {
                    isJS: false,
                    name: funcname_id,
                    resource: -1,
                    relevantForJS: false,
                    file_name: filename_id,
                    line_number,
                    column_number: None,
                });
                self.frame_table.add(Frame {
                    address: None,
                    inlineDepth: 0,
                    category,
                    subcategory: 0,
                    func: function_id,
                    nativeSymbol: None,
                    innerWindowID: None,
                    implementation: None,
                    line: line_number, // TODO: Is this the right line number?
                    column: None,
                });

                // We need to allocate to put this in the function table.
                let func_uid = match func_uid {
                    FunctionCall::Native { .. } => FunctionCall::Native {
                        function_name: std::borrow::Cow::Owned(name.to_string()),
                        module: std::borrow::Cow::Owned(file_or_module.to_string()),
                    },
                    FunctionCall::Python { line, .. } => FunctionCall::Python {
                        function_name: std::borrow::Cow::Owned(name.to_string()),
                        module: std::borrow::Cow::Owned(file_or_module.to_string()),
                        line,
                    },
                };

                self.functions.insert(func_uid, function_id);
                function_id
            }
        };
        // Frames and functions must be equivalent.
        let frame_id = FrameIndex(function_id.0);

        // This stack belongs under the previous caller.  To find it, we look at the last caller
        // sample to find the stack attached to that sample.
        let prefix = self.calls.last().map(|&x| *self.samples.get_stack(x));

        // Ensure we only emit one StackTable entry per stack, as there can be duplicates caused by
        // things like loops.
        // This will increase profile generation memory usage, but decreases the emitted profile
        // size by a substantial amount (>25% on some testcases).
        let stack_id = match self.existing_stacks.get(&(function_id, prefix)) {
            Some(&id) => id,
            None => {
                let stack_id = self.stack_table.add(Stack {
                    frame: frame_id,
                    subcategory: 0,
                    category,
                    prefix,
                });
                self.existing_stacks.insert((function_id, prefix), stack_id);
                stack_id
            }
        };

        // Emit a sample and a stacktrace.
        let time = self.offset_time(time);
        let samples_id = self.samples.add(Sample {
            event_delay: Some(0),
            stack: stack_id,
            time,
            weight: None, // We'll overwrite this during the return.
        });
        self.calls.push(samples_id);

        time
    }

    fn emit_return(&mut self, func_name: &str, time: FirefoxTime) -> FirefoxTime {
        // Find the call that this return corresponds to and update its duration.  During startuo,
        // we'll observe returns with an empty callstack - we ignore this.
        while let Some(call_id) = self.calls.pop() {
            // Given the call on top of the stack, verify that it was the one we expected.  If
            // it's wrong, we'll keep unwinding the stack until we find the correct call,
            // marking all unwound functions as ending at the same time.
            let matches = {
                // Track the return of functions but not the call.
                let caller_id = *self.func_table.get_name(
                    *self
                        .frame_table
                        .get_func(*self.stack_table.get_frame(*self.samples.get_stack(call_id))),
                );

                // Given our potential caller's id, we're most likely a match if the
                // StringIndex for the function we're returning matches the caller id.
                self.string_array
                    .get_full(func_name)
                    .map_or(false, |(id, _)| id == caller_id)
            };

            // Mark the parent as having returned.
            self.samples.weight[call_id.0] = Some(time - self.samples.get_time(call_id));

            if matches {
                // We've found the function we were supposed to be returning from.
                break;
            }

            log::trace!(
                "Returning from `{}` which was not the most recently called function",
                func_name
            );
        }

        time
    }

    /// Convert the given `FunctionTrace` into a meaningful format and attach it to the current
    /// thread.
    pub fn add_trace(&mut self, trace: FunctionTrace) {
        let event_time = match trace {
            FunctionTrace::Call {
                time,
                func_name,
                filename,
                linenumber,
            } => self.emit_call(&func_name, time, &filename, Some(linenumber)),
            FunctionTrace::NativeCall {
                time,
                func_name,
                module_name,
            } => self.emit_call(&func_name, time, &module_name, None),
            FunctionTrace::Return { time, func_name }
            | FunctionTrace::NativeReturn { time, func_name } => {
                self.emit_return(&func_name, self.offset_time(time))
            }
            FunctionTrace::Exception {
                time,
                exception_type,
                exception_value,
                filename,
                linenumber,
            } => {
                // TODO: We could report a large amount of information here such as
                // values of locals. For inspiration, see a django exception report.
                let unique_type = self.unique_string(&exception_type);
                let time = self.offset_time(time);

                self.markers.add(Marker {
                    data: serde_json::json!({
                        "type": "Exception",
                        "exeption": exception_type,
                        "value": exception_value,
                        "module": format!("{}:{}", filename, linenumber)
                    }),
                    name: unique_type,
                    startTime: time,
                    endTime: None,
                    phase: 0,
                    category: CategoryIndex::Exceptions,
                });

                time
            }
            FunctionTrace::Log {
                time,
                log_type,
                log_value,
            } => {
                let time = self.offset_time(time);
                let unique_type = self.unique_string(&log_type);

                self.markers.add(Marker {
                    data: serde_json::json!({
                        "type":   "Log",
                        "origin": log_type,
                        "value":   log_value,
                    }),
                    name: unique_type,
                    startTime: time,
                    endTime: None,
                    phase: 0,
                    category: CategoryIndex::Logs,
                });

                time
            }
            FunctionTrace::Import { time, module_name } => {
                let time = self.offset_time(time);
                let unique_type = self.unique_string("Import");

                self.markers.add(Marker {
                    data: serde_json::json!({
                        "type":   "Import",
                        "module":   module_name,
                    }),
                    name: unique_type,
                    startTime: time,
                    endTime: None,
                    phase: 0,
                    category: CategoryIndex::Imports,
                });

                time
            }
            FunctionTrace::Allocation { time, details } => {
                let time = self.offset_time(time);
                let stack = self.calls.last().map(|&x| *self.samples.get_stack(x));

                match details {
                    AllocationDetails::Alloc { bytes, addr } => {
                        self.allocations.insert(addr, bytes);

                        self.native_allocations.add(NativeAllocation {
                            time,
                            weight: bytes as isize,
                            stack,
                            memory_address: addr,
                            thread_id: self.tid,
                        });
                    }
                    AllocationDetails::Realloc {
                        bytes,
                        old_addr,
                        new_addr,
                    } => {
                        // Remove the old allocation first.
                        // If we don't find the allocation to free, it's either a bug or the
                        // allocation occurred before we started tracing.  Assume the latter.
                        let old_bytes = self.allocations.remove(&old_addr).unwrap_or(0);
                        self.native_allocations.add(NativeAllocation {
                            time,
                            weight: -(old_bytes as isize),
                            stack,
                            memory_address: old_addr,
                            thread_id: self.tid,
                        });

                        // Now add the new one.
                        self.allocations.insert(new_addr, bytes);
                        self.native_allocations.add(NativeAllocation {
                            time,
                            weight: bytes as isize,
                            stack,
                            memory_address: new_addr,
                            thread_id: self.tid,
                        });
                    }
                    AllocationDetails::Free { old_addr } => {
                        // If we don't find the allocation to free, it's either a bug or the
                        // allocation occurred before we started tracing.  Assume the latter.
                        let bytes = self.allocations.remove(&old_addr).unwrap_or(0);

                        self.native_allocations.add(NativeAllocation {
                            time,
                            weight: -(bytes as isize),
                            stack,
                            memory_address: old_addr,
                            thread_id: self.tid,
                        });
                    }
                }

                time
            }
            FunctionTrace::RegisterThread(_) => unreachable!(),
        };

        // Track what the last event that occurred on our thread was.
        self.process_shutdown_time = self.process_shutdown_time.max(event_time);
    }
}