use std::{
    borrow::Cow,
    collections::HashMap,
    fmt::Display,
    fs::{self, File},
    io::{self, BufWriter, Write},
    path::{Path, PathBuf},
    sync::mpsc,
};

use okane_core::load;

use chrono::{Datelike, Days, NaiveDate, Weekday};
use pretty_decimal::PrettyDecimal;
use rust_decimal::Decimal;
use rust_decimal_macros::dec;

/// Metadata containing the reference to the generated input.
pub struct ExampleInput<T: FileSink> {
    rootdir: PathBuf,
    rootfile: PathBuf,
    pricedbfile: PathBuf,
    cleanup: bool,
    sink: T,
    params: InputParams,
}

impl<T: FileSink> Drop for ExampleInput<T> {
    fn drop(&mut self) {
        if self.cleanup {
            self.sink.cleanup(&self.rootdir);
        } else {
            log::debug!("{} is not set, leave the test data as-is", CLEANUP_KEY);
        }
    }
}

/// Error when creating [ExampleInput].
#[derive(Debug, thiserror::Error)]
pub enum ExampleInputError {
    #[error("failed to perform IO")]
    IO(#[from] std::io::Error),
    #[error("invlaid environment variable")]
    EnvVar(#[from] std::env::VarError),
}

const CLEANUP_KEY: &str = "OKANE_BENCH_CLEANUP";
const CREATE_KEY: &str = "OKANE_BENCH_CREATE";
const BENCH_TARGET_KEY: &str = "OKANE_BENCH_TARGET";

fn bool_from_env(env_key: &str) -> bool {
    !std::env::var(env_key)
        .unwrap_or_else(|err| match err {
            std::env::VarError::NotPresent => "".to_string(),
            std::env::VarError::NotUnicode(_) => {
                log::warn!("{env_key} was invalid unicode: {err}",);
                // invalid unicode = not empty.
                "1".to_string()
            }
        })
        .is_empty()
}

const YEAR_BEGIN: i32 = 2015;
const NUM_THREADS: usize = 2;

#[derive(Debug, Clone, Copy)]
pub struct InputParams {
    pub name: &'static str,
    pub num_years: i32,
    pub num_sub_files: usize,
    pub num_transactions_per_file: usize,
    pub sample_size: Option<usize>,
    pub more_commodity: bool,
}

impl Display for InputParams {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{0}", self.name)
    }
}

impl InputParams {
    pub fn params_from_env() -> impl Iterator<Item = Self> {
        let mut target = std::env::var(BENCH_TARGET_KEY).unwrap_or_default();
        if target.is_empty() {
            target = Self::middle_more_commodity().name.to_string();
        }
        [
            Self::small(),
            Self::middle(),
            Self::middle_more_commodity(),
            Self::large(),
        ]
        .into_iter()
        .scan(false, move |done, item| {
            if *done {
                return None;
            }
            if item.name == target {
                *done = true;
            }
            Some(item)
        })
    }

    pub fn small() -> Self {
        Self {
            name: "small_5y10a200t",
            num_years: 5,
            num_sub_files: 10,
            num_transactions_per_file: 200,
            sample_size: None,
            more_commodity: false,
        }
    }

    pub fn middle() -> Self {
        Self {
            name: "middle_10y16a500t",
            num_years: 10,
            num_sub_files: 16,
            num_transactions_per_file: 500,
            sample_size: None,
            more_commodity: false,
        }
    }

    pub fn middle_more_commodity() -> Self {
        Self {
            name: "middle_more_commodity_10y16a500t",
            num_years: 10,
            num_sub_files: 16,
            num_transactions_per_file: 500,
            sample_size: None,
            more_commodity: true,
        }
    }

    pub fn large() -> Self {
        Self {
            name: "large_20y30a800t",
            num_years: 20,
            num_sub_files: 30,
            num_transactions_per_file: 800,
            sample_size: Some(20),
            more_commodity: false,
        }
    }

    pub fn num_transactions(&self) -> u64 {
        (self.num_years as u64)
            * (self.num_sub_files as u64)
            * (self.num_transactions_per_file as u64)
    }

