extern crate hercules_cg;
extern crate hercules_ir;
extern crate postcard;
extern crate serde;
extern crate take_mut;

use std::cell::RefCell;
use std::collections::HashMap;
use std::env::temp_dir;
use std::fs::File;
use std::io::Write;
use std::iter::zip;
use std::process::{Command, Stdio};

use self::serde::Deserialize;

use self::hercules_cg::*;
use self::hercules_ir::*;

use crate::*;

/*
 * Passes that can be run on a module.
 */
#[derive(Debug, Clone, Deserialize)]
pub enum Pass {
    DCE,
    CCP,
    GVN,
    Forkify,
    PhiElim,
    ForkGuardElim,
    Predication,
    SROA,
    Verify,
    // Parameterized over whether analyses that aid visualization are necessary.
    // Useful to set to false if displaying a potentially broken module.
    Xdot(bool),
    SchedXdot,
    // Parameterized over output directory and module name.
    Codegen(String, String),
    // Parameterized over where to serialize module to.
    Serialize(String),
}

/*
 * Manages passes to be run on an IR module. Transparently handles analysis
 * requirements for optimizations.
 */
#[derive(Debug, Clone)]
pub struct PassManager {
    module: Module,

    // Passes to run.
    passes: Vec<Pass>,

    // Cached analysis results.
    pub def_uses: Option<Vec<ImmutableDefUseMap>>,
    pub reverse_postorders: Option<Vec<Vec<NodeID>>>,
    pub typing: Option<ModuleTyping>,
    pub control_subgraphs: Option<Vec<Subgraph>>,
    pub doms: Option<Vec<DomTree>>,
    pub postdoms: Option<Vec<DomTree>>,
    pub fork_join_maps: Option<Vec<HashMap<NodeID, NodeID>>>,
    pub fork_join_nests: Option<Vec<HashMap<NodeID, Vec<NodeID>>>>,
    pub loops: Option<Vec<LoopTree>>,
    pub antideps: Option<Vec<Vec<(NodeID, NodeID)>>>,
    pub bbs: Option<Vec<Vec<NodeID>>>,

    // Current plan.
    pub plans: Option<Vec<Plan>>,

    // Store the manifest of a compiled object.
    pub manifests: Option<HashMap<String, Manifest>>,
}

impl PassManager {
    pub fn new(module: Module) -> Self {
        PassManager {
            module,
            passes: vec![],
            def_uses: None,
            reverse_postorders: None,
            typing: None,
            control_subgraphs: None,
            doms: None,
            postdoms: None,
            fork_join_maps: None,
            fork_join_nests: None,
            loops: None,
            antideps: None,
            bbs: None,
            plans: None,
            manifests: None,
        }
    }

    pub fn add_pass(&mut self, pass: Pass) {
        self.passes.push(pass);
    }

    pub fn make_def_uses(&mut self) {
        if self.def_uses.is_none() {
            self.def_uses = Some(self.module.functions.iter().map(def_use).collect());
        }
    }

    pub fn make_reverse_postorders(&mut self) {
        if self.reverse_postorders.is_none() {
            self.make_def_uses();
            self.reverse_postorders = Some(
                self.def_uses
                    .as_ref()
                    .unwrap()
                    .iter()
                    .map(reverse_postorder)
                    .collect(),
            );
        }
    }

    pub fn make_typing(&mut self) {
        if self.typing.is_none() {
            self.make_reverse_postorders();
            self.typing = Some(
                typecheck(&mut self.module, self.reverse_postorders.as_ref().unwrap()).unwrap(),
            );
        }
    }

    pub fn make_control_subgraphs(&mut self) {
        if self.control_subgraphs.is_none() {
            self.make_def_uses();
            self.control_subgraphs = Some(
                zip(&self.module.functions, self.def_uses.as_ref().unwrap())
                    .map(|(function, def_use)| control_subgraph(function, def_use))
                    .collect(),
            );
        }
    }

    pub fn make_doms(&mut self) {
        if self.doms.is_none() {
            self.make_control_subgraphs();
            self.doms = Some(
                self.control_subgraphs
                    .as_ref()
                    .unwrap()
                    .iter()
                    .map(|subgraph| dominator(subgraph, NodeID::new(0)))
                    .collect(),
            );
        }
    }

