use bt_analytics::PerformanceMetrics;
use bt_core::{
    BacktestConfig, BacktestRunner, BatchResultLite, ParamGrid, RunManifest, SimpleBacktestRunner,
};
use bt_strategy::{compile_strategy, StrategyDef};
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use serde::Serialize;

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

// ---------------------------------------------------------------------------
// Legacy JSON-only endpoint (kept for backward compat)
// ---------------------------------------------------------------------------

#[derive(Debug, Serialize)]
struct RunSummary {
    manifest: RunManifest,
    metrics: PerformanceMetrics,
    trade_count: usize,
    warning_count: usize,
}

#[pyfunction]
pub fn run_json(
    py: Python<'_>,
    strategy_json: &str,
    config_json: &str,
    store: &PyDataStore,
) -> PyResult<String> {
    let strategy: StrategyDef =
        serde_json::from_str(strategy_json).map_err(super::map_json_error)?;
    let config: BacktestConfig =
        serde_json::from_str(config_json).map_err(super::map_json_error)?;
    let compiled = compile_strategy(&strategy).map_err(super::map_bt_error)?;

    let data_store = store.store().clone();

    let result = py.allow_threads(|| {
        let runner = SimpleBacktestRunner::new();
        runner.run(&compiled, &config, data_store.as_ref())
    });
    let result = result.map_err(super::map_bt_error)?;

    let summary = RunSummary {
        manifest: result.manifest,
        metrics: result.metrics,
        trade_count: result.trades.num_rows(),
        warning_count: result.warnings.len(),
    };

    serde_json::to_string_pretty(&summary).map_err(super::map_json_error)
}

// ---------------------------------------------------------------------------
// Arrow-native endpoint (returns PyBacktestResult with zero-copy data)
// ---------------------------------------------------------------------------

#[pyfunction]
pub fn run(
    py: Python<'_>,
    strategy_json: &str,
    config_json: &str,
    store: &PyDataStore,
) -> PyResult<PyBacktestResult> {
    let strategy: StrategyDef =
        serde_json::from_str(strategy_json).map_err(super::map_json_error)?;
    let config: BacktestConfig =
        serde_json::from_str(config_json).map_err(super::map_json_error)?;
    let compiled = compile_strategy(&strategy).map_err(super::map_bt_error)?;

    let data_store = store.store().clone();

    let result = py.allow_threads(|| {
        let runner = SimpleBacktestRunner::new();
        runner.run(&compiled, &config, data_store.as_ref())
    });
    let result = result.map_err(super::map_bt_error)?;

    Ok(PyBacktestResult::new(result))
}

// ---------------------------------------------------------------------------
// Test helper: run from a parquet file (no metadata DB needed)
// ---------------------------------------------------------------------------

#[pyfunction]
pub fn run_with_parquet(
    py: Python<'_>,
    strategy_json: &str,
    config_json: &str,
    parquet_path: &str,
    version_id: &str,
) -> PyResult<PyBacktestResult> {
    use bt_kernel::concat_batches;
    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
    use std::fs::File;

    let file = File::open(parquet_path)
        .map_err(|e| PyValueError::new_err(format!("cannot open parquet: {e}")))?;
    let reader = ParquetRecordBatchReaderBuilder::try_new(file)
        .map_err(|e| PyValueError::new_err(format!("invalid parquet: {e}")))?
        .build()
        .map_err(|e| PyValueError::new_err(format!("reader build: {e}")))?;

    let mut batches = Vec::new();
    for batch in reader {
        batches.push(batch.map_err(|e| PyValueError::new_err(format!("read batch: {e}")))?);
    }

    if batches.is_empty() {
        return Err(PyValueError::new_err("parquet file is empty"));
    }

    let bars = if batches.len() == 1 {
        batches.remove(0)
    } else {
        let schema = batches[0].schema();
        concat_batches(schema, &batches).map_err(super::map_bt_error)?
    };

    let store = InMemoryStore::new(bars, version_id);

    let strategy: StrategyDef =
        serde_json::from_str(strategy_json).map_err(super::map_json_error)?;
    let config: BacktestConfig =
        serde_json::from_str(config_json).map_err(super::map_json_error)?;
    let compiled = compile_strategy(&strategy).map_err(super::map_bt_error)?;

    let result = py.allow_threads(|| {
        let runner = SimpleBacktestRunner::new();
        runner.run(&compiled, &config, &store)
    });
    let result = result.map_err(super::map_bt_error)?;

    Ok(PyBacktestResult::new(result))
}

