cmtc 0.1.2

use cmtc::{run_firtool, to_fir_pipeline, Visitor};
use cmtir::utils::{indent, print_to};
use indexmap::IndexMap;
use std::{collections::VecDeque, path::Path};
use test_utils::pre_sim_pipeline;

/// Backend for Verilator
///
/// SV code of synthesized modules is generated by firtool,
/// so this backend only needs to generate C++ code for tb modules.
use super::*;

pub mod instance;
use instance::*;

mod rule_run;
pub use rule_run::*;

mod expr;
use expr::*;

mod cpp;
use cpp::*;

mod cmake;

/// Verilator Backend
pub struct SimVerilator {
  circuit: ir::Circuit,
  sim_data: interface::SimData,
}

pub struct Instance {
  pub full_path: String,
  pub module_name: String,
}

pub struct VerilatorRun {
  pub name: String,
  pub instances: Vec<Instance>,
  pub schedule_runs: Vec<ScheduleRun>,
  pub name_id_map: IndexMap<String, usize>,
  pub invoke_relations: Vec<Vec<bool>>,
  pub is_always: Vec<bool>,
  pub circuit: ir::Circuit,
}

impl VerilatorRun {
  pub fn build(&self) -> anyhow::Result<VerilatorCpp> {
    let instance_impls: Vec<InstanceImpl> = self
      .instances
      .iter()
      .map(|i| i.to_impl(&self.circuit))
      .collect();

    let includes = instance_impls.iter().map(|i| i.to_include()).collect();
    let instances_decl = instance_impls.iter().map(|i| i.to_decl()).collect();

    let instances_construct =
      instance_impls.iter().map(|i| i.to_construct()).collect();

    let instances_trace_on =
      instance_impls.iter().map(|i| i.to_trace_on()).collect();

    let instances_delete =
      instance_impls.iter().map(|i| i.to_delete()).collect();

    let instances_tick = tick_in_cpp(&instance_impls);

    let instances_reset = reset_in_cpp(&instance_impls);

    let cpp_rule_runs: Vec<CppRuleRun> = self
      .schedule_runs
      .iter()
      .map(|schedule| schedule.rules.clone())
      .flatten()
      .collect();

    let rule_guards = cpp_rule_runs
      .iter()
      .map(|rule| rule.to_guard_cpp(&self.name_id_map))
      .chain(instance_impls.iter().map(|i| i.to_guard_cpp()))
      .collect();

    let rule_runs = cpp_rule_runs
      .iter()
      .map(|rule| rule.to_rule_cpp(&self.name_id_map))
      .chain(instance_impls.iter().map(|i| i.to_rule_cpp()))
      .collect();

    let dispatch = dispatch_in_cpp(&self.name_id_map);

    let schedule_runs =
      schedule_runs_in_cpp(&self.schedule_runs, &self.name_id_map);

    let all_relations = all_relations_in_cpp(&self.invoke_relations);
    let is_always = is_always_in_cpp(&self.is_always);
    let has_run = format!(
      "{{{}}}",
      self
        .is_always
        .iter()
        .map(|_| "false".to_string())
        .collect::<Vec<_>>()
        .join(", ")
    );
    Ok(VerilatorCpp {
      name: self.name.clone(),
      includes,
      instances_decl,
      instances_construct,
      instances_trace_on,
      instances_delete,
      instances_tick,
      instances_reset,
      rule_guards,
      rule_runs,
      dispatch,
      schedule_runs,
      all_relations,
      is_always,
      has_run,
    })
  }
}

fn dispatch_in_cpp(name_id_map: &IndexMap<String, usize>) -> String {
  let mut swithes = String::new();

  for (name, id) in name_id_map.iter() {
    swithes.push_str(&format!("case {}:\n", id));
    swithes.push_str(&format!("  r_{}();\n", name));
    swithes.push_str("  break;\n");
  }

  swithes.push_str("  default:\n");
  swithes.push_str("    break;\n");

  format!(
    "void dispatch(int id) {{\n  switch (id) {{\n{}\n    }}\n  }}\n",
    indent(swithes, 4)
  )
}