    pub fn make_postdoms(&mut self) {
        if self.postdoms.is_none() {
            self.make_control_subgraphs();
            self.postdoms = Some(
                zip(
                    self.control_subgraphs.as_ref().unwrap().iter(),
                    self.module.functions.iter(),
                )
                .map(|(subgraph, function)| dominator(subgraph, NodeID::new(function.nodes.len())))
                .collect(),
            );
        }
    }

    pub fn make_fork_join_maps(&mut self) {
        if self.fork_join_maps.is_none() {
            self.make_control_subgraphs();
            self.fork_join_maps = Some(
                zip(
                    self.module.functions.iter(),
                    self.control_subgraphs.as_ref().unwrap().iter(),
                )
                .map(|(function, subgraph)| fork_join_map(function, subgraph))
                .collect(),
            );
        }
    }

    pub fn make_fork_join_nests(&mut self) {
        if self.fork_join_nests.is_none() {
            self.make_doms();
            self.make_fork_join_maps();
            self.fork_join_nests = Some(
                zip(
                    self.module.functions.iter(),
                    zip(
                        self.doms.as_ref().unwrap().iter(),
                        self.fork_join_maps.as_ref().unwrap().iter(),
                    ),
                )
                .map(|(function, (dom, fork_join_map))| {
                    compute_fork_join_nesting(function, dom, fork_join_map)
                })
                .collect(),
            );
        }
    }

    pub fn make_loops(&mut self) {
        if self.loops.is_none() {
            self.make_control_subgraphs();
            self.make_doms();
            self.make_fork_join_maps();
            let control_subgraphs = self.control_subgraphs.as_ref().unwrap().iter();
            let doms = self.doms.as_ref().unwrap().iter();
            let fork_join_maps = self.fork_join_maps.as_ref().unwrap().iter();
            self.loops = Some(
                zip(control_subgraphs, zip(doms, fork_join_maps))
                    .map(|(control_subgraph, (dom, fork_join_map))| {
                        loops(control_subgraph, NodeID::new(0), dom, fork_join_map)
                    })
                    .collect(),
            );
        }
    }

    pub fn make_antideps(&mut self) {
        if self.antideps.is_none() {
            self.make_def_uses();
            self.antideps = Some(
                zip(
                    self.def_uses.as_ref().unwrap().iter(),
                    self.module.functions.iter(),
                )
                .map(|(def_use, function)| antideps(function, def_use))
                .collect(),
            );
        }
    }

    pub fn make_bbs(&mut self) {
        if self.bbs.is_none() {
            self.make_def_uses();
            self.make_reverse_postorders();
            self.make_doms();
            self.make_antideps();
            self.make_loops();
            self.make_fork_join_maps();
            let def_uses = self.def_uses.as_ref().unwrap().iter();
            let reverse_postorders = self.reverse_postorders.as_ref().unwrap().iter();
            let doms = self.doms.as_ref().unwrap().iter();
            let antideps = self.antideps.as_ref().unwrap().iter();
            let loops = self.loops.as_ref().unwrap().iter();
            let fork_join_maps = self.fork_join_maps.as_ref().unwrap().iter();
            self.bbs = Some(
                zip(
                    self.module.functions.iter(),
                    zip(
                        def_uses,
                        zip(
                            reverse_postorders,
                            zip(doms, zip(antideps, zip(loops, fork_join_maps))),
                        ),
                    ),
                )
                .map(
                    |(
                        function,
                        (def_use, (reverse_postorder, (dom, (antideps, (loops, fork_join_map))))),
                    )| {
                        gcm(
                            function,
                            def_use,
                            reverse_postorder,
                            dom,
                            antideps,
                            loops,
                            fork_join_map,
                            &mut vec![None; function.nodes.len()],
                        )
                    },
                )
                .collect(),
            );
        }
    }

    pub fn set_plans(&mut self, plans: Vec<Plan>) {
        self.plans = Some(plans);
    }

    pub fn make_plans(&mut self) {
        if self.plans.is_none() {
            self.make_def_uses();
            self.make_reverse_postorders();
            self.make_fork_join_maps();
            self.make_fork_join_nests();
            self.make_bbs();
            let def_uses = self.def_uses.as_ref().unwrap().iter();
            let reverse_postorders = self.reverse_postorders.as_ref().unwrap().iter();
            let fork_join_maps = self.fork_join_maps.as_ref().unwrap().iter();
            let fork_join_nests = self.fork_join_nests.as_ref().unwrap().iter();
            let bbs = self.bbs.as_ref().unwrap().iter();
            self.plans = Some(
                zip(
                    self.module.functions.iter(),
                    zip(
                        def_uses,
                        zip(
                            reverse_postorders,
                            zip(fork_join_maps, zip(fork_join_nests, bbs)),
                        ),
                    ),
                )
                .map(
                    |(
                        function,
                        (def_use, (reverse_postorder, (fork_join_map, (fork_join_nest, bb)))),
                    )| {
                        default_plan(
                            function,
                            &self.module.dynamic_constants,
                            def_use,
                            reverse_postorder,
                            fork_join_map,
                            fork_join_nest,
                            bb,
                        )
                    },
                )
                .collect(),
            );
        }
    }

