pub mod parser;
pub(crate) mod dimacs_ext;

use std::{ffi::OsStr, fs, io::Read, path::PathBuf, io::BufReader, collections::BTreeMap};

use clap::Parser;
use dimacs::parse_dimacs;
use log::{warn, debug, info, error};
use parser::LratParser;

use std::fs::File;

use crate::{dimacs_ext::ToDimacs as _, parser::LratDerivation};

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

    /// Proof of unsatisfiability in LRAT format
    proof_path: PathBuf,

    /// Output JSON string with duplicate count distribution
    #[arg(short = 'd', long)]
    count_duplicates: bool,
}

fn main() {
    env_logger::init();
    let args = Args::parse();
    // 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 =
        match parse_dimacs(&dimacs_formula).expect("Should have been able to parse the formula") {
            dimacs::Instance::Cnf { clauses, .. } => clauses.to_vec(),
            _ => todo!("Arbitrary SAT formulas are not supported yet"),
        };
    let mut lrat_parser = LratParser::from(&cnf[..]);
    for line in lrat_parser.parse(BufReader::new(File::open(&args.proof_path).expect("Should have been able to open the proof file"))) {
        match line {
            Err(line) => {
                error!("Failed to parse the line {line} as an LRAT step");
                panic!("{line} is not a valid step");
            }
            Ok(LratDerivation::Delete) => {
                // Deletions are irrelevant, keep going
            },
            Ok(LratDerivation::Derive { clause, is_rat, chain_length }) => {
                if is_rat {
                    warn!("Discovered a RAT lemma {}", clause.to_dimacs());
                }
                else {
                    debug!("Derived lemma {} with {chain_length} resolutions", clause.to_dimacs());
                }
            }
        }
    }
    info!("Derived {} lemmas from {} axioms with {} RAT lemmas",
        lrat_parser.count_derived(),
        lrat_parser.count_axioms(),
        lrat_parser.count_rat_steps()
    );
    if args.count_duplicates {
        let total_repeated = lrat_parser.repetitions().map(|(_, count)| count).sum::<usize>();
        let mut count_distribution = BTreeMap::<usize, usize>::default();
        for count in lrat_parser.repetitions().map(|(_, count)| count) {
            *count_distribution.entry(count).or_default() += 1;
        }
        if total_repeated > 0 {
            info!("Removed {total_repeated} repeated clause derivations");
            for (num_repetitions, num_clauses) in &count_distribution {
                info!("- {} clauses have been repeated {} times", num_clauses, num_repetitions);
            }
        }
        let repetitions_str = format!("[{}]", count_distribution.iter()
            .map(|(reps, count)| format!(r#"{{"reps":{reps},"cls":{count}}}"#))
            .collect::<Vec<_>>()
            .join(",")
        );
        println!("{} {} {} {repetitions_str}",
            lrat_parser.count_axioms(), lrat_parser.count_derived(), lrat_parser.count_rat_steps()
        );
    } else{ 
        println!("{} {} {}",
            lrat_parser.count_axioms(), lrat_parser.count_derived(), lrat_parser.count_rat_steps()
        );
    }
}
