//! PyO3 bindings for bt-research functions.

use arrow::array::{Array as _, Float64Array};
use pyo3::prelude::*;
use pyo3::types::PyDict;

use bt_core::orchestrator::BacktestConfig;
use bt_research::monte_carlo::{run_monte_carlo, MonteCarloConfig};
use bt_research::stability::{run_stability, StabilityConfig};
use bt_research::sweep_2d::{run_sweep_2d, Sweep2dConfig};
use bt_research::walk_forward::{run_walk_forward, WalkForwardConfig};
use bt_strategy::compiler::compile_strategy;
use bt_strategy::compiler::StrategyDef;

use crate::data::PyDataStore;
use crate::result::PyBacktestResult;

#[pyfunction]
#[pyo3(signature = (strategy_json, wf_config_json, config_json, store))]
pub fn py_run_walk_forward(
    py: Python<'_>,
    strategy_json: &str,
    wf_config_json: &str,
    config_json: &str,
    store: &PyDataStore,
) -> PyResult<PyObject> {
    let strategy_def: StrategyDef =
        serde_json::from_str(strategy_json).map_err(super::map_json_error)?;
    let wf_config: WalkForwardConfig =
        serde_json::from_str(wf_config_json).map_err(super::map_json_error)?;
    let config: BacktestConfig =
        serde_json::from_str(config_json).map_err(super::map_json_error)?;
    let plan = compile_strategy(&strategy_def).map_err(super::map_bt_error)?;
    let data_store = store.store().clone();

    let result =
        py.allow_threads(|| run_walk_forward(&plan, &wf_config, &config, data_store.as_ref()));
    let result = result.map_err(super::map_bt_error)?;

    let json = serde_json::to_string(&result).map_err(super::map_json_error)?;
    let dict: PyObject = pyo3::types::PyModule::import(py, "json")?
        .call_method1("loads", (json,))?
        .into();
    Ok(dict)
}

#[pyfunction]
#[pyo3(signature = (strategy_json, sweep_config_json, config_json, store))]
pub fn py_run_sweep_2d(
    py: Python<'_>,
    strategy_json: &str,
    sweep_config_json: &str,
    config_json: &str,
    store: &PyDataStore,
) -> PyResult<PyObject> {
    let strategy_def: StrategyDef =
        serde_json::from_str(strategy_json).map_err(super::map_json_error)?;
    let sweep_config: Sweep2dConfig =
        serde_json::from_str(sweep_config_json).map_err(super::map_json_error)?;
    let config: BacktestConfig =
        serde_json::from_str(config_json).map_err(super::map_json_error)?;
    let plan = compile_strategy(&strategy_def).map_err(super::map_bt_error)?;
    let data_store = store.store().clone();

    let result =
        py.allow_threads(|| run_sweep_2d(&plan, &sweep_config, &config, data_store.as_ref()));
    let result = result.map_err(super::map_bt_error)?;

    let json = serde_json::to_string(&result).map_err(super::map_json_error)?;
    let dict: PyObject = pyo3::types::PyModule::import(py, "json")?
        .call_method1("loads", (json,))?
        .into();
    Ok(dict)
}

#[pyfunction]
#[pyo3(signature = (strategy_json, stability_config_json, config_json, store))]
pub fn py_run_stability(
    py: Python<'_>,
    strategy_json: &str,
    stability_config_json: &str,
    config_json: &str,
    store: &PyDataStore,
) -> PyResult<PyObject> {
    let strategy_def: StrategyDef =
        serde_json::from_str(strategy_json).map_err(super::map_json_error)?;
    let stability_config: StabilityConfig =
        serde_json::from_str(stability_config_json).map_err(super::map_json_error)?;
    let config: BacktestConfig =
        serde_json::from_str(config_json).map_err(super::map_json_error)?;
    let plan = compile_strategy(&strategy_def).map_err(super::map_bt_error)?;
    let data_store = store.store().clone();

    let result =
        py.allow_threads(|| run_stability(&plan, &stability_config, &config, data_store.as_ref()));
    let result = result.map_err(super::map_bt_error)?;

    let json = serde_json::to_string(&result).map_err(super::map_json_error)?;
    let dict: PyObject = pyo3::types::PyModule::import(py, "json")?
        .call_method1("loads", (json,))?
        .into();
    Ok(dict)
}