    fn years(&self) -> impl Iterator<Item = Year> {
        (YEAR_BEGIN..YEAR_BEGIN + self.num_years).map(Year)
    }

    /// Returns a `[start, end)` date window covering `length_pct` percent of
    /// the dataset's date span, starting `start_pct` percent in from the
    /// beginning. The dataset spans `[YEAR_BEGIN-01-01, (YEAR_BEGIN+num_years)-01-01)`.
    fn date_window(&self, start_pct: u32, length_pct: u32) -> (NaiveDate, NaiveDate) {
        debug_assert!(start_pct + length_pct <= 100);
        let begin = NaiveDate::from_ymd_opt(YEAR_BEGIN, 1, 1).unwrap();
        let end = NaiveDate::from_ymd_opt(YEAR_BEGIN + self.num_years, 1, 1).unwrap();
        let total_days = (end - begin).num_days() as u64;
        let start_offset = total_days * start_pct as u64 / 100;
        let end_offset = total_days * (start_pct + length_pct) as u64 / 100;
        let start = begin.checked_add_days(Days::new(start_offset)).unwrap();
        let stop = begin.checked_add_days(Days::new(end_offset)).unwrap();
        (start, stop)
    }

    /// First 10% of the dataset's date span.
    pub fn date_range_first(&self) -> (NaiveDate, NaiveDate) {
        self.date_window(0, 10)
    }

    /// Middle 20% of the dataset's date span (40%..60%).
    pub fn date_range_middle(&self) -> (NaiveDate, NaiveDate) {
        self.date_window(40, 20)
    }

    /// Last 30% of the dataset's date span.
    pub fn date_range_last(&self) -> (NaiveDate, NaiveDate) {
        self.date_window(70, 30)
    }
}

impl<T: FileSink + Send> ExampleInput<T> {
    /// Creates an example used for benchmarks.
    /// Created example is left as-is, unless `OKANE_BENCH_CLEANUP` is set.
    /// If `OKANE_BENCH_CLEANUP` is set,
    /// * Always recreate the input.
    /// * Clean up the created input.
    pub fn new(subdir: &Path, params: InputParams) -> Result<Self, ExampleInputError> {
        let cleanup = bool_from_env(CLEANUP_KEY);
        let create = bool_from_env(CREATE_KEY);
        let rootdir = T::rootdir().join(subdir);
        let rootfile = rootdir.join("root.ledger");
        let pricedbfile = rootdir.join("price.db");
        if !create && !cleanup && T::shortcut(&rootfile) && T::shortcut(&pricedbfile) {
            return Ok(Self {
                rootdir,
                rootfile,
                pricedbfile,
                cleanup,
                sink: T::new(),
                params,
            });
        }
        let before_input = std::time::Instant::now();
        T::initialize(&rootdir)?;
        let mut tasks = Vec::new();
        let (tx, rx) = mpsc::channel();
        for i1 in 0..NUM_THREADS {
            let thread_tx = tx.clone();
            let rootdir = rootdir.clone();
            tasks.push(std::thread::spawn(move || {
                let ret = || -> Result<T, io::Error> {
                    let mut sink = T::new();
                    for i2 in 0..params.num_sub_files / NUM_THREADS {
                        let i = i1 * params.num_sub_files / NUM_THREADS + i2;
                        let dirname = dir_name(i);
                        let subdirpath = rootdir.join(&dirname);
                        T::create_subdir(&subdirpath)?;
                        let mut total: i64 = 0;
                        for year in params.years() {
                            total = prepare_leaf_file(
                                &mut sink,
                                &subdirpath.join(leaf_file(year)),
                                &params,
                                i,
                                year,
                                total,
                            )?;
                        }
                        prepare_middle_file(&mut sink, &rootdir, &dirname, &params)?;
                    }
                    Ok(sink)
                }();
                thread_tx.send(ret).expect("send must not fail");
            }));
        }
        let mut sink = T::new();
        prepare_root_file(&mut sink, &rootfile, 0..params.num_sub_files)?;
        prepare_price_db_file(&mut sink, &pricedbfile, YEAR_BEGIN, params.num_years)?;
        for _ in 0..tasks.len() {
            let another = rx
                .recv_timeout(std::time::Duration::from_secs(150))
                .expect("Can't wait 150 seconds deadline on the recv task")?;
            sink.merge(another);
        }
        for jh in tasks.into_iter() {
            jh.join().expect("thread join must not fail");
        }
        let after_input = std::time::Instant::now();
        let duration = after_input - before_input;
        log::debug!("input creation took {:.3} seconds", duration.as_secs_f64());
        Ok(ExampleInput {
            rootdir,
            rootfile,
            pricedbfile,
            cleanup,
            sink,
            params,
        })
    }