fn schedule_runs_in_cpp(
  schedule_runs: &Vec<ScheduleRun>,
  name_id_map: &IndexMap<String, usize>,
) -> String {
  let mut s = String::new();
  s.push_str("{");
  for schedule in schedule_runs.iter() {
    s.push_str(&format!(
      "{{{}}}",
      schedule
        .rules
        .iter()
        .map(|r| name_id_map[&r.full_path].to_string())
        .collect::<Vec<_>>()
        .join(", ")
    ));
    s.push_str(",");
  }
  s.push_str("}");
  s
}

fn all_relations_in_cpp(invoke_relations: &Vec<Vec<bool>>) -> String {
  let mut s = String::new();
  s.push_str("{");
  for row in invoke_relations.iter() {
    s.push_str(&format!(
      "{{{}}}",
      row
        .iter()
        .map(|b| b.to_string())
        .collect::<Vec<_>>()
        .join(", ")
    ));
    s.push_str(",");
  }
  s.push_str("}");
  s
}

fn is_always_in_cpp(is_always: &Vec<bool>) -> String {
  let mut s = String::new();
  s.push_str("{");
  s.push_str(&format!(
    "{}",
    is_always
      .iter()
      .map(|b| b.to_string())
      .collect::<Vec<_>>()
      .join(", ")
  ));
  s.push_str("}");
  s
}

trait ToCpp {
  fn to_cpp(&self) -> String;
}

impl ToCpp for ir::Cmp {
  fn to_cpp(&self) -> String {
    match self {
      ir::Cmp::Eq => "==".to_string(),
      ir::Cmp::Neq => "!=".to_string(),
      ir::Cmp::Lt => "<".to_string(),
      ir::Cmp::Leq => "<=".to_string(),
      ir::Cmp::Gt => ">".to_string(),
      ir::Cmp::Geq => ">=".to_string(),
    }
  }
}

pub struct ScheduleRun {
  pub rules: Vec<CppRuleRun>,
}

impl ScheduleRun {
  pub fn new() -> Self {
    Self { rules: vec![] }
  }

  pub fn append(&mut self, rule_run: CppRuleRun) {
    self.rules.push(rule_run);
  }

  pub fn dump(&self) -> String {
    self
      .rules
      .iter()
      .map(|r| r.dump())
      .collect::<Vec<_>>()
      .join("")
  }
}

fn concat_path(path: &str, name: &str) -> String {
  format!("{}{}{}", path, if path.is_empty() { "" } else { "_" }, name)
}

impl SimVerilator {
  pub fn new(circuit: ir::Circuit, sim_data: interface::SimData) -> Self {
    Self { circuit, sim_data }
  }

  fn iterate_tb<InitT, UpdateT, T>(
    &self,
    init: InitT,
    update: UpdateT,
  ) -> anyhow::Result<T>
  where
    InitT: FnOnce() -> T,
    UpdateT: Fn(&mut T, &str, &str, &ModuleRun) -> anyhow::Result<()>,
  {
    log::debug!("Iterating over tb modules");

    let mut result = init();
    let mut fifo = VecDeque::new();

    let top_module_name = self.circuit.top_module_name.clone();
    log::debug!("\tTop module name: {}", top_module_name);

    fifo.push_back((top_module_name, "".to_string()));

    while !fifo.is_empty() {
      let (cur_module_name, cur_path) = fifo.pop_front().unwrap();

      log::debug!(
        "\tCurrent module name: {}, path: {}",
        cur_module_name,
        cur_path
      );

      let cur_module_run: &ModuleRun = self
        .sim_data
        .module_run(&cur_module_name)
        .ok_or(anyhow::anyhow!("Module {} not found", cur_module_name))?;

      update(&mut result, &cur_module_name, &cur_path, cur_module_run)?;

      // iterate over instances and add those tb instances into fifo
      for InstDef { name, module, .. } in cur_module_run.instances() {
        if self
          .circuit
          .module(module)
          .ok_or(anyhow::anyhow!("Module {} not found", module))?
          .is_tb()
        {
          fifo.push_back((module.clone(), concat_path(&cur_path, name)));
        }
      }
    }

    Ok(result)
  }

