pub mod mus;
pub mod preprocess;
mod profiling;
mod proof;
mod propositional;
mod search;
mod subproblem;

use std::{
    ffi::OsStr,
    fs::{self, File, OpenOptions},
    io::prelude::*,
    path::PathBuf,
    process::{Command, Stdio},
    str::FromStr,
    time::{Duration, Instant},
};

use mus::{
    bnb::BnbMusEngine, forqes::Forqes, ifelse::IfElseEngine, trivial::TrivialBounder, MusEngine,
};
use proof::CompleteArgs;
use rand::{rngs::SmallRng, SeedableRng};
use search::{ProofMinimizer, ProofMinimizerState};
use subproblem::PartitionOrder;

use clap::Parser;
use dimacs::{parse_dimacs, Clause};
use log::{error, info, warn};

static VERSION: &str = concat!("v", env!("CARGO_PKG_VERSION"), "+", env!("GIT_HASH"));
static FORQES_URL: &str = "https://alexeyignatiev.github.io/assets/tools/forqes-linux-x86-64";

#[derive(Parser, Debug)]
#[command(about, long_about = None)]
#[command(version = VERSION)]
struct Args {
    /// Unsatisfiable CNF formula in DIMACS format
    formula_path: PathBuf,

    /// Random seed
    #[arg(long)]
    seed: Option<u64>,

    /// Switch the clause order in subproblem partitioning to ascending length
    #[arg(long)]
    ascending_length_partition: bool,

    /// Break ties between equal-length clauses in favor of less occurrences of literals
    #[arg(long)]
    ascending_frequency: bool,

    /// Limit the number of subproblems created in every branch-and-bound iteration.
    /// Setting this value makes the algorithm incomplete but drastically speeds up
    /// the exploration for problems with hundreds of clauses.
    #[arg(long)]
    max_subproblems_per_step: Option<usize>,

    /// Limit the maximum number of subproblems in the queue.
    /// Setting this value makes the algorithm incomplete but reduces the risk of
    /// halting with out-of-memory error.
    #[arg(long)]
    max_total_subproblems: Option<usize>,

    /// Do not use any search-based SMUS bounders during search
    #[arg(long)]
    trivial_bounding: bool,

    /// Use assignment cover strengthening for search-less SMUS bounds
    #[arg(long)]
    use_cover_bound: bool,

    /// Maximum formula size to trigger brute-force SMUS bounder
    #[arg(long, default_value_t = 32)]
    max_brute_force: usize,

    /// Feed input formula to Forqes via stdin rather than temporary files.
    /// Setting this option speeds up the individual bounding calls,
    /// especially the faster ones, but risks triggering a bug in Forqes'
    /// handling of stdio, which may lead to faulty results.
    #[arg(long)]
    forqes_no_file_io: bool,

    /// Path to Forqes executable; by default assume that it is available in PATH.
    #[arg(long)]
    forqes_binary: Option<PathBuf>,

    /// Randomize variable order in completion calls
    #[arg(short = 's', long)]
    shuffle_sat: bool,

    /// Use CaDiCaL for proof production
    #[arg(long)]
    cadical: bool,

    /// Exponential backoff: multiply the gap between the consecutive completion calls by
    /// this number after every completion call
    #[arg(long, default_value_t = 1.0)]
    complete_backoff_base: f64,

    /// Time limit for a single bounding call, ms
    #[arg(long, default_value_t = 1000)]
    bounding_time_limit_ms: u64,

    /// Time limit for the entire run, seconds
    #[arg(short = 't', long)]
    time_limit_seconds: Option<u64>,

    /// Output file for the best discovered proofs; defaults to stdout
    #[arg(short = 'o', long)]
    output_file: Option<PathBuf>,
}

impl Args {
    fn bounding_time_limit(&self) -> Duration {
        Duration::from_millis(self.bounding_time_limit_ms)
    }

    fn time_limit(&self) -> Option<Duration> {
        self.time_limit_seconds.map(Duration::from_secs)
    }

    fn partition_order(&self) -> PartitionOrder {
        if self.ascending_length_partition {
            if self.ascending_frequency {
                PartitionOrder::LengthAscFreqAsc
            } else {
                PartitionOrder::LengthAscFreqDsc
            }
        } else {
            if self.ascending_frequency {
                PartitionOrder::LengthDscFreqAsc
            } else {
                PartitionOrder::LengthDscFreqDsc
            }
        }
    }

    fn rng(&self) -> SmallRng {
        if let Some(seed) = self.seed {
            SmallRng::seed_from_u64(seed)
        } else {
            SmallRng::from_entropy()
        }
    }

