use std::{
    collections::BTreeMap,
    io::Write,
    path::PathBuf,
    process::{Command, Stdio},
    str::FromStr,
    time::Instant,
};

use dimacs::Clause;
use itertools::Itertools;
use log::{error, info, warn};
use tap::Pipe;
use tempfile::NamedTempFile;

use crate::propositional::to_i64;

use super::{MusEngine, MusStatus};

#[derive(Debug)]
pub struct Forqes<FallbackMus: MusEngine + Default> {
    pub executable: PathBuf,
    use_file_io: bool,
    cache: Vec<(Vec<Clause>, MusStatus)>,
    formula_size: BTreeMap<usize, usize>,
    mandatory_size: BTreeMap<usize, usize>,
    pub fallback: FallbackMus,
}

impl<FallbackMus: MusEngine + Default> From<PathBuf> for Forqes<FallbackMus> {
    fn from(executable: PathBuf) -> Self {
        Forqes {
            executable,
            ..Default::default()
        }
    }
}

impl<FallbackMus: MusEngine + Default> Default for Forqes<FallbackMus> {
    fn default() -> Self {
        Self {
            executable: PathBuf::from_str("forqes").expect("failed to create PathBuf"),
            cache: vec![],
            formula_size: Default::default(),
            mandatory_size: Default::default(),
            use_file_io: true,
            fallback: Default::default(),
        }
    }
}

impl<FallbackMus: MusEngine + Default> MusEngine for Forqes<FallbackMus> {
    fn optimize_with_deadline(
        &mut self,
        formula: &[Clause],
        mandatory: &[Clause],
        deadline: Option<Instant>,
    ) -> super::MusStatus {
        if let Some(empty) = formula.iter().find(|clause| clause.len() == 0) {
            return MusStatus {
                smallest_subset: vec![empty.clone()],
                lower_bound: 1,
            };
        }
        let fallback_status = self
            .fallback
            .optimize_with_deadline(formula, mandatory, deadline);
        match self.try_optimize(formula, mandatory, deadline) {
            Ok(status) => MusStatus {
                smallest_subset: status.smallest_subset,
                lower_bound: status.lower_bound.max(fallback_status.lower_bound),
            },
            Err(e) => {
                error!("Forqes has crashed with error {e:?}");
                fallback_status
            }
        }
    }

    fn report(&self) {
        info!("Formula size distribution: {:?}", self.formula_size);
        info!("Mandatory size distribution: {:?}", self.mandatory_size);
    }
}

impl<FallbackMus: MusEngine + Default> Forqes<FallbackMus> {
    pub fn disable_file_io(&mut self) {
        self.use_file_io = false;
    }

    fn try_optimize(
        &mut self,
        formula: &[dimacs::Clause],
        mandatory: &[dimacs::Clause],
        deadline: Option<Instant>,
    ) -> std::io::Result<super::MusStatus> {
        self.mandatory_size
            .entry(mandatory.len())
            .and_modify(|c| *c += 1)
            .or_insert(1);
        self.formula_size
            .entry(formula.len())
            .and_modify(|c| *c += 1)
            .or_insert(1);
        if let Some(status) = self.lookup(formula) {
            Ok(status)
        } else {
            let mut status = self.call_binary(formula, deadline)?;
            self.cache
                .push((formula.iter().cloned().collect_vec(), status.clone()));
            status.lower_bound = status.lower_bound.max(mandatory.len());
            Ok(status)
        }
    }

    fn lookup(&self, formula: &[dimacs::Clause]) -> Option<MusStatus> {
        self.cache
            .iter()
            .find_map(|(key, value)| if key == formula { Some(value) } else { None })
            .cloned()
    }

    fn call_binary(
        &mut self,
        formula: &[dimacs::Clause],
        deadline: Option<Instant>,
    ) -> std::io::Result<super::MusStatus> {
        let cutoff = deadline.map(|d| (Instant::now() - d).as_secs_f32().ceil() as u32);
        let dimacs = self.prepare_input(formula);
        let output = Command::new(&self.executable)
            .arg("-vv")
            .pipe(|cmd| {
                if let Some(cutoff) = cutoff {
                    cmd.arg("-t").arg(format!("{}", cutoff.max(1)))
                } else {
                    cmd
                }
            })
            .stdout(Stdio::piped())
            .pipe(|cmd| {
                let output = if self.use_file_io {
                    let mut formula_file = NamedTempFile::new()?;
                    write!(formula_file, "{}", dimacs)?;
                    cmd.arg(formula_file.path()).spawn()?.wait_with_output()?
                } else {
                    let mut child = cmd.stdin(Stdio::piped()).spawn()?;
                    let stdin = child.stdin.as_mut().unwrap();
                    stdin.write_all(dimacs.as_bytes())?;
                    child.wait_with_output()?
                };
                Ok::<std::process::Output, std::io::Error>(output)
            })?;
        let mut status = MusStatus {
            smallest_subset: formula.iter().cloned().collect_vec(),
            lower_bound: 1,
        };
        for line in String::from_utf8_lossy(&output.stdout).split('\n') {
            let prefixes = ["c disj MDSes: ", "c constr: ", "o "];
            for prefix in prefixes {
                if line.starts_with(prefix) {
                    if let Some(lower_bound_str) = line.strip_prefix(prefix) {
                        if let Ok(lower_bound) = usize::from_str(lower_bound_str) {
                            status.lower_bound = status.lower_bound.max(lower_bound);
                        } else {
                            warn!("Failed to parse the line `{line}` in Forqes output as a bound");
                        }
                    }
                }
            }
        }
        Ok(status)
    }

    fn prepare_input(&self, formula: &[Clause]) -> String {
        let n_variables = formula
            .iter()
            .flat_map(|clause| clause.lits().iter().map(|lit| lit.var().0))
            .max()
            .unwrap_or(0);
        let n_clauses = formula.len();
        let mut dimacs = String::new();
        dimacs.push_str(format!("p cnf {n_variables} {n_clauses}\n").as_str());
        for clause in formula {
            let mut clause_str = String::new();
            for lit in clause.lits() {
                clause_str.push_str(format!("{} ", to_i64(lit)).as_str());
            }
            clause_str.push_str("0\n");
            dimacs.push_str(clause_str.as_str());
        }
        dimacs
    }
}
