use anyhow::{bail, Context, Result};
use clap::Parser;
use glob::glob;
use pimcore::json_to_instruction::json_to_executor;
use pimcore::tracing::TRACER;
use serde_json::Value;
use std::fs;
use std::io::Write;
use std::path::PathBuf;

/// Program to simulate core execution configuration
#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// The folder containing the configuration files
    #[arg(short, long)]
    folder: Option<PathBuf>,

    /// Override path for the config.json file
    #[arg(long)]
    config: Option<PathBuf>,

    /// Override path for the core*.json files
    #[arg(long, value_delimiter = ',', num_args = 1..)]
    cores: Option<Vec<PathBuf>>,

    /// Override path for the memory file (.txt or .bin)
    #[arg(long)]
    memory: Option<PathBuf>,

    /// Path to output file
    #[arg(short, long)]
    output: PathBuf,

    /// Comma separated list of (address,size) for memory output dump
    #[arg(short, long, value_delimiter = ',', num_args = 1.., value_name = "ADDR,SIZE")]
    dump: Vec<i32>,
}

fn main() -> Result<()> {
    let args = Args::parse();

    let config_json = retrive_config(&args)?;
    let core_jsons = retrive_cores(&args)?;
    let memory = retrive_memory(&args)?;
    let mut executor = json_to_executor::json_to_executor(config_json, core_jsons.iter());
    populate_crossbar(&args, &mut executor);
    set_memory(&mut executor, memory);
    TRACER
        .lock()
        .unwrap()
        .init(executor.cpu().num_core(), args.output.clone());
    executor.execute();
    dump_memory(executor, &args)?;
    Ok(())
}

fn populate_crossbar(args: &Args, executor: &mut pimcore::Executable) {
    let num_cores = executor.cpu_mut().num_core();

    if let Some(folder) = args.folder.as_ref() {
        for core_idx in 0..num_cores {
            let core_folder = folder.join(format!("core_{}", core_idx));

            if !core_folder.is_dir() {
                continue;
            }

            let mut bin_files: Vec<(u32, std::path::PathBuf)> = std::fs::read_dir(&core_folder)
                .expect("Failed to read core directory")
                .filter_map(|entry| {
                    let path = entry.ok()?.path();
                    let file_name = path.file_name()?.to_str()?;

                    if file_name.starts_with("crossbar_") && file_name.ends_with(".bin") {
                        let num_str = &file_name[9..file_name.len() - 4];
                        let num = num_str.parse::<u32>().ok()?;
                        Some((num, path))
                    } else {
                        None
                    }
                })
                .collect();
            bin_files.sort_by_key(|&(num, _)| num);
            let core = executor.cpu_mut().core(core_idx);
            let (_memory, crossbars) = core.get_memory_crossbar();

            for (i, path) in bin_files {
                let bytes = std::fs::read(path).expect("Failed to read binary file");
                crossbars
                    .get_mut(i as usize)
                    .unwrap()
                    .execute_store(&bytes)
                    .unwrap();
            }
        }
    }
}

fn dump_memory(mut executor: pimcore::Executable, args: &Args) -> Result<()> {
    let dumps: Vec<(i32, i32)> = args
        .dump
        .chunks_exact(2)
        .map(|chunk| (chunk[0], chunk[1]))
        .collect();
    let mut out_file = fs::OpenOptions::new()
        .create(true)
        .write(true)
        .truncate(true)
        .open(&args.output)
        .with_context(|| format!("cannot open file {:?} for writing", args.output))?;

    for (address, size) in dumps {
        out_file.write_all(executor.cpu_mut().host().load::<u8>(address, size).unwrap()[0])?;
    }
    Ok(())
}

fn set_memory(executor: &mut pimcore::Executable, memory: Vec<u8>) {
    executor.cpu_mut().host().execute_store(0, &memory).unwrap();
}

fn retrive_memory(args: &Args) -> Result<Vec<u8>> {
    let memory_path = if let Some(mem_override) = &args.memory {
        mem_override.clone()
    } else if let Some(folder) = &args.folder.as_ref() {
        let bin_path = folder.join("memory.bin");
        let txt_path = folder.join("memory.txt");

        if bin_path.exists() {
            bin_path
        } else if txt_path.exists() {
            txt_path
        } else {
            bail!(
                "No memory file (memory.bin or memory.txt) found in {:?}",
                folder
            );
        }
    } else {
        bail!("Either --memory or --folder must be provided.");
    };
    let memory_vector: Vec<u8> = if let Some(ext) = memory_path.extension() {
        if ext == "bin" {
            fs::read(&memory_path).context("Failed to read binary memory file")?
        } else {
            let content =
                fs::read_to_string(&memory_path).context("Failed to read ascii memory file")?;
            let clean_hex: String = content.chars().filter(|c| !c.is_whitespace()).collect();
            hex::decode(&clean_hex).context("Failed to decode ASCII hex string")?
        }
    } else {
        bail!("Memory file has no extension, cannot determine type");
    };
    Ok(memory_vector)
}

fn retrive_cores(args: &Args) -> Result<Vec<Value>, anyhow::Error> {
    let mut core_jsons: Vec<Value> = Vec::new();
    if let Some(cores_override) = &args.cores {
        for core in cores_override {
            let content = fs::read_to_string(core)
                .with_context(|| format!("Failed to read core file: {:?}", cores_override))?;
            let json: Value =
                serde_json::from_str(&content).context("Failed to parse core json override")?;
            core_jsons.push(json);
        }
    } else if let Some(folder) = args.folder.as_ref() {
        let pattern = folder.join("core*.json");
        let pattern_str = pattern.to_str().context("Invalid path encoding")?;
        let mut paths: Vec<_> = glob(pattern_str)?.map(|x| x.unwrap()).collect();
        paths.sort_by_cached_key(|x| {
            let mut x = x.file_stem().expect("Extracting the stem").to_str().expect("File not utf-8");
            x = &x[5..];
            x.parse::<i32>().unwrap()
        });

        if paths.is_empty() {
            bail!("No core*.json files found in {:?}", folder);
        }
        for entry in paths {
            let path = entry;
            let content = fs::read_to_string(&path)
                .with_context(|| format!("Failed to read core file: {:?}", path))?;
            let json: Value = serde_json::from_str(&content)
                .with_context(|| format!("Failed to parse JSON in {:?}", path))?;
            core_jsons.push(json);
        }
    } else {
        bail!("Either --core or --folder must be provided to find core definitions.");
    }
    Ok(core_jsons)
}

fn retrive_config(args: &Args) -> Result<Value, anyhow::Error> {
    let config_path: PathBuf = {
        let override_path = args.config.as_ref();
        let folder = args.folder.as_ref();
        let filename = "config.json";
        if let Some(path) = override_path {
            path.clone()
        } else if let Some(folder) = folder {
            let path = folder.join(filename);
            if path.exists() {
                path
            } else {
                bail!("File {} not found in folder {:?}", filename, folder);
            }
        } else {
            bail!(
                "Cannot resolve {}: no override path and no input folder provided.",
                filename
            );
        }
    };
    let config_content = fs::read_to_string(&config_path)
        .with_context(|| format!("Failed to read config file: {:?}", config_path))?;
    let config_json: Value =
        serde_json::from_str(&config_content).context("Failed to parse config.json")?;
    Ok(config_json)
}