/// Load a parquet file into AlignedData (no backtest, no DataStore needed).
/// Use with `run_on_aligned()` to benchmark the engine without I/O.
#[pyfunction]
pub fn load_parquet_as_aligned(
    py: Python<'_>,
    config_json: &str,
    parquet_path: &str,
    version_id: &str,
) -> PyResult<PyAlignedData> {
    use bt_kernel::concat_batches;
    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
    use std::fs::File;

    let file = File::open(parquet_path)
        .map_err(|e| PyValueError::new_err(format!("cannot open parquet: {e}")))?;
    let reader = ParquetRecordBatchReaderBuilder::try_new(file)
        .map_err(|e| PyValueError::new_err(format!("invalid parquet: {e}")))?
        .build()
        .map_err(|e| PyValueError::new_err(format!("reader build: {e}")))?;

    let mut batches = Vec::new();
    for batch in reader {
        batches.push(batch.map_err(|e| PyValueError::new_err(format!("read batch: {e}")))?);
    }

    if batches.is_empty() {
        return Err(PyValueError::new_err("parquet file is empty"));
    }

    let bars = if batches.len() == 1 {
        batches.remove(0)
    } else {
        let schema = batches[0].schema();
        concat_batches(schema, &batches).map_err(super::map_bt_error)?
    };

    let store = InMemoryStore::new(bars, version_id);
    let config: BacktestConfig =
        serde_json::from_str(config_json).map_err(super::map_json_error)?;

    let (aligned, universe, version, _, _) = py
        .allow_threads(|| SimpleBacktestRunner::load_and_align(&config, &store))
        .map_err(super::map_bt_error)?;

    Ok(PyAlignedData {
        aligned: Arc::new(aligned),
        universe,
        data_version: version,
    })
}

// ---------------------------------------------------------------------------
// Parameter sweep (parallel via rayon, GIL released)
// ---------------------------------------------------------------------------

#[pyfunction]
#[pyo3(signature = (strategy_json, param_grid_json, config_json, store, max_parallelism=0))]
pub fn run_sweep(
    py: Python<'_>,
    strategy_json: &str,
    param_grid_json: &str,
    config_json: &str,
    store: &PyDataStore,
    max_parallelism: usize,
) -> PyResult<Vec<PyBacktestResult>> {
    let strategy: StrategyDef =
        serde_json::from_str(strategy_json).map_err(super::map_json_error)?;
    let param_grid: ParamGrid =
        serde_json::from_str(param_grid_json).map_err(super::map_json_error)?;
    let config: BacktestConfig =
        serde_json::from_str(config_json).map_err(super::map_json_error)?;
    let compiled = compile_strategy(&strategy).map_err(super::map_bt_error)?;

    let data_store = store.store().clone();

    let threads = if max_parallelism == 0 {
        num_cpus()
    } else {
        max_parallelism
    };

    let results = py.allow_threads(|| {
        let runner = SimpleBacktestRunner::new();
        runner.run_sweep(
            &compiled,
            &param_grid,
            &config,
            data_store.as_ref(),
            threads,
        )
    });
    let results = results.map_err(super::map_bt_error)?;

    Ok(results.into_iter().map(PyBacktestResult::new).collect())
}

// ---------------------------------------------------------------------------
// Sweep lite: parameter grid → metrics only (fastest sweep path)
// ---------------------------------------------------------------------------

#[pyfunction]
#[pyo3(signature = (strategy_json, param_grid_json, config_json, store, max_parallelism=0))]
pub fn run_sweep_lite(
    py: Python<'_>,
    strategy_json: &str,
    param_grid_json: &str,
    config_json: &str,
    store: &PyDataStore,
    max_parallelism: usize,
) -> PyResult<Vec<PyBatchResultLite>> {
    let strategy: StrategyDef =
        serde_json::from_str(strategy_json).map_err(super::map_json_error)?;
    let param_grid: ParamGrid =
        serde_json::from_str(param_grid_json).map_err(super::map_json_error)?;
    let config: BacktestConfig =
        serde_json::from_str(config_json).map_err(super::map_json_error)?;
    let compiled = compile_strategy(&strategy).map_err(super::map_bt_error)?;

    let data_store = store.store().clone();

    let threads = if max_parallelism == 0 {
        num_cpus()
    } else {
        max_parallelism
    };

    let results = py.allow_threads(|| {
        let runner = SimpleBacktestRunner::new();
        runner.run_sweep_lite(
            &compiled,
            &param_grid,
            &config,
            data_store.as_ref(),
            threads,
        )
    });
    let results = results.map_err(super::map_bt_error)?;

    Ok(results
        .into_iter()
        .map(|r| PyBatchResultLite { inner: r })
        .collect())
}

