//! 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::stochastic;
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))
}

/// Stochastic simulation via SDE expression DSL.
///
/// `sim_config_json` must be a JSON-serialised config. Example:
/// ```json
/// {
///   "model": {
///     "name": "garch_jd",
///     "drift": "mu",
///     "diffusion": "sqrt(h)",
///     "jump_intensity": "lambda",
///     "jump_size": "normal(mu_j, sigma_j)",
///     "state_vars": {"h": 1e-4},
///     "state_update": {"h": "omega + alpha * (ret - mu) ** 2 + beta * h"},
///     "params": {"mu": 0.08, "omega": 1e-6, "alpha": 0.1, "beta": 0.85,
///                "lambda": 5.0, "mu_j": -0.02, "sigma_j": 0.04}
///   },
///   "n_paths": 10000,
///   "n_steps": 252,
///   "dt": 0.003968,
///   "s0": 100.0,
///   "confidence_levels": [0.05, 0.25, 0.50, 0.75, 0.95],
///   "rng_seed": 42,
///   "store_paths": false
/// }
/// ```
///
/// For presets, pass `"preset": "gbm"` instead of full model definition.
#[pyfunction]
#[pyo3(signature = (sim_config_json,))]
pub fn py_run_stochastic(py: Python<'_>, sim_config_json: &str) -> PyResult<PyObject> {
    // Deserialise the combined config (model def + sim params).
    let raw: serde_json::Value =
        serde_json::from_str(sim_config_json).map_err(super::map_json_error)?;

    // Resolve the model definition: either a preset name or a full SdeModelDef.
    let model_def = resolve_model_def(&raw).map_err(super::map_bt_error)?;

    // Parse simulation config fields.
    let sim_config = stochastic::StochasticSimConfig {
        n_paths: raw["n_paths"].as_u64().unwrap_or(1000) as usize,
        n_steps: raw["n_steps"].as_u64().unwrap_or(252) as usize,
        dt: raw["dt"].as_f64().unwrap_or(1.0 / 252.0),
        s0: raw["s0"].as_f64().unwrap_or(100.0),
        confidence_levels: raw
            .get("confidence_levels")
            .and_then(|v| serde_json::from_value(v.clone()).ok())
            .unwrap_or_else(|| vec![0.05, 0.10, 0.25, 0.50, 0.75, 0.90, 0.95]),
        rng_seed: raw.get("rng_seed").and_then(|v| v.as_u64()),
        store_paths: raw
            .get("store_paths")
            .and_then(|v| v.as_bool())
            .unwrap_or(false),
        precision: raw
            .get("precision")
            .and_then(|v| v.as_str())
            .unwrap_or("f64")
            .to_string(),
    };

    // Compile SDE expressions.
    let program = stochastic::compile_sde(&model_def).map_err(super::map_bt_error)?;

    // Device dispatch: "cuda" → GPU, anything else → CPU (Rayon).
    let device = raw
        .get("device")
        .and_then(|v| v.as_str())
        .unwrap_or("cpu");

    let result = match device {
        #[cfg(feature = "cuda")]
        "cuda" | "gpu" => {
            py.allow_threads(|| stochastic::run_stochastic_gpu(&program, &sim_config))
        }
        #[cfg(not(feature = "cuda"))]
        "cuda" | "gpu" => {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "CUDA support not compiled. Rebuild with: maturin develop --release --features cuda",
            ));
        }
        _ => py.allow_threads(|| stochastic::run_stochastic(&program, &sim_config)),
    };
    let mut result = result.map_err(super::map_bt_error)?;

    // Set model name.
    result.model_name = model_def.name.clone();

    // Extract paths before JSON serialisation.
    let paths_data = result.paths.take();

    let json = serde_json::to_string(&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);
            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())
}

/// Resolve model definition from JSON: supports preset name or full def.
fn resolve_model_def(raw: &serde_json::Value) -> bt_kernel::Result<stochastic::SdeModelDef> {
    // Check for preset: {"preset": "gbm", "params": {...}}
    if let Some(preset_name) = raw.get("preset").and_then(|v| v.as_str()) {
        let mut def = stochastic::preset(preset_name)?;
        // Override params if provided at top level
        if let Some(params) = raw.get("params").and_then(|v| v.as_object()) {
            for (k, v) in params {
                if let Some(f) = v.as_f64() {
                    def.params.insert(k.clone(), f);
                }
            }
        }
        return Ok(def);
    }

    // Check for model object: {"model": {...}}
    if let Some(model) = raw.get("model") {
        // Check if model is a preset string
        if let Some(preset_name) = model.get("preset").and_then(|v| v.as_str()) {
            let mut def = stochastic::preset(preset_name)?;
            if let Some(params) = model.get("params").and_then(|v| v.as_object()) {
                for (k, v) in params {
                    if let Some(f) = v.as_f64() {
                        def.params.insert(k.clone(), f);
                    }
                }
            }
            return Ok(def);
        }
        // Full model definition
        let def: stochastic::SdeModelDef = serde_json::from_value(model.clone())
            .map_err(|e| bt_kernel::BtError::Strategy(format!("invalid model definition: {e}")))?;
        return Ok(def);
    }

    // Try parsing the whole thing as an SdeModelDef
    let def: stochastic::SdeModelDef = serde_json::from_value(raw.clone())
        .map_err(|e| bt_kernel::BtError::Strategy(format!("invalid stochastic config: {e}")))?;
    Ok(def)
}
