use crate::utils::misc::{create_ticker, split_ticker};
use crate::utils::qdf::*;
use anyhow;
use log;
use polars::datatypes::DataType;
use polars::prelude::*;
use rayon::prelude::*;
use std::error::Error;

/// The number of concurrent file loads to perform.
const CONCURRENT_FILE_LOADS: usize = 500;

fn _file_base_name(file_path: String) -> String {
    std::path::Path::new(&file_path.clone())
        .file_stem()
        .unwrap()
        .to_str()
        .unwrap()
        .to_string()
}

/// Load a Quantamental DataFrame from a CSV file.
/// The CSV must be named in the format `cid_xcat.csv` (`ticker.csv`).
/// The DataFrame must have a `real_date` column along with additional value columns.
pub fn load_quantamental_dataframe(
    file_path: String,
) -> Result<DataFrame, Box<dyn std::error::Error>> {
    // get the file base name
    let base_file_name = _file_base_name(file_path.clone().into());

    // if filename does not have _ then it is not a Quantamental DataFrame
    if !base_file_name.contains('_') {
        return Err("The file name must be in the format `cid_xcat.csv` (`ticker.csv`)".into());
    }

    let ticker = base_file_name.split('.').collect::<Vec<&str>>()[0];
    let (cid, xcat) = split_ticker(ticker.to_string())?;

    let mut df = CsvReadOptions::default()
        .try_into_reader_with_file_path(Some(file_path.to_string().into()))
        .unwrap()
        .finish()
        .unwrap();

    let err = "The dataframe must have a `real_date` column and atleast 1 additional value column";
    if df.column("real_date").is_err() || df.width() < 2 {
        return Err(err.into());
    }

    // check if the first item in the real_date column has a dash or not
    let has_dashes = df
        .column("real_date")
        .unwrap()
        .cast(&DataType::String)?
        .get(0)
        .unwrap()
        .to_string()
        .contains('-');

    let date_format = if has_dashes { "%Y-%m-%d" } else { "%Y%m%d" };
    // let real_date_col = df
    //     .column("real_date".into())
    //     .unwrap()
    //     .cast(&DataType::Date)?;
    let real_date_col = df
        .column("real_date")
        .unwrap()
        // .str()?
        .cast(&DataType::String)?
        .str()?
        .as_date(Some(date_format), false)
        .map_err(|e| format!("Failed to parse date with format {}: {}", date_format, e))?;

    df.with_column(real_date_col)?;
    df.with_column(Series::new("cid".into(), vec![cid; df.height()]))?;
    df.with_column(Series::new("xcat".into(), vec![xcat; df.height()]))?;

    sort_qdf_columns(&mut df)?;

    Ok(df)
}

fn collect_paths_recursively<P: AsRef<std::path::Path>>(path: P) -> std::io::Result<Vec<String>> {
    let mut paths = Vec::new();

    for entry in std::fs::read_dir(path)? {
        let entry = entry?;
        let path = entry.path();

        if path.is_dir() {
            // Recurse into the directory and append the results
            paths.extend(collect_paths_recursively(&path)?);
        }
        // Add the path to the vector
        paths.push(path.to_string_lossy().to_string());
    }

    Ok(paths)
}

fn _load_qdf_thread_safe(file_path: &str) -> Result<DataFrame, Box<dyn Error + Send + Sync>> {
    let res = load_quantamental_dataframe(file_path.to_string());
    res.map_err(|e| {
        anyhow::Error::msg(e.to_string())
            .context("Failed to load quantamental dataframe")
            .into()
    })
}
pub fn load_qdf_from_download_bank(
    folder_path: String,
    cids: Option<Vec<String>>,
    xcats: Option<Vec<String>>,
    tickers: Option<Vec<String>>,
) -> Result<DataFrame, Box<dyn std::error::Error>> {
    let rcids = cids.unwrap_or_else(|| Vec::new());
    let rxcats = xcats.unwrap_or_else(|| Vec::new());
    let rtickers = tickers.unwrap_or_else(|| Vec::new());

    let files = collect_paths_recursively(folder_path)?;
    log::info!("Found {} files", files.len());

    // filter files that are not csv files
    let files = files
        .iter()
        .filter(|f| f.ends_with(".csv"))
        .collect::<Vec<&String>>();

    log::info!("Found {} csv files", files.len());

    let mut rel_files = Vec::new();
    for file in files {
        let base_file_name: String = _file_base_name(file.into())
            .split('.')
            .collect::<Vec<&str>>()[0]
            .into();
        let (cid, xcat) = match split_ticker(base_file_name.clone()) {
            Ok((cid, xcat)) => (cid, xcat),
            Err(_) => continue,
        };
        rel_files.push((file, cid, xcat));
    }

    log::info!("Found {} relevant ticker files", rel_files.len());

    let load_files = rel_files
        .iter()
        .filter(|(_, cid, xcat)| {
            let f1 = rcids.len() > 0 && rcids.contains(&cid);
            let f2 = rxcats.len() > 0 && rxcats.contains(&xcat);
            let f3 = rtickers.len() > 0 && rtickers.contains(&create_ticker(cid, xcat));
            f1 | f2 | f3
        })
        .map(|(file, _, _)| *file)
        .collect::<Vec<&String>>();

    // print number of files to load
    log::info!("Loading {} files", load_files.len());

    if load_files.len() == 0 {
        return Err("No files to load".into());
    }
    if load_files.len() == 1 {
        let dfx = load_quantamental_dataframe(load_files[0].to_string()).unwrap();
        return Ok(dfx);
    }

    let load_files = load_files.iter().map(|s| s.as_str()).collect::<Vec<&str>>();
    let qdf_batches = load_files
        .chunks(CONCURRENT_FILE_LOADS)
        .collect::<Vec<&[&str]>>();

    let mut results = Vec::new();
    let mut curr_batch = 1;
    let total_batches = qdf_batches.len();
    for batch in qdf_batches {
        let qdf_list = batch
            .par_iter()
            .map(|file| _load_qdf_thread_safe(file).unwrap().lazy())
            .collect::<Vec<LazyFrame>>();
        results.extend(qdf_list);
        curr_batch += 1;
        log::info!("Loaded {}/{} batches", curr_batch, total_batches);
    }

    log::info!("Loaded {} files", results.len());
    let res_df = concat(results, UnionArgs::default())
        .unwrap()
        .collect()
        .unwrap();

    log::info!(
        "Loaded dataframe with {} rows and {} columns",
        res_df.height(),
        res_df.width()
    );

    Ok(res_df)
}