// ---------------------------------------------------------------------------
// Batch run: many different strategies, shared data loading (rayon)
// ---------------------------------------------------------------------------

#[pyfunction]
#[pyo3(signature = (strategy_jsons, config_json, store, max_parallelism=0))]
pub fn run_batch(
    py: Python<'_>,
    strategy_jsons: Vec<String>,
    config_json: &str,
    store: &PyDataStore,
    max_parallelism: usize,
) -> PyResult<Vec<PyBacktestResult>> {
    let config: BacktestConfig =
        serde_json::from_str(config_json).map_err(super::map_json_error)?;

    let compiled: Vec<_> = strategy_jsons
        .iter()
        .map(|json| {
            let def: StrategyDef = serde_json::from_str(json).map_err(super::map_json_error)?;
            compile_strategy(&def).map_err(super::map_bt_error)
        })
        .collect::<PyResult<_>>()?;

    let data_store = store.store().clone();

    let threads = if max_parallelism == 0 {
        num_cpus()
    } else {
        max_parallelism
    };

    let results = py.allow_threads(|| {
        let runner = SimpleBacktestRunner::new();
        runner.run_batch(&compiled, &config, data_store.as_ref(), threads)
    });
    let results = results.map_err(super::map_bt_error)?;

    Ok(results.into_iter().map(PyBacktestResult::new).collect())
}

// ---------------------------------------------------------------------------
// Batch lite: metrics only, no Arrow output (much less memory)
// ---------------------------------------------------------------------------

#[pyclass(name = "BatchResultLite")]
#[derive(Clone)]
pub struct PyBatchResultLite {
    inner: BatchResultLite,
}

#[pymethods]
impl PyBatchResultLite {
    #[getter]
    fn strategy_name(&self) -> &str {
        &self.inner.strategy_name
    }

    #[getter]
    fn final_equity(&self) -> f64 {
        self.inner.final_equity
    }

    #[getter]
    fn trade_count(&self) -> usize {
        self.inner.trade_count
    }

    #[getter]
    fn metrics(&self) -> PyResult<PyObject> {
        Python::with_gil(|py| {
            let json = serde_json::to_string(&self.inner.metrics).map_err(super::map_json_error)?;
            let dict: PyObject = pyo3::types::PyModule::import(py, "json")?
                .call_method1("loads", (json,))?
                .into();
            Ok(dict)
        })
    }

    #[getter]
    fn profile(&self) -> PyResult<PyObject> {
        Python::with_gil(|py| {
            let json = serde_json::to_string(&self.inner.profile).map_err(super::map_json_error)?;
            let dict: PyObject = pyo3::types::PyModule::import(py, "json")?
                .call_method1("loads", (json,))?
                .into();
            Ok(dict)
        })
    }

    fn __repr__(&self) -> String {
        format!(
            "BatchResultLite(name={}, equity={:.2}, trades={})",
            self.inner.strategy_name, self.inner.final_equity, self.inner.trade_count
        )
    }
}