    pub fn run_passes(&mut self) {
        for pass in self.passes.clone().iter() {
            match pass {
                Pass::DCE => {
                    self.make_def_uses();
                    let def_uses = self.def_uses.as_ref().unwrap();
                    for idx in 0..self.module.functions.len() {
                        let constants_ref =
                            RefCell::new(std::mem::take(&mut self.module.constants));
                        let dynamic_constants_ref =
                            RefCell::new(std::mem::take(&mut self.module.dynamic_constants));
                        let types_ref = RefCell::new(std::mem::take(&mut self.module.types));
                        let mut editor = FunctionEditor::new(
                            &mut self.module.functions[idx],
                            &constants_ref,
                            &dynamic_constants_ref,
                            &types_ref,
                            &def_uses[idx],
                        );
                        dce(&mut editor);

                        self.module.constants = constants_ref.take();
                        self.module.dynamic_constants = dynamic_constants_ref.take();
                        self.module.types = types_ref.take();

                        let edits = &editor.edits();
                        if let Some(plans) = self.plans.as_mut() {
                            repair_plan(&mut plans[idx], &self.module.functions[idx], edits);
                        }
                        let grave_mapping = self.module.functions[idx].delete_gravestones();
                        if let Some(plans) = self.plans.as_mut() {
                            plans[idx].fix_gravestones(&grave_mapping);
                        }
                    }
                    self.clear_analyses();
                }
                Pass::CCP => {
                    self.make_def_uses();
                    self.make_reverse_postorders();
                    let def_uses = self.def_uses.as_ref().unwrap();
                    let reverse_postorders = self.reverse_postorders.as_ref().unwrap();
                    for idx in 0..self.module.functions.len() {
                        ccp(
                            &mut self.module.functions[idx],
                            &mut self.module.constants,
                            &def_uses[idx],
                            &reverse_postorders[idx],
                        );
                    }
                    self.legacy_repair_plan();
                    self.clear_analyses();
                }
                Pass::GVN => {
                    self.make_def_uses();
                    let def_uses = self.def_uses.as_ref().unwrap();
                    for idx in 0..self.module.functions.len() {
                        let constants_ref =
                            RefCell::new(std::mem::take(&mut self.module.constants));
                        let dynamic_constants_ref =
                            RefCell::new(std::mem::take(&mut self.module.dynamic_constants));
                        let types_ref = RefCell::new(std::mem::take(&mut self.module.types));
                        let mut editor = FunctionEditor::new(
                            &mut self.module.functions[idx],
                            &constants_ref,
                            &dynamic_constants_ref,
                            &types_ref,
                            &def_uses[idx],
                        );
                        gvn(&mut editor);

                        self.module.constants = constants_ref.take();
                        self.module.dynamic_constants = dynamic_constants_ref.take();
                        self.module.types = types_ref.take();

                        let edits = &editor.edits();
                        if let Some(plans) = self.plans.as_mut() {
                            repair_plan(&mut plans[idx], &self.module.functions[idx], edits);
                        }
                        let grave_mapping = self.module.functions[idx].delete_gravestones();
                        if let Some(plans) = self.plans.as_mut() {
                            plans[idx].fix_gravestones(&grave_mapping);
                        }
                    }
                    self.clear_analyses();
                }
                Pass::Forkify => {
                    self.make_def_uses();
                    self.make_loops();
                    let def_uses = self.def_uses.as_ref().unwrap();
                    let loops = self.loops.as_ref().unwrap();
                    for idx in 0..self.module.functions.len() {
                        forkify(
                            &mut self.module.functions[idx],
                            &self.module.constants,
                            &mut self.module.dynamic_constants,
                            &def_uses[idx],
                            &loops[idx],
                        )
                    }
                    self.legacy_repair_plan();
                    self.clear_analyses();
                }
                Pass::PhiElim => {
                    for function in self.module.functions.iter_mut() {
                        phi_elim(function);
                    }
                    self.legacy_repair_plan();
                    self.clear_analyses();
                }
                Pass::ForkGuardElim => {
                    self.make_def_uses();
                    self.make_fork_join_maps();
                    let def_uses = self.def_uses.as_ref().unwrap();
                    let fork_join_maps = self.fork_join_maps.as_ref().unwrap();
                    for idx in 0..self.module.functions.len() {
                        fork_guard_elim(
                            &mut self.module.functions[idx],
                            &self.module.constants,
                            &fork_join_maps[idx],
                            &def_uses[idx],
                        )
                    }
                    self.legacy_repair_plan();
                    self.clear_analyses();
                }
                Pass::Predication => {
                    self.make_def_uses();
                    self.make_reverse_postorders();
                    self.make_doms();
                    self.make_fork_join_maps();
                    self.make_plans();
                    let def_uses = self.def_uses.as_ref().unwrap();
                    let reverse_postorders = self.reverse_postorders.as_ref().unwrap();
                    let doms = self.doms.as_ref().unwrap();
                    let fork_join_maps = self.fork_join_maps.as_ref().unwrap();
                    let plans = self.plans.as_ref().unwrap();
                    for idx in 0..self.module.functions.len() {
                        predication(
                            &mut self.module.functions[idx],
                            &def_uses[idx],
                            &reverse_postorders[idx],
                            &doms[idx],
                            &fork_join_maps[idx],
                            &plans[idx].schedules,
                        )
                    }
                    self.legacy_repair_plan();
                    self.clear_analyses();
                }
                Pass::SROA => {
                    self.make_def_uses();
                    self.make_reverse_postorders();
                    self.make_typing();
                    let def_uses = self.def_uses.as_ref().unwrap();
                    let reverse_postorders = self.reverse_postorders.as_ref().unwrap();
                    let typing = self.typing.as_ref().unwrap();
                    for idx in 0..self.module.functions.len() {
                        sroa(
                            &mut self.module.functions[idx],
                            &def_uses[idx],
                            &reverse_postorders[idx],
                            &typing[idx],
                            &self.module.types,
                            &mut self.module.constants,
                        );
                    }
                    self.legacy_repair_plan();
                    self.clear_analyses();
                }
                Pass::Verify => {
                    let (
                        def_uses,
                        reverse_postorders,
                        typing,
                        subgraphs,
                        doms,
                        postdoms,
                        fork_join_maps,
                    ) = verify(&mut self.module)
                        .expect("PANIC: Failed to verify Hercules IR module.");

                    // Verification produces a bunch of analysis results that
                    // may be useful for later passes.
                    self.def_uses = Some(def_uses);
                    self.reverse_postorders = Some(reverse_postorders);
                    self.typing = Some(typing);
                    self.control_subgraphs = Some(subgraphs);
                    self.doms = Some(doms);
                    self.postdoms = Some(postdoms);
                    self.fork_join_maps = Some(fork_join_maps);

                    // Verify the plan, if it exists.
                    if let Some(plans) = &self.plans {
                        for idx in 0..self.module.functions.len() {
                            plans[idx].verify_partitioning(
                                &self.module.functions[idx],
                                &self.def_uses.as_ref().unwrap()[idx],
                                &self.fork_join_maps.as_ref().unwrap()[idx],
                            );
                        }
                    }
                }
                Pass::Xdot(force_analyses) => {
                    self.make_reverse_postorders();
                    if *force_analyses {
                        self.make_doms();
                        self.make_fork_join_maps();
                        self.make_bbs();
                        self.make_plans();
                    }
                    xdot_module(
                        &self.module,
                        self.reverse_postorders.as_ref().unwrap(),
                        self.doms.as_ref(),
                        self.fork_join_maps.as_ref(),
                        self.bbs.as_ref(),
                        self.plans.as_ref(),
                    );
                }
                Pass::SchedXdot => {
                    self.make_def_uses();
                    self.make_typing();
                    self.make_control_subgraphs();
                    self.make_fork_join_maps();
                    self.make_fork_join_nests();
                    self.make_antideps();
                    self.make_bbs();
                    self.make_plans();

                    let smodule = sched_compile(
                        &self.module,
                        self.def_uses.as_ref().unwrap(),
                        self.typing.as_ref().unwrap(),
                        self.control_subgraphs.as_ref().unwrap(),
                        self.fork_join_maps.as_ref().unwrap(),
                        self.fork_join_nests.as_ref().unwrap(),
                        self.antideps.as_ref().unwrap(),
                        self.bbs.as_ref().unwrap(),
                        self.plans.as_ref().unwrap(),
                    );

                    xdot_sched_module(&smodule);
                }
                Pass::Codegen(output_dir, module_name) => {
                    self.make_def_uses();
                    self.make_typing();
                    self.make_control_subgraphs();
                    self.make_fork_join_maps();
                    self.make_fork_join_nests();
                    self.make_antideps();
                    self.make_bbs();
                    self.make_plans();

                    let smodule = sched_compile(
                        &self.module,
                        self.def_uses.as_ref().unwrap(),
                        self.typing.as_ref().unwrap(),
                        self.control_subgraphs.as_ref().unwrap(),
                        self.fork_join_maps.as_ref().unwrap(),
                        self.fork_join_nests.as_ref().unwrap(),
                        self.antideps.as_ref().unwrap(),
                        self.bbs.as_ref().unwrap(),
                        self.plans.as_ref().unwrap(),
                    );

                    let mut llvm_ir = String::new();
                    for manifest in smodule.manifests.values() {
                        for partition_manifest in manifest.partitions.iter() {
                            let function = &smodule.functions[&partition_manifest.name];
                            match partition_manifest.device {
                                DeviceManifest::CPU { parallel_launch: _ } => {
                                    cpu_compile(function, partition_manifest, &mut llvm_ir).unwrap()
                                }
                                _ => todo!(),
                            }
                        }
                    }
                    println!("{}", llvm_ir);

                    // Write the LLVM IR into a temporary file.
                    let mut tmp_path = temp_dir();
                    tmp_path.push(format!("{}.ll", module_name));
                    let mut file = File::create(&tmp_path)
                        .expect("PANIC: Unable to open output LLVM IR file.");
                    file.write_all(llvm_ir.as_bytes())
                        .expect("PANIC: Unable to write output LLVM IR file contents.");
                    println!("{}", tmp_path.display());

                    // Compile LLVM IR into an ELF object file.
                    let output_archive = format!("{}/lib{}.a", output_dir, module_name);
                    let mut clang_process = Command::new("clang")
                        .arg(&tmp_path)
                        .arg("--emit-static-lib")
                        .arg("-O3")
                        .arg("-march=native")
                        .arg("-o")
                        .arg(&output_archive)
                        .stdin(Stdio::piped())
                        .stdout(Stdio::piped())
                        .spawn()
                        .unwrap();
                    assert!(clang_process.wait().unwrap().success());
                    println!("{}", output_archive);

                    // Package manifest into a file.
                    let hman_contents: Vec<u8> = postcard::to_allocvec(&smodule.manifests).unwrap();
                    let mut file = File::create(format!("{}/{}.hman", output_dir, module_name))
                        .expect("PANIC: Unable to open output manifest file.");
                    file.write_all(&hman_contents)
                        .expect("PANIC: Unable to write output manifest file contents.");
                    self.manifests = Some(smodule.manifests);
                    println!("{:?}", self.manifests);
                }
                Pass::Serialize(output_file) => {
                    let module_contents: Vec<u8> = postcard::to_allocvec(&self.module).unwrap();
                    let mut file = File::create(&output_file)
                        .expect("PANIC: Unable to open output module file.");
                    file.write_all(&module_contents)
                        .expect("PANIC: Unable to write output module file contents.");
                }
            }
        }
    }

    fn legacy_repair_plan(&mut self) {
        // Cleanup the module after passes. Delete gravestone nodes. Repair the
        // plans.
        for idx in 0..self.module.functions.len() {
            let grave_mapping = self.module.functions[idx].delete_gravestones();
            let plans = &mut self.plans;
            let functions = &self.module.functions;
            if let Some(plans) = plans.as_mut() {
                take_mut::take(&mut plans[idx], |plan| {
                    plan.repair(&functions[idx], &grave_mapping)
                });
            }
        }
    }

    fn clear_analyses(&mut self) {
        self.def_uses = None;
        self.reverse_postorders = None;
        self.typing = None;
        self.control_subgraphs = None;
        self.doms = None;
        self.postdoms = None;
        self.fork_join_maps = None;
        self.fork_join_nests = None;
        self.loops = None;
        self.antideps = None;
        self.bbs = None;

        // Don't clear the plan - this is repaired, not reconstructed.
    }

    pub fn get_module(self) -> Module {
        self.module
    }

    pub fn get_manifests(self) -> HashMap<String, Manifest> {
        self.manifests.unwrap()
    }
}