/// Monte Carlo simulation on an existing backtest result.
///
/// `mc_config_json` must be a JSON-serialised `MonteCarloConfig`. Example:
/// ```json
/// {
///   "n_paths": 1000,
///   "method": {"type": "return_bootstrap"},
///   "confidence_levels": [0.05, 0.25, 0.50, 0.75, 0.95],
///   "rng_seed": 42,
///   "initial_capital": 10000.0,
///   "store_paths": false
/// }
/// ```
/// For `block_bootstrap`:
/// ```json
/// {"type": "block_bootstrap", "block_size": 20}
/// ```
/// For `trade_bootstrap`:
/// ```json
/// {"type": "trade_bootstrap"}
/// ```
#[pyfunction]
#[pyo3(signature = (result, mc_config_json))]
pub fn py_run_monte_carlo(
    py: Python<'_>,
    result: &PyBacktestResult,
    mc_config_json: &str,
) -> PyResult<PyObject> {
    let config: MonteCarloConfig =
        serde_json::from_str(mc_config_json).map_err(super::map_json_error)?;

    // Extract data from result before entering allow_threads (requires GIL).
    let equity = result.equity_as_f64();
    let trade_pnls = result.trade_pnls();

    let mc_result = py.allow_threads(|| run_monte_carlo(&equity, &trade_pnls, &config));
    let mut mc_result = mc_result.map_err(super::map_bt_error)?;

    // Extract paths before JSON serialisation to avoid serialising millions of
    // floats through a text round-trip (slow). We return paths as a flat
    // PyArrow Float64Array + n_steps so the caller can reshape (numpy/arrow).
    let paths_data = mc_result.paths.take(); // sets field to None before JSON

    let json = serde_json::to_string(&mc_result).map_err(super::map_json_error)?;
    let dict = pyo3::types::PyModule::import(py, "json")?.call_method1("loads", (json,))?;
    let dict = dict.downcast::<PyDict>()?;

    match paths_data {
        None => {
            dict.set_item("paths", py.None())?;
            dict.set_item("paths_n_steps", py.None())?;
        }
        Some(paths) => {
            let n_steps = paths.first().map(|p| p.len()).unwrap_or(0);
            // Flatten row-major: paths[path_idx][step] → flat Vec<f64>
            let flat: Vec<f64> = paths.into_iter().flatten().collect();
            let arrow_arr = Float64Array::from(flat);
            use arrow::pyarrow::ToPyArrow;
            let py_arr = arrow_arr.into_data().to_pyarrow(py)?;
            dict.set_item("paths", py_arr)?;
            dict.set_item("paths_n_steps", n_steps)?;
        }
    }

    Ok(dict.clone().into_any().unbind())
}

#[pyfunction]
#[pyo3(signature = (manifest_json, strategy_json, store))]
pub fn py_replay(
    py: Python<'_>,
    manifest_json: &str,
    strategy_json: &str,
    store: &PyDataStore,
) -> PyResult<crate::result::PyBacktestResult> {
    let manifest: bt_core::RunManifest =
        serde_json::from_str(manifest_json).map_err(super::map_json_error)?;
    let strategy_def: StrategyDef =
        serde_json::from_str(strategy_json).map_err(super::map_json_error)?;
    let plan = compile_strategy(&strategy_def).map_err(super::map_bt_error)?;
    let data_store = store.store().clone();

    let result =
        py.allow_threads(|| bt_core::manifest::replay(&manifest, &plan, data_store.as_ref()));
    let result = result.map_err(super::map_bt_error)?;

    Ok(crate::result::PyBacktestResult::new(result))
}
