use clap::Parser;
#[cfg(any(feature = "exact_cbc", feature = "exact_highs"))]
use clap::ValueEnum;
use egg::{FromOpError, RecExpr, RecExprParseError};
use eqmap::{
    asic::{CellAnalysis, CellLang, CellRpt, asic_rewrites},
    driver::{SynthRequest, logger_init, process_string_expression, simple_reader},
    rewrite::RewriteManager,
    verilog::SVModule,
};
use log::{debug, info, warn};
use std::path::PathBuf;

fn get_main_runner(
    s: &str,
) -> Result<SynthRequest<CellLang, CellAnalysis>, RecExprParseError<FromOpError>> {
    let expr: RecExpr<CellLang> = s.parse()?;
    let rules = asic_rewrites();

    Ok(SynthRequest::default()
        .with_asserts()
        .with_expr(expr)
        .with_rules(rules)
        .with_k(4)
        .without_progress_bar()
        .with_joint_limits(20, 20_000, 30))
}

#[allow(dead_code)]
/// parse an expression, simplify it with DSD and at most 4 fan-in, and pretty print it back out
fn simplify(s: &str) -> String {
    let mut req = get_main_runner(s).unwrap();
    req.synth::<CellRpt>().unwrap().get_expr().to_string()
}

#[allow(dead_code)]
/// parse an expression, simplify it with DSD and at most 4 fan-in, and pretty print it back out
fn simplify_w_proof(s: &str) -> String {
    let mut req = get_main_runner(s).unwrap().with_proof();
    req.synth::<CellRpt>().unwrap().get_expr().to_string()
}

#[cfg(any(feature = "exact_cbc", feature = "exact_highs"))]
#[derive(Debug, Clone, ValueEnum)]
enum Solver {
    #[cfg(feature = "exact_cbc")]
    Cbc,
    #[cfg(feature = "exact_highs")]
    Highs,
}

/// ASIC Technology Mapping Optimization with E-Graphs
#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
    /// Path to input file. If not provided, reads from stdin
    input: Option<PathBuf>,

    /// If provided, output a condensed JSON file with the e-graph
    #[cfg(feature = "graph_dumps")]
    #[arg(long)]
    dump_graph: Option<PathBuf>,

    /// If provided, use rules compiled from file instead of built-in rules
    #[arg(long)]
    rules: Option<PathBuf>,

    /// Use a cost model that weighs the cells by exact area
    #[arg(short = 'a', long, default_value_t = false)]
    area: bool,

    /// Do not verify the functionality of the output
    #[arg(short = 'f', long, default_value_t = false)]
    verify: bool,

    /// Only do algebraic simplification
    #[arg(short = 'c', long, default_value_t = false)]
    canonicalize: bool,

    /// Opt a specific LUT expr instead of from file
    #[arg(long)]
    command: Option<String>,

    /// Perform an exact extraction using ILP (much slower)
    #[cfg(any(feature = "exact_cbc", feature = "exact_highs"))]
    #[arg(long, value_enum)]
    exact: Option<Solver>,

    /// Print explanations (this generates a proof and runs longer)
    #[arg(short = 'v', long, default_value_t = false)]
    verbose: bool,

    /// Output Verilog
    #[arg(long, default_value_t = false)]
    verilog: bool,

    /// Extract for minimum circuit depth
    #[arg(long, default_value_t = false)]
    min_depth: bool,

    /// Max fan in size allowed for extracted Cells
    #[arg(short = 'k', long, default_value_t = 4)]
    k: usize,

    /// Build/extraction timeout in seconds
    #[arg(short = 't', long)]
    timeout: Option<u64>,

    /// Maximum number of nodes in graph
    #[arg(short = 's', long)]
    node_limit: Option<usize>,

    /// Maximum number of rewrite iterations
    #[arg(short = 'n', long)]
    iter_limit: Option<usize>,
}