#[pyfunction]
#[pyo3(signature = (strategy_jsons, config_json, store, max_parallelism=0))]
pub fn run_batch_lite(
    py: Python<'_>,
    strategy_jsons: Vec<String>,
    config_json: &str,
    store: &PyDataStore,
    max_parallelism: usize,
) -> PyResult<Vec<PyBatchResultLite>> {
    let config: BacktestConfig =
        serde_json::from_str(config_json).map_err(super::map_json_error)?;

    let compiled: Vec<_> = strategy_jsons
        .iter()
        .map(|json| {
            let def: StrategyDef = serde_json::from_str(json).map_err(super::map_json_error)?;
            compile_strategy(&def).map_err(super::map_bt_error)
        })
        .collect::<PyResult<_>>()?;

    let data_store = store.store().clone();

    let threads = if max_parallelism == 0 {
        num_cpus()
    } else {
        max_parallelism
    };

    let results = py.allow_threads(|| {
        let runner = SimpleBacktestRunner::new();
        runner.run_batch_lite(&compiled, &config, data_store.as_ref(), threads)
    });
    let results = results.map_err(super::map_bt_error)?;

    Ok(results
        .into_iter()
        .map(|r| PyBatchResultLite { inner: r })
        .collect())
}

fn num_cpus() -> usize {
    std::thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(4)
}

// ---------------------------------------------------------------------------
// Pre-loaded aligned data (load once, run many times)
// ---------------------------------------------------------------------------

use bt_core::alignment::{slice_aligned_bars, AlignedBars};
use bt_kernel::{SymbolId, TimeRange, Timestamp};
use std::sync::Arc;

/// Opaque handle to pre-loaded and aligned bar data.
/// Allows loading data once and running multiple backtests without reloading.
#[pyclass(name = "AlignedData")]
pub struct PyAlignedData {
    aligned: Arc<AlignedBars>,
    universe: Vec<SymbolId>,
    data_version: String,
}

#[pymethods]
impl PyAlignedData {
    /// Number of bars (master timestamps).
    #[getter]
    fn num_bars(&self) -> usize {
        self.aligned.master_timestamps.len()
    }

    /// Number of symbols.
    #[getter]
    fn num_symbols(&self) -> usize {
        self.universe.len()
    }

    /// Slice to a sub-range [start_ns, end_ns). Returns a new AlignedData.
    fn slice(&self, start_ns: i64, end_ns: i64) -> PyResult<PyAlignedData> {
        let range =
            TimeRange::new(Timestamp(start_ns), Timestamp(end_ns)).map_err(super::map_bt_error)?;
        let sliced = slice_aligned_bars(&self.aligned, range).map_err(super::map_bt_error)?;
        Ok(PyAlignedData {
            aligned: Arc::new(sliced),
            universe: self.universe.clone(),
            data_version: self.data_version.clone(),
        })
    }

    fn __repr__(&self) -> String {
        format!(
            "AlignedData(bars={}, symbols={})",
            self.num_bars(),
            self.num_symbols()
        )
    }
}

/// Load and align bar data (phases 1-2 only). Returns an opaque handle
/// that can be passed to `run_on_aligned()` multiple times.
#[pyfunction]
pub fn load_and_align(
    py: Python<'_>,
    config_json: &str,
    store: &PyDataStore,
) -> PyResult<PyAlignedData> {
    let config: BacktestConfig =
        serde_json::from_str(config_json).map_err(super::map_json_error)?;
    let data_store = store.store().clone();

    let (aligned, universe, version, _, _) = py
        .allow_threads(|| SimpleBacktestRunner::load_and_align(&config, data_store.as_ref()))
        .map_err(super::map_bt_error)?;

    Ok(PyAlignedData {
        aligned: Arc::new(aligned),
        universe,
        data_version: version,
    })
}

/// Run a backtest on pre-loaded aligned data (phases 3-6 only).
/// Skips disk I/O entirely — useful for diagnostics, sweeps, etc.
#[pyfunction]
pub fn run_on_aligned(
    py: Python<'_>,
    strategy_json: &str,
    config_json: &str,
    aligned: &PyAlignedData,
) -> PyResult<PyBacktestResult> {
    let strategy: StrategyDef =
        serde_json::from_str(strategy_json).map_err(super::map_json_error)?;
    let config: BacktestConfig =
        serde_json::from_str(config_json).map_err(super::map_json_error)?;
    let compiled = compile_strategy(&strategy).map_err(super::map_bt_error)?;

    let aligned_ref = aligned.aligned.clone();
    let universe = aligned.universe.clone();
    let version = aligned.data_version.clone();

    let result = py
        .allow_threads(move || {
            let runner = SimpleBacktestRunner::new();
            runner.run_on_aligned(&compiled, &config, &aligned_ref, &universe, &version)
        })
        .map_err(super::map_bt_error)?;

    Ok(PyBacktestResult::new(result))
}