    pub fn rootpath(&self) -> &Path {
        &self.rootfile
    }

    pub fn new_loader(&self) -> load::Loader<T::FileSystem> {
        load::Loader::new(self.rootfile.clone(), self.sink.clone_as_filesystem())
    }

    pub fn pricedbpath(&self) -> &Path {
        &self.pricedbfile
    }

    pub fn num_transactions(&self) -> u64 {
        self.params.num_transactions()
    }
}

pub trait FileSink: Send + Sized + 'static {
    type FileSystem: load::FileSystem;

    fn new_example(
        subdir: &Path,
        params: InputParams,
    ) -> Result<ExampleInput<Self>, ExampleInputError> {
        ExampleInput::<Self>::new(subdir, params)
    }

    fn new() -> Self;

    fn rootdir() -> PathBuf;

    fn clone_as_filesystem(&self) -> Self::FileSystem;

    fn initialize(rootdir: &Path) -> Result<(), std::io::Error>;

    fn create_subdir(subdir: &Path) -> Result<(), std::io::Error>;

    fn shortcut(rootfile: &Path) -> bool;

    fn merge(&mut self, other: Self);

    /// Gives the writer for the given path.
    fn writer<'a>(
        &'a mut self,
        path: &Path,
    ) -> Result<Box<dyn std::io::Write + 'a>, std::io::Error>;

    fn cleanup(&self, rootdir: &Path);
}

pub struct RealFileSink;

impl FileSink for RealFileSink {
    type FileSystem = load::ProdFileSystem;

    fn new() -> Self {
        RealFileSink
    }

    fn rootdir() -> PathBuf {
        PathBuf::from(env!("CARGO_TARGET_TMPDIR"))
    }

    fn clone_as_filesystem(&self) -> Self::FileSystem {
        load::ProdFileSystem
    }

    fn shortcut(rootfile: &Path) -> bool {
        match fs::metadata(rootfile) {
            Err(error) => {
                log::warn!(
                    "std::fs::metadata() failed on {}, retry creation: {}",
                    rootfile.to_string_lossy(),
                    error
                );

                false
            }
            Ok(_) => true,
        }
    }

    fn initialize(rootdir: &Path) -> Result<(), std::io::Error> {
        fs::remove_dir_all(rootdir).or_else(|e| match e.kind() {
            io::ErrorKind::NotFound => Ok(()),
            _ => Err(e),
        })?;
        fs::create_dir_all(rootdir)
    }

    fn create_subdir(subdir: &Path) -> Result<(), std::io::Error> {
        fs::create_dir(subdir)
    }

    fn merge(&mut self, _other: Self) {}

    fn writer<'a>(
        &'a mut self,
        path: &Path,
    ) -> Result<Box<dyn std::io::Write + 'a>, std::io::Error> {
        Ok(Box::new(BufWriter::new(File::create(path)?)))
    }

    fn cleanup(&self, rootdir: &Path) {
        let _ignore = std::fs::remove_dir_all(rootdir).inspect_err(|x| {
            log::error!(
                "failed to clean up the directory {}: {}",
                rootdir.display(),
                x
            );
        });
    }
}

pub struct FakeFileSink {
    files: HashMap<PathBuf, Vec<u8>>,
}

impl FileSink for FakeFileSink {
    type FileSystem = load::FakeFileSystem;