fn main() -> std::io::Result<()> {
    let args = Args::parse();
    logger_init(args.verbose);

    if cfg!(debug_assertions) {
        warn!("Debug assertions are enabled");
    }

    let buf = simple_reader(args.command, args.input)?;

    let path = if let Some(p) = args.rules {
        p
    } else {
        let root = match std::env::var("EQMAP_ROOT") {
            Ok(root) => PathBuf::from(root),
            Err(_) => std::env::current_exe()?
                .parent()
                .unwrap()
                .parent()
                .unwrap()
                .parent()
                .unwrap()
                .to_path_buf(),
        };
        root.join("rules/asic.celllang")
    };

    info!("Loading rewrite rules from {path:?}");

    let mut rules = RewriteManager::<CellLang, CellAnalysis>::new();
    let file = std::fs::File::open(path)?;
    rules.parse_rules(file).map_err(std::io::Error::other)?;

    if args.canonicalize {
        rules.enable_category("algebraic");
    } else {
        let categories = rules.categories().cloned().collect::<Vec<_>>();
        for cat in categories {
            rules.enable_category(&cat);
        }
    }

    debug!(
        "Running with {} rewrite rules. Hash: {}",
        rules.num_active(),
        rules.rules_hash()
    );

    let rules = rules.active_rules();

    let req = SynthRequest::default()
        .with_rules(rules)
        .without_canonicalization();

    let req = match (args.timeout, args.node_limit, args.iter_limit) {
        (None, None, None) => req.with_joint_limits(2, 20_000, 9),
        (Some(t), None, None) => req.time_limited(t),
        (None, Some(n), None) => req.node_limited(n),
        (None, None, Some(i)) => req.iter_limited(i),
        (Some(t), Some(n), Some(i)) => req.with_joint_limits(t, n, i),
        _ => {
            return Err(std::io::Error::other(
                "Invalid build constraints (Use none, one, or three build constraints)",
            ));
        }
    };

    let req = if args.verbose { req.with_proof() } else { req };

    #[cfg(feature = "graph_dumps")]
    let req = match args.dump_graph {
        Some(p) => req.with_graph_dump(p),
        None => req,
    };

    let req = if args.canonicalize {
        req.with_ast_size()
    } else if args.min_depth {
        req.with_min_depth()
    } else if args.area {
        req.with_area()
    } else {
        req.with_k(args.k)
    };

    #[cfg(any(feature = "exact_cbc", feature = "exact_highs"))]
    let req = if let Some(solver) = &args.exact {
        let timeout = args.timeout.unwrap_or(600);
        match solver {
            #[cfg(feature = "exact_cbc")]
            Solver::Cbc => req.with_cbc(timeout),
            #[cfg(feature = "exact_highs")]
            Solver::Highs => req.with_highs(timeout),
        }
    } else {
        req
    };

    for line in buf.lines() {
        let result =
            process_string_expression::<CellLang, _, CellRpt>(line, req.clone(), !args.verify)?;
        if !result.is_empty() {
            if args.verilog {
                let module = SVModule::from_cells(
                    result.get_expr().to_owned(),
                    "top".to_string(),
                    Vec::new(),
                );
                let module = module.map_err(std::io::Error::other)?;
                print!(
                    "/* Generated by {} {} */\n\n{}",
                    env!("CARGO_PKG_NAME"),
                    env!("CARGO_PKG_VERSION"),
                    module
                );
            } else {
                println!("{result}");
            }
        }
    }
    Ok(())
}

#[test]
fn simple_tests() {
    assert_eq!(simplify("(AND a b)"), "(AND2_X1 a b)");
    assert_eq!(simplify("(INV a)"), "(INV_X1 a)");
    assert_eq!(simplify("(AND a true)"), "a");
}

#[test]
fn cell_rpt() {
    let mut req = get_main_runner("(INV a)").unwrap().with_report();
    let result = req.synth::<CellRpt>().unwrap();
    let rpt = result.write_report_to_string();
    assert!(rpt.is_ok());
    let rpt = rpt.unwrap();
    assert!(rpt.contains("name"));
}
