mod entry_book;

use std::collections::HashMap;
use std::sync::Arc;
use std::sync::LazyLock;
use std::sync::RwLock;

use cirru_parser::Cirru;

use crate::calcit::{self, Calcit, CalcitTypeAnnotation, DYNAMIC_TYPE};
use crate::data::cirru::code_to_calcit;
use crate::snapshot;
use crate::snapshot::Snapshot;
use crate::util::string::extract_pkg_from_ns;

pub use entry_book::EntryBook;

pub type ProgramEvaledData = EntryBook<EntryBook<Calcit>>;

/// definition entry with code and documentation
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProgramDefEntry {
  pub code: Calcit,
  pub schema: Arc<CalcitTypeAnnotation>,
  pub doc: Arc<str>,
  pub examples: Vec<Cirru>,
}

/// information extracted from snapshot
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProgramFileData {
  pub import_map: HashMap<Arc<str>, Arc<ImportRule>>,
  pub defs: HashMap<Arc<str>, ProgramDefEntry>,
}

type ImportMapPair = (Arc<str>, Arc<ImportRule>);

/// defRule: ns def
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ImportRule {
  /// ns imported via `:as`
  NsAs(Arc<str>),
  /// (ns, def) imported via `:refer`
  NsReferDef(Arc<str>, Arc<str>),
  /// ns imported via `:default`, js only
  NsDefault(Arc<str>),
}

pub type ProgramCodeData = HashMap<Arc<str>, ProgramFileData>;

/// data of program running
static PROGRAM_EVALED_DATA_STATE: LazyLock<RwLock<ProgramEvaledData>> = LazyLock::new(|| RwLock::new(EntryBook::default()));
/// raw code information before program running
pub static PROGRAM_CODE_DATA: LazyLock<RwLock<ProgramCodeData>> = LazyLock::new(|| RwLock::new(HashMap::new()));

fn extract_import_rule(nodes: &Cirru) -> Result<Vec<ImportMapPair>, String> {
  match nodes {
    Cirru::Leaf(_) => Err(String::from("Expected import rule in expr")),
    Cirru::List(rule_nodes) => {
      let mut xs = rule_nodes.to_owned();
      match xs.first() {
        // strip leading `[]` symbols
        Some(Cirru::Leaf(s)) if &**s == "[]" => xs = xs[1..4].to_vec(),
        // allow using comment
        Some(Cirru::Leaf(s)) if &**s == ";" => return Ok(vec![]),
        _ => (),
      }
      match (&xs[0], &xs[1], &xs[2]) {
        (Cirru::Leaf(ns), x, Cirru::Leaf(alias)) if x.eq_leaf(":as") => Ok(vec![(
          (*alias.to_owned()).into(),
          Arc::new(ImportRule::NsAs((*ns.to_owned()).into())),
        )]),
        (Cirru::Leaf(ns), x, Cirru::Leaf(alias)) if x.eq_leaf(":default") => Ok(vec![(
          (*alias.to_owned()).into(),
          Arc::new(ImportRule::NsDefault((*ns.to_owned()).into())),
        )]),
        (Cirru::Leaf(ns), x, Cirru::List(ys)) if x.eq_leaf(":refer") => {
          let mut rules: Vec<(Arc<str>, Arc<ImportRule>)> = Vec::with_capacity(ys.len());
          for y in ys {
            match y {
              Cirru::Leaf(s) if &**s == "[]" => (), // `[]` symbol are ignored
              Cirru::Leaf(s) => rules.push((
                (*s.to_owned()).into(),
                Arc::new(ImportRule::NsReferDef((*ns.to_owned()).into(), (*s.to_owned()).into())),
              )),
              Cirru::List(_defs) => return Err(format!("invalid refer values, {y}")),
            }
          }
          Ok(rules)
        }
        (_, x, _) if x.eq_leaf(":as") => Err(format!("invalid import rule: {nodes}")),
        (_, x, _) if x.eq_leaf(":default") => Err(format!("invalid default rule: {nodes}")),
        (_, x, _) if x.eq_leaf(":refer") => Err(format!("invalid import rule: {nodes}")),
        _ if xs.len() != 3 => Err(format!("expected import rule has length 3: {nodes}")),
        _ => Err(String::from("unknown rule")),
      }
    }
  }
}

