use pyo3::prelude::*;
use pyo3_polars::PyDataFrame;

/// Python wrapper for [`crate::utils::qdf`] module.
#[allow(deprecated)]
#[pymodule]
pub fn qdf(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(load_qdf, m)?)?;
    m.add_function(wrap_pyfunction!(load_qdf_from_download_bank, m)?)?;
    m.add_function(wrap_pyfunction!(reduce_dataframe, m)?)?;
    m.add_function(wrap_pyfunction!(update_dataframe, m)?)?;
    m.add_function(wrap_pyfunction!(pivot_dataframe_by_ticker, m)?)?;
    Ok(())
}

/// Python wrapper for loading a Quantamental DataFrame from a CSV file.
/// See [`crate::utils::qdf::load_quantamental_dataframe`] for full documentation.
#[pyfunction]
pub fn load_qdf(file_path: String) -> PyResult<PyDataFrame> {
    Ok(PyDataFrame(
        crate::utils::qdf::load_quantamental_dataframe(file_path).unwrap(),
    ))
}

/// Python wrapper for loading a Quantamental DataFrame from a download bank.
/// See [`crate::utils::qdf::load::load_qdf_from_download_bank`] for full documentation.
#[pyfunction]
pub fn load_qdf_from_download_bank(
    folder_path: String,
    cids: Option<Vec<String>>,
    xcats: Option<Vec<String>>,
    tickers: Option<Vec<String>>,
) -> PyResult<PyDataFrame> {
    Ok(PyDataFrame(
        crate::utils::qdf::load::load_qdf_from_download_bank(folder_path, cids, xcats, tickers)
            .unwrap(),
    ))
}

/// Python wrapper for reduce_dataframe
/// See [`crate::utils::qdf::reduce_df::reduce_dataframe`] for full documentation.
#[pyfunction]
pub fn reduce_dataframe(
    df: PyDataFrame,
    cids: Option<Vec<String>>,
    xcats: Option<Vec<String>>,
    metrics: Option<Vec<String>>,
    start: Option<String>,
    end: Option<String>,
    intersect: Option<bool>,
) -> PyResult<PyDataFrame> {
    Ok(PyDataFrame(
        crate::utils::qdf::reduce_df::reduce_dataframe(
            df.into(),
            cids,
            xcats,
            metrics,
            start,
            end,
            intersect.unwrap_or(false),
        )
        .unwrap(),
    ))
}

/// Python wrapper for update_dataframe
/// See [`crate::utils::qdf::update_df::update_dataframe`] for full documentation.
#[pyfunction]
pub fn update_dataframe(
    df: PyDataFrame,
    df_add: PyDataFrame,
    xcat_replace: Option<bool>,
) -> PyResult<PyDataFrame> {
    let xcat_replace = xcat_replace.unwrap_or(false);
    Ok(PyDataFrame(
        crate::utils::qdf::update_df::update_dataframe(&df.into(), &df_add.into(), xcat_replace)
            .unwrap(),
    ))
}

/// Python wrapper for pivoting a dataframe by ticker.
/// See [`crate::utils::qdf::pivots::pivot_dataframe_by_ticker`] for full documentation.
#[pyfunction]
pub fn pivot_dataframe_by_ticker(df: PyDataFrame, metric: Option<String>) -> PyResult<PyDataFrame> {
    Ok(PyDataFrame(
        crate::utils::qdf::pivots::pivot_dataframe_by_ticker(df.into(), metric).unwrap(),
    ))
}