    fn configure_solver(&self) -> ProofMinimizer<SmallRng> {
        let bounder = self.configure_bounder();
        let mut solver = ProofMinimizer::<SmallRng>::new(
            bounder,
            self.complete_args(),
            self.partition_order(),
            self.rng(),
        );
        solver.complete_backoff(self.complete_backoff_base);
        if let Some(max_split) = self.max_subproblems_per_step {
            solver.limit_branching(max_split);
        }
        if let Some(max_queue) = self.max_total_subproblems {
            solver.limit_queue(max_queue);
        }
        solver
    }

    fn configure_bounder(&self) -> Box<dyn MusEngine> {
        let trivial_bounder = self.configure_trivial_bounder();
        if self.trivial_bounding {
            Box::<TrivialBounder>::from(trivial_bounder)
        } else {
            self.ensure_forqes_binary();
            let max_brute_force = self.max_brute_force;
            let mut bounder =
                IfElseEngine::<BnbMusEngine<TrivialBounder>, Forqes<TrivialBounder>, _>::from(
                    move |formula: &[Clause], _: &_| formula.len() <= max_brute_force,
                );
            bounder.with_duration(self.bounding_time_limit());
            bounder.if_engine_mut().sub_bounder = trivial_bounder.clone();
            bounder.else_engine_mut().executable = self.forqes_path();
            bounder.else_engine_mut().fallback = trivial_bounder;
            if self.forqes_no_file_io {
                bounder.else_engine_mut().disable_file_io();
            }
            Box::new(bounder)
        }
    }

    fn configure_trivial_bounder(&self) -> TrivialBounder {
        TrivialBounder {
            use_cover_bound: self.use_cover_bound,
        }
    }

    fn forqes_path(&self) -> PathBuf {
        self.forqes_binary
            .clone()
            .unwrap_or(PathBuf::from_str("forqes").expect("failed to create the `forqes` PathBuf"))
    }

    fn ensure_forqes_binary(&self) {
        let output = Command::new(self.forqes_path())
            .arg("--help")
            .stdout(Stdio::piped())
            .output();
        match output {
            Ok(out) => {
                if String::from_utf8_lossy(&out.stdout).contains("Formula Quintessence Extractor") {
                    info!("Forqes executable has been detected");
                } else {
                    warn!("Forqes executable has not returned the expected help message");
                    warn!("Make sure you have obtained the binary from {}", FORQES_URL);
                }
            }
            Err(_) => {
                error!(
                    "Failed to run the Forqes binary; please download the binary from {}",
                    FORQES_URL
                );
            }
        }
    }

    fn complete_args(&self) -> CompleteArgs {
        CompleteArgs {
            randomize: self.shuffle_sat,
            use_cadical: self.cadical,
        }
    }
}

fn main() -> Result<(), ()> {
    env_logger::init();
    info!("Solver version: {VERSION}");
    let args = Args::parse();
    info!("Launched with the follwing arguments: {args:?}");
    // Parse the input CNF formula
    let dimacs_formula = if args
        .formula_path
        .extension()
        .and_then(OsStr::to_str)
        .unwrap_or_default()
        == "xz"
    {
        let mut decoder = xz::read::XzDecoder::new(
            File::open(&args.formula_path).expect("Should have been able to open the file"),
        );
        let mut dimacs_formula_buffer = String::new();
        decoder
            .read_to_string(&mut dimacs_formula_buffer)
            .expect("Should have been able to read the file");
        dimacs_formula_buffer
    } else {
        fs::read_to_string(&args.formula_path).expect("Should have been able to read the file")
    };
    let cnf = parse_dimacs(&dimacs_formula).expect("Should have been able to parse the formula");
    let clauses = if let dimacs::Instance::Cnf {
        clauses,
        num_vars: _,
    } = cnf
    {
        preprocess::preprocess(clauses.as_ref())
    } else {
        todo!("Arbitrary propositional formulas are not supported, please use CNF")
    };
    // Configure and run the solver
    let mut solver = args.configure_solver();
    let ProofMinimizerState { proof, bound } = solver.run(
        &clauses,
        args.time_limit()
            .map(|time_limit| Instant::now() + time_limit),
    );
    // Output the best proof parameters and stats once the search has completed
    if let Some(proof) = proof {
        info!("Best proof has {} steps", proof.len());
        if let Some(output_file) = args.output_file {
            // Write the best proof in the given file
            let writer = OpenOptions::new()
                .create(true)
                .truncate(true)
                .write(true)
                .append(false)
                .open(output_file.as_path())
                .unwrap();
            info!("Sending the output to {}", output_file.display());
            proof.write_tracecheck(writer).unwrap();
        } else {
            // Write the best proof to the standard output
            proof.write_tracecheck(std::io::stdout()).unwrap();
        }
    } else {
        warn!("Failed to find a proof");
    }
    info!("Proof size lower bound is {}", bound);
    Ok(())
}