fn extract_import_map(nodes: &Cirru, ns_name: &str) -> Result<HashMap<Arc<str>, Arc<ImportRule>>, String> {
  match nodes {
    Cirru::Leaf(_) => unreachable!("Expected expr for ns"),
    Cirru::List(xs) => match (xs.first(), xs.get(1), xs.get(2)) {
      // Too many clones
      (Some(x), Some(Cirru::Leaf(_)), Some(Cirru::List(xs))) if x.eq_leaf("ns") => {
        if !xs.is_empty() && xs[0].eq_leaf(":require") {
          let mut ys: HashMap<Arc<str>, Arc<ImportRule>> = HashMap::with_capacity(xs.len());
          for x in xs.iter().skip(1) {
            let rules = extract_import_rule(x).map_err(|e| format!("in namespace '{ns_name}': {e}"))?;
            for (target, rule) in rules {
              ys.insert(target, rule);
            }
          }
          Ok(ys)
        } else {
          Ok(HashMap::new())
        }
      }
      _ if xs.len() < 3 => Ok(HashMap::new()),
      _ => {
        let preview = cirru_parser::format(&[nodes.clone()], true.into()).unwrap_or_else(|_| format!("{nodes:?}"));
        let preview_short = if preview.len() > 200 {
          format!("{}...", &preview[..200])
        } else {
          preview
        };
        Err(format!("invalid ns form in '{ns_name}':\n{preview_short}"))
      }
    },
  }
}

fn extract_file_data(file: &snapshot::FileInSnapShot, ns: Arc<str>) -> Result<ProgramFileData, String> {
  let import_map = extract_import_map(&file.ns.code, &ns)?;
  let mut defs: HashMap<Arc<str>, ProgramDefEntry> = HashMap::with_capacity(file.defs.len());
  for (def, entry) in &file.defs {
    let at_def = def.to_owned();
    let code = code_to_calcit(&entry.code, &ns, &at_def, vec![])?;
    let schema = entry.schema.clone();
    let doc = Arc::from(entry.doc.as_str());
    defs.insert(
      def.to_owned().into(),
      ProgramDefEntry {
        code,
        schema,
        doc,
        examples: entry.examples.clone(),
      },
    );
  }
  Ok(ProgramFileData { import_map, defs })
}

pub fn extract_program_data(s: &Snapshot) -> Result<ProgramCodeData, String> {
  // Register the program lookup functions in type_annotation so it can resolve
  // imported type definitions without a circular module dependency.
  calcit::register_program_lookups(lookup_evaled_def, lookup_def_code);

  let mut xs: ProgramCodeData = HashMap::with_capacity(s.files.len());

  for (ns, file) in &s.files {
    let file_info = extract_file_data(file, ns.to_owned().into())?;
    xs.insert(ns.to_owned().into(), file_info);
  }

  Ok(xs)
}

// lookup without cloning
pub fn has_def_code(ns: &str, def: &str) -> bool {
  let program_code = { PROGRAM_CODE_DATA.read().expect("read program code") };
  match &program_code.get(ns) {
    Some(v) => v.defs.contains_key(def),
    None => false,
  }
}

pub fn lookup_def_code(ns: &str, def: &str) -> Option<Calcit> {
  let program_code = { PROGRAM_CODE_DATA.read().expect("read program code") };
  let file = program_code.get(ns)?;
  let entry = file.defs.get(def)?;
  Some(entry.code.to_owned())
}