  /// Extract schedule runs from tb module
  /// it starts from the top module, and then recursively extract schedule runs
  /// of its submodules, until meeting a module that is not a tb module
  pub fn tb_schedule_runs(&self) -> anyhow::Result<Vec<ScheduleRun>> {
    log::debug!("Flattening rules");

    let init = || vec![];

    let update = |schedule_runs: &mut Vec<ScheduleRun>,
                  cur_module_name: &str,
                  cur_path: &str,
                  cur_module_run: &ModuleRun| {
      // extract schedule
      let mut schedule_run = ScheduleRun::new();
      for inst_rule in cur_module_run.schedule() {
        schedule_run.append(CppRuleRun::create(
          cur_module_run
            .rule_run(&inst_rule.rule_name)
            .ok_or(anyhow::anyhow!("Rule {} not found", inst_rule.rule_name))?
            .clone(),
          cur_module_name.to_string(),
          cur_path.to_string(),
          inst_rule.rule_name.to_string(),
          &self.circuit,
        )?);
      }
      log::debug!("\t\tSchedule run: \n{}", schedule_run.dump());

      schedule_runs.push(schedule_run);
      Ok(())
    };

    self.iterate_tb(init, update)
  }

  pub fn sv_instances(&self) -> anyhow::Result<Vec<Instance>> {
    let init = || vec![];
    let update = |sv_instances: &mut Vec<Instance>,
                  _cur_module_name: &str,
                  cur_path: &str,
                  cur_module_run: &ModuleRun| {
      // extract non-tb instances
      for InstDef {
        name,
        module: module_name,
        ..
      } in cur_module_run.instances()
      {
        let module = self
          .circuit
          .module(module_name)
          .ok_or(anyhow::anyhow!("Module {} not found", module_name))?;
        if !module.is_tb() {
          let sv_instance = Instance {
            full_path: concat_path(&cur_path, name),
            module_name: module_name.clone(),
          };
          log::debug!("\t\tInstance: {}", sv_instance.full_path);
          sv_instances.push(sv_instance);
        }
      }
      Ok(())
    };
    self.iterate_tb(init, update)
  }

  pub fn build_verilator_run(&self) -> anyhow::Result<VerilatorRun> {
    let tb_schedule_runs = self.tb_schedule_runs()?;
    let sv_instances = self.sv_instances()?;
    let mut is_always = vec![];

    let mut name_id_map = IndexMap::new();
    for schedule in tb_schedule_runs.iter() {
      for rule in schedule.rules.iter() {
        name_id_map.insert(rule.full_path.clone(), name_id_map.len());
        is_always.push(!rule.need_invoke);
      }
    }

    let mut invoke_relations =
      vec![vec![false; name_id_map.len()]; name_id_map.len()];

    for schedule in tb_schedule_runs.iter() {
      for rule in schedule.rules.iter() {
        let rule_id = name_id_map[&rule.full_path];
        for invoke in rule.invokes.iter() {
          if let Some(invoke_id) = name_id_map.get(invoke) {
            invoke_relations[*invoke_id][rule_id] = true;
          }
        }
      }
    }

    Ok(VerilatorRun {
      name: self.circuit.top_module_name.clone(),
      instances: sv_instances,
      schedule_runs: tb_schedule_runs,
      name_id_map,
      invoke_relations,
      is_always,
      circuit: self.circuit.clone(),
    })
  }

  pub fn elaborate(&self, path: impl AsRef<Path>) -> anyhow::Result<()> {
    log::info!("Elaborating verilator run");
    let verilator_run = self.build_verilator_run()?;

    let verilator_tb = format!("{}", verilator_run.build()?.to_cpp());

    let verilator_tb_path = path.as_ref().join("tb.cpp");
    print_to(verilator_tb, Some(&verilator_tb_path));

    let verilator_cmake = verilator_run.to_cmake();
    let verilator_cmake_path = path.as_ref().join("CMakeLists.txt");
    print_to(verilator_cmake, Some(&verilator_cmake_path));

    Ok(())
  }
}

pub fn create_verilator_workspace(
  circuit: ir::Circuit,
  path: impl AsRef<Path>,
) -> anyhow::Result<()> {
  let path = path.as_ref();
  std::fs::create_dir_all(path)?;

  let design_file_dir = path.join("hw");
  std::fs::create_dir_all(design_file_dir.clone())?;

  let (_cmtir, fir) = to_fir_pipeline(circuit.clone())?;
  // println!("{}", cmtir);
  run_firtool(fir, design_file_dir)?;

  let mut circuit = pre_sim_pipeline(circuit)?;
  let mut sim_itfc = SimItfc::new();
  sim_itfc.apply_pass(&mut circuit)?;
  let sim_verilator = SimVerilator::new(circuit, sim_itfc.into_inner());
  sim_verilator.elaborate(path)?;

  Ok(())
}