    fn new() -> Self {
        Self {
            files: HashMap::default(),
        }
    }

    fn rootdir() -> PathBuf {
        PathBuf::from("path/to/")
    }

    fn clone_as_filesystem(&self) -> Self::FileSystem {
        self.files.clone().into()
    }

    fn shortcut(_rootfile: &Path) -> bool {
        false
    }

    fn initialize(_rootdir: &Path) -> Result<(), std::io::Error> {
        Ok(())
    }

    fn create_subdir(_subdir: &Path) -> Result<(), std::io::Error> {
        Ok(())
    }

    fn merge(&mut self, other: Self) {
        for (k, v) in other.files {
            self.files.insert(k, v);
        }
    }

    fn writer<'a>(
        &'a mut self,
        path: &Path,
    ) -> Result<Box<dyn std::io::Write + 'a>, std::io::Error> {
        Ok(Box::new(self.files.entry(path.to_owned()).or_default()))
    }

    fn cleanup(&self, _rootdir: &Path) {}
}

fn prepare_price_db_file<T: FileSink>(
    sink: &mut T,
    pricedbfile: &Path,
    year_begin: i32,
    num_years: i32,
) -> Result<(), std::io::Error> {
    let mut w = sink.writer(pricedbfile)?;
    let begin = NaiveDate::from_ymd_opt(year_begin, 1, 1).unwrap();
    let end = NaiveDate::from_ymd_opt(year_begin + num_years, 1, 1).unwrap();
    for d in begin.iter_days() {
        if d >= end {
            break;
        }
        if d.weekday() == Weekday::Sat || d.weekday() == Weekday::Sun {
            continue;
        }
        writeln!(w, "P {} USD {} JPY", d.format("%Y/%m/%d"), usdjpy(d))?;
        writeln!(w, "P {} CHF {} JPY", d.format("%Y/%m/%d"), chfjpy(d))?;
    }
    w.flush()?;
    Ok(())
}

fn usdjpy(date: NaiveDate) -> Decimal {
    let rand = (date.year() * 5 + date.month() as i32) % 97
        + ((date.month() * 7 + date.day() * 13) as i32) % 13;
    dec!(120) + Decimal::from(rand)
}

fn chfjpy(date: NaiveDate) -> Decimal {
    let rand = (date.year() * 5 + date.month() as i32) % 97
        + ((date.month() * 7 + date.day() * 13) as i32) % 13;
    dec!(160) + Decimal::from(rand)
}

fn prepare_root_file<T: FileSink>(
    sink: &mut T,
    rootfile: &Path,
    dirs: std::ops::Range<usize>,
) -> Result<(), std::io::Error> {
    let mut w = sink.writer(rootfile)?;
    for dir in dirs {
        writeln!(w, "include {}.ledger", dir_name(dir))?;
    }
    Ok(())
}

fn prepare_middle_file<T: FileSink>(
    sink: &mut T,
    rootdir: &Path,
    dirname: &str,
    params: &InputParams,
) -> Result<(), std::io::Error> {
    let mut target = PathBuf::from(rootdir);
    target.push(format!("{}.ledger", dirname));
    let mut w = sink.writer(&target)?;
    for year in params.years() {
        let leaf = leaf_file(year);
        writeln!(w, "include {dirname}/{leaf}")?;
    }
    Ok(())
}

fn payee(dir: usize, year: Year, i: usize) -> String {
    format!(
        "Random payee {}",
        ((dir as i32) * 31 + year.0 * 17 + (i as i32) * 7) % 13
    )
}

struct Amount<'a>(PrettyDecimal, Cow<'a, str>);

impl<'a> Display for Amount<'a> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{0} {1}", self.0, self.1)
    }
}