pub fn lookup_def_schema(ns: &str, def: &str) -> Arc<CalcitTypeAnnotation> {
  let program_code = { PROGRAM_CODE_DATA.read().expect("read program code") };
  let file = match program_code.get(ns) {
    Some(f) => f,
    None => return DYNAMIC_TYPE.clone(),
  };
  let entry = match file.defs.get(def) {
    Some(e) => e,
    None => return DYNAMIC_TYPE.clone(),
  };
  entry.schema.clone()
}

/// lookup documentation for a definition from program data
pub fn lookup_def_doc(ns: &str, def: &str) -> Option<String> {
  let program_code = { PROGRAM_CODE_DATA.read().expect("read program code") };
  let file = program_code.get(ns)?;
  let entry = file.defs.get(def)?;
  if entry.doc.is_empty() { None } else { Some(entry.doc.to_string()) }
}

/// lookup examples for a definition from program data
pub fn lookup_def_examples(ns: &str, def: &str) -> Option<Vec<Cirru>> {
  let program_code = { PROGRAM_CODE_DATA.read().expect("read program code") };
  let file = program_code.get(ns)?;
  let entry = file.defs.get(def)?;
  if entry.examples.is_empty() {
    None
  } else {
    Some(entry.examples.clone())
  }
}

pub fn lookup_def_target_in_import(ns: &str, def: &str) -> Option<Arc<str>> {
  let program = { PROGRAM_CODE_DATA.read().expect("read program code") };
  let file = program.get(ns)?;
  let import_rule = file.import_map.get(def)?;
  match &**import_rule {
    ImportRule::NsReferDef(ns, _def) => Some(ns.to_owned()),
    ImportRule::NsAs(_ns) => None,
    ImportRule::NsDefault(_ns) => None,
  }
}

pub fn lookup_ns_target_in_import(ns: &str, alias: &str) -> Option<Arc<str>> {
  let program = { PROGRAM_CODE_DATA.read().expect("read program code") };
  let file = program.get(ns)?;
  let import_rule = file.import_map.get(alias)?;
  match &**import_rule {
    ImportRule::NsReferDef(_ns, _def) => None,
    ImportRule::NsAs(ns) => Some(ns.to_owned()),
    ImportRule::NsDefault(_ns) => None,
  }
}

/// try if we can load coord from evaled data
pub fn tip_coord(ns: &str, def: &str) -> Option<(u16, u16)> {
  let program = { PROGRAM_EVALED_DATA_STATE.read().expect("read program code") };
  match program.lookup(ns) {
    Some(file) => file.0.lookup(def).map(|(_, v)| (file.1, v)),
    None => None,
  }
}

// imported via :default
pub fn lookup_default_target_in_import(at_ns: &str, alias: &str) -> Option<Arc<str>> {
  let program = { PROGRAM_CODE_DATA.read().expect("read program code") };
  let file = program.get(at_ns)?;
  let import_rule = file.import_map.get(alias)?;
  match &**import_rule {
    ImportRule::NsReferDef(_ns, _def) => None,
    ImportRule::NsAs(_ns) => None,
    ImportRule::NsDefault(ns) => Some(ns.to_owned()),
  }
}

/// lookup and return value
pub fn lookup_evaled_def(ns: &str, def: &str) -> Option<Calcit> {
  let s2 = PROGRAM_EVALED_DATA_STATE.read().expect("read program data");
  s2.lookup(ns)?.0.lookup(def).map(|p| p.0).cloned()
}

pub fn load_by_index(ns_idx: u16, ns: &str, def_idx: u16, def: &str) -> Option<Calcit> {
  let s2 = PROGRAM_EVALED_DATA_STATE.read().expect("read program data");
  let (file, ns_cache) = s2.load(ns_idx).ok()?;
  if ns == ns_cache {
    let (value, def_cache) = file.load(def_idx).ok()?;
    if def == def_cache { Some(value.to_owned()) } else { None }
  } else {
    None
  }
}

