use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use std::fs::File;
use std::io::Write;
use std::time::{Duration, SystemTime};
use serde::Serialize;
use serde_repr::Serialize_repr;
use crate::function_trace::*;
use indexmap::set::IndexSet;

// 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 for doing so.
pub struct FirefoxProfileThreadId {
    tid: usize,
    start_duration: Duration,
}

#[derive(Serialize, Debug)]
pub struct FirefoxProfile {
    meta: Metadata,
    counters: Vec<Unused>,
    threads: Vec<FirefoxThread>,

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

///////////////////////////////////////////////////////////////////////////////
// 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)]
#[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_name: ident, $($element: ident: $ty: ty),*) => {
        paste::item! {
            #[derive(Serialize, Debug, Copy, Clone)]
            struct [<$struct_name Index>] (usize);

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

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

            impl [<$struct_name Table>] {
                fn new() -> Self {
                    Self {
                        $($element: Vec::new()),*,
                        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)
                }

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


struct_table!(Sample,
    event_delay: u32,
    stack: StackIndex,
    time: FirefoxTime,
    duration: Option<FirefoxTime>
);
struct_table!(Stack,
    frame: FrameIndex,
    category: CategoryIndex,
    subcategory: u32,
    prefix: Option<StackIndex>
);
struct_table!(Frame,
    address: Option<u32>, // TODO: What's the proper type here?
    category: CategoryIndex,
    subcategory: u32,
    func: FunctionIndex,
    innerWindowID: Option<Unused>,
    implementation: Option<Unused>,
    line: Option<u32>,
    column: Option<u32>,
    optimizations: Option<Unused>
);
struct_table!(Function,
    address: Option<Unused>,
    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!(NativeAllocation,
    time: FirefoxTime,
    duration: isize,  // This is actually bytes
    stack: Option<StackIndex>,
    memory_address: usize,
    thread_id: usize
);
struct_table!(Marker,
    data: serde_json::Value,
    name: StringIndex,
    time: FirefoxTime,
    category: CategoryIndex
);
struct_table!(Resource,
    lib: Unused,
    name: Unused,
    host: Unused,
    r#type: Unused
);

///////////////////////////////////////////////////////////////////////////////
// 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,
    process_type: u32,  // Should be 0
    categories: &'static [Category; CategoryIndex::Length as usize],
    stackwalk: u32, // Should be 0
    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,

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


#[derive(Serialize, Debug)]
#[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,
    process_name: String,
    is_js_tracer: bool,
    pid: usize,
    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)
    ///////////////////////////////////////////////////////////////////////////
    libs: Vec<Unused>,
    resource_table: ResourceTable,

    ///////////////////////////////////////////////////////////////////////////
    // 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<String, 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,

    #[serde(skip_serializing)]
    // 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.
    allocations: HashMap<usize, usize>,
}

impl FirefoxProfile {
    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 3 Feb 2020
                version: 19,
                preprocessed_profile_version: 28,
                symbolicated: true,
                physical_cpus: num_cpus::get_physical(),
                logical_cpus: num_cpus::get(),

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

            thread_count: 0,
        }
    }

    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);
    }

    pub fn export(&mut self, mut output_dir: PathBuf) -> std::io::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"));

        // 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 = HashSet::new();

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

            registered_pids.insert(thread.pid);
        }

        // Output the current profile to the requested directory.
        let output_file = {
            let mut dir = output_dir.clone();
            let now = chrono::Local::now();

            dir.push(format!("functiontrace.{}.json", now.format("%F-%T.%3f")));
            dir
        };

        let json = serde_json::to_string(&self)?;
        let bytes = json.as_bytes();
        File::create(&output_file)?.write_all(bytes)?;

        // Link `functiontrace.latest.json` to this file.  We need to remove any existing symlink
        // before doing so.
        output_dir.push("functiontrace.latest.json");
        let _ = std::fs::remove_file(&output_dir);
        std::os::unix::fs::symlink(&output_file, &output_dir)?;

        println!("\n[FunctionTrace] Wrote profile data to {} ({})\n", output_file.display(), bytesize::to_string(bytes.len() as u64, false));
        Ok(())
    }
}

impl FirefoxThread {
    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,
            unregister_time: None,
            paused_ranges: Vec::new(),
            name: "<filled in at export>",
            process_name: registration.program_name,
            is_js_tracer: true,
            pid: registration.pid,
            tid: thread.tid,
            samples: SampleTable::new(),
            markers: MarkerTable::new(),
            stack_table: StackTable::new(),
            frame_table: FrameTable::new(),
            func_table: FunctionTable::new(),
            string_array: IndexSet::new(),
            native_allocations: NativeAllocationTable::new(),

            libs: Vec::new(),
            resource_table: ResourceTable::new(),

            calls: Vec::new(),
            functions: HashMap::new(),
            start_duration: thread.start_duration,
            allocations: HashMap::new(),
        }
    }

    // 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, linenumber: Option<u32>) -> FirefoxTime {
        // Only native functions don't have known line numbers.
        let native = linenumber.is_none();

        // Create a UID for this function to avoid creating multiple for it.
        let func_uid = if native {
            format!("{}:{}:native", name, file_or_module)
        } else {
            format!("{} ({}:{})", name, file_or_module, linenumber.expect("Checked above"))
        };

        // 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 {
                    address: None,
                    isJS: false,
                    name: funcname_id,
                    resource: -1,
                    relevantForJS: false,
                    file_name: filename_id,
                    line_number: linenumber,
                    column_number: None
                });
                self.frame_table.add(Frame {
                    address: None,
                    category: if native { CategoryIndex::Native } else { CategoryIndex::Python },
                    subcategory: 0,
                    func: function_id,
                    innerWindowID: None,
                    implementation: None,
                    line: linenumber, // TODO: Is this the right line number?
                    column: None,
                    optimizations: None
                });

                self.functions.insert(func_uid, function_id);
                function_id
            }
        };

        // Emit a sample and a stacktrace.
        let time = self.offset_time(time);
        // TODO: We don't need a new stack for each sample - many are lkkely duplicated.
        let stack_id = self.stack_table.add(Stack {
            frame: FrameIndex(function_id.0),
            category: if native { CategoryIndex::Native } else { CategoryIndex::Python },
            subcategory: 0,
            prefix: self.calls.last().map(|&x| StackIndex(x.0))
        });
        let samples_id = self.samples.add(Sample {
            event_delay: 0,
            stack: stack_id,
            time,
            duration: 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.name[self.frame_table.func[self.stack_table.frame[self.samples.stack[call_id.0].0].0].0];

                // 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.duration[call_id.0] = Some(time - self.samples.time[call_id.0]);

            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
    }

    // Returns true if we should keep tracing.
    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} => self.emit_return(&func_name, self.offset_time(time)),
            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":   "Log",
                        "name":   exception_value,
                        "module": format!("{}:{}", filename, linenumber)
                    }),
                    name: unique_type,
                    time,
                    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":   "Text",
                        "name":   log_value,
                    }),
                    name: unique_type,
                    time,
                    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":   "Log",
                        "name":   module_name,
                        "module": "import",
                    }),
                    name: unique_type,
                    time,
                    category: CategoryIndex::Imports,
                });

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

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

                        self.native_allocations.add(NativeAllocation {
                            time,
                            duration: 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,
                            duration: -(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,
                            duration: 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,
                            duration: -(bytes as isize),
                            stack,
                            memory_address: old_addr,
                            thread_id: self.tid
                        });
                    },
                }

                time
            },
            msg => panic!("Invalid message variant: {:?}", msg)
        };

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