/// Creates the leaf file, returns the total amount of the asset.
fn prepare_leaf_file<T: FileSink>(
    sink: &mut T,
    target: &Path,
    params: &InputParams,
    dir: usize,
    year: Year,
    mut total: i64,
) -> Result<i64, std::io::Error> {
    let initial_week = NaiveDate::from_ymd_opt(YEAR_BEGIN, 1, 7).unwrap();
    let multi_commodity =
        |date: NaiveDate| date.weekday().number_from_monday() <= 3 || date < initial_week;
    let mut w = sink.writer(target)?;
    for i in 0..params.num_transactions_per_file {
        let ordinal = (i * 365 / params.num_transactions_per_file + 1) as u32;
        let date = NaiveDate::from_yo_opt(year.0, ordinal).expect("must be a valid date");
        let date_fmt = date.format("%Y/%m/%d");
        let payee = payee(dir, year, i);
        let pseudo = ((dir as i32) * 31 + year.0 * 17 + (i as i32) * 7) % 13;
        let (other_account, other_amount, amount): (Cow<str>, _, _) = match pseudo {
            0..=10 => {
                if pseudo >= 9 && params.more_commodity && date.weekday() == Weekday::Mon {
                    let pseudo2 = ((dir as i32) * 31 + year.0 * 17 + (i as i32) * 7) % 26;
                    let pseudo3 = ((dir as i32) * 31 + year.0 * 17 + (i as i32) * 7) / 26 % 26;
                    let commodity = format!(
                        "COMMODITY_{}{}",
                        nth_alphabet(pseudo2),
                        nth_alphabet(pseudo3)
                    );
                    total -= 3000;
                    (
                        format!("Expenses:Type{}", pseudo).into(),
                        Amount(PrettyDecimal::comma3dot(dec!(100)), commodity.into()),
                        Amount(PrettyDecimal::comma3dot(dec!(-3000)), "JPY".into()),
                    )
                } else {
                    total -= 1234;
                    (
                        format!("Expenses:Type{}", pseudo).into(),
                        Amount(PrettyDecimal::comma3dot(dec!(1234)), "JPY".into()),
                        Amount(PrettyDecimal::comma3dot(dec!(-1234)), "JPY".into()),
                    )
                }
            }
            // limits the currency conversion events to Mon/Tue/Wed.
            11 if multi_commodity(date) => {
                total += 400000;
                (
                    "Income:Salary".into(),
                    Amount(
                        PrettyDecimal::comma3dot(dec!(-3000.00)),
                        Cow::Borrowed("USD"),
                    ),
                    Amount(PrettyDecimal::comma3dot(dec!(400000)), "JPY".into()),
                )
            }
            11 => {
                total += 400000;
                (
                    "Income:Salary".into(),
                    Amount(
                        PrettyDecimal::comma3dot(dec!(-400000)),
                        Cow::Borrowed("JPY"),
                    ),
                    Amount(PrettyDecimal::comma3dot(dec!(400000)), "JPY".into()),
                )
            }
            // limits the currency conversion frequency to Mon/Tue/Wed.
            12 if multi_commodity(date) => {
                total -= 14000;
                (
                    "Assets:Account99".into(),
                    Amount(PrettyDecimal::comma3dot(dec!(100.00)), "CHF".into()),
                    Amount(PrettyDecimal::comma3dot(dec!(-14000)), "JPY".into()),
                )
            }
            12 => {
                total -= 14000;
                (
                    "Assets:Account99".into(),
                    Amount(PrettyDecimal::comma3dot(dec!(14000)), "JPY".into()),
                    Amount(PrettyDecimal::comma3dot(dec!(-14000)), "JPY".into()),
                )
            }
            _ => unreachable!("this can't happen"),
        };
        writeln!(w, "{date_fmt} {payee}",)?;
        writeln!(w, "  Assets:Account{dir:02}     {amount} = {total} JPY",)?;
        writeln!(w, "  {other_account}     {other_amount}",)?;
        writeln!(w)?;
    }
    Ok(total)
}

fn nth_alphabet(i: i32) -> char {
    char::from_u32('a' as u32 + i as u32 % 26).unwrap()
}

fn dir_name(i: usize) -> String {
    format!("sub-{:02}", i)
}

fn leaf_file(year: Year) -> String {
    format!("{:04}.ledger", year.0)
}

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
struct Year(i32);