// Dirty mutating global states
pub fn write_evaled_def(ns: &str, def: &str, value: Calcit) -> Result<(), String> {
  // println!("writing {} {}", ns, def);
  let mut program = PROGRAM_EVALED_DATA_STATE.write().expect("read program data");

  match program.lookup_mut(ns) {
    Some(file_entry) => {
      file_entry.0.insert(Arc::from(def), value);
    }
    None => {
      let mut file_entry = EntryBook::default();
      file_entry.insert(Arc::from(def), value);
      program.insert(Arc::from(ns), file_entry);
    }
  }

  Ok(())
}

// take a snapshot for codegen
pub fn clone_evaled_program() -> ProgramEvaledData {
  let program = &PROGRAM_EVALED_DATA_STATE.read().expect("read program data");
  (**program).to_owned()
}

pub fn apply_code_changes(changes: &snapshot::ChangesDict) -> Result<(), String> {
  let mut program_code = PROGRAM_CODE_DATA.write().expect("open program code");
  let coord0 = vec![];

  for (ns, file) in &changes.added {
    program_code.insert(ns.to_owned(), extract_file_data(file, ns.to_owned())?);
  }
  for ns in &changes.removed {
    program_code.remove(ns);
  }
  for (ns, info) in &changes.changed {
    // println!("handling ns: {:?} {}", ns, program_code.contains_key(ns));
    let file = program_code.get_mut(ns).ok_or_else(|| format!("can not load ns: {ns}"))?;
    if let Some(v) = &info.ns {
      file.import_map = extract_import_map(v, ns)?;
    }
    for (def, code) in &info.added_defs {
      let calcit_code = code_to_calcit(code, ns, def, coord0.to_owned())?;
      let entry = ProgramDefEntry {
        code: calcit_code,
        schema: DYNAMIC_TYPE.clone(),
        doc: Arc::from(""), // No doc info in changes, use empty string
        examples: vec![],   // No examples info in changes, use empty vector
      };
      file.defs.insert(def.to_owned().into(), entry);
    }
    for def in &info.removed_defs {
      file.defs.remove(def.as_str());
    }
    for (def, code) in &info.changed_defs {
      let calcit_code = code_to_calcit(code, ns, def, coord0.to_owned())?;
      let (schema, doc, examples) = match file.defs.get(def.as_str()) {
        Some(existing) => (existing.schema.clone(), existing.doc.clone(), existing.examples.clone()),
        None => (DYNAMIC_TYPE.clone(), Arc::from(""), Vec::new()),
      };
      file.defs.insert(
        def.to_owned().into(),
        ProgramDefEntry {
          code: calcit_code,
          schema,
          doc,
          examples,
        },
      );
    }
  }

  Ok(())
}

/// clear evaled data after reloading
pub fn clear_all_program_evaled_defs(init_ns: Arc<str>, reload_ns: Arc<str>, reload_libs: bool) -> Result<(), String> {
  let mut program = PROGRAM_EVALED_DATA_STATE.write().expect("open program data");
  if reload_libs {
    (*program).clear();
  } else {
    // reduce changes of libs. could be dirty in some cases
    let init_pkg = extract_pkg_from_ns(init_ns.to_owned()).ok_or_else(|| format!("failed to extract pkg from: {init_ns}"))?;
    let reload_pkg = extract_pkg_from_ns(reload_ns.to_owned()).ok_or_else(|| format!("failed to extract pkg from: {reload_ns}"))?;
    let mut to_remove: Vec<Arc<str>> = vec![];
    let xs = program.keys();
    for k in xs {
      if k == &init_pkg || k == &reload_pkg || k.starts_with(&format!("{init_pkg}.")) || k.starts_with(&format!("{reload_pkg}.")) {
        to_remove.push(k.to_owned());
      } else {
        continue;
      }
    }
    for k in to_remove {
      (*program).remove(&k);
    }
  }
  Ok(())
}
