use std::collections::HashMap;
use std::sync::Arc;

use arrow::array::{
    ArrayRef, BooleanArray, Float64Array, TimestampNanosecondArray, UInt32Array, UInt8Array,
};
use arrow::record_batch::RecordBatch;
use bt_core::{
    BacktestConfig, BacktestRunner, BarInterval, ExecutionConfig, ExecutionPrice, FeeConfig,
    FillModel, PositionSizingMode, SimpleBacktestRunner, SlippageConfig, StrategyMetadata,
    StrategyPlan,
};
use bt_data::{DataStore, DatasetVersion, WriteOptions};
use bt_expr::{compile_expr, Expr, ScalarValue};
use bt_kernel::{BtError, Result, SymbolId, TimeRange, Timestamp};
use criterion::{black_box, criterion_group, criterion_main, Criterion};

struct BenchStore {
    bars: RecordBatch,
    version: DatasetVersion,
}

impl BenchStore {
    fn new(rows: usize) -> Self {
        let bars = synthetic_bars(rows);
        Self {
            bars,
            version: DatasetVersion {
                id: "bench_v1".to_string(),
                dataset: "bars_1m".to_string(),
                created_at: Timestamp::from_nanos(1),
                schema_hash: "bench".to_string(),
                row_count: rows as u64,
                time_range: TimeRange::new(
                    Timestamp::from_nanos(0),
                    Timestamp::from_nanos(rows as i64 * 60_000_000_000),
                )
                .expect("valid range"),
                symbols: vec![SymbolId(1)],
                source_hash: "bench".to_string(),
                metadata: serde_json::json!({}),
            },
        }
    }
}

impl DataStore for BenchStore {
    fn load_bars(
        &self,
        _symbol: SymbolId,
        _range: TimeRange,
        _version: Option<&str>,
    ) -> Result<RecordBatch> {
        Ok(self.bars.clone())
    }

    fn load_features(
        &self,
        _feature_set: &str,
        _symbol: SymbolId,
        _range: TimeRange,
        _columns: Option<&[&str]>,
    ) -> Result<RecordBatch> {
        Err(BtError::Unsupported("benchmark store".to_string()))
    }

    fn load_exo(&self, _name: &str, _range: TimeRange) -> Result<RecordBatch> {
        Err(BtError::Unsupported("benchmark store".to_string()))
    }

    fn list_versions(&self, _dataset: &str) -> Result<Vec<DatasetVersion>> {
        Ok(vec![self.version.clone()])
    }

    fn active_version(&self, _dataset: &str) -> Result<DatasetVersion> {
        Ok(self.version.clone())
    }

    fn write_bars(
        &self,
        _symbol: SymbolId,
        _batch: &RecordBatch,
        _opts: WriteOptions,
    ) -> Result<DatasetVersion> {
        Err(BtError::Unsupported("benchmark store".to_string()))
    }
}

fn benchmark_single_asset_runner(c: &mut Criterion) {
    let rows = 50_000;
    let store = BenchStore::new(rows);
    let runner = SimpleBacktestRunner::new();
    let strategy = benchmark_strategy();
    let config = BacktestConfig {
        universe: vec![SymbolId(1)],
        time_range: TimeRange::new(
            Timestamp::from_nanos(0),
            Timestamp::from_nanos(rows as i64 * 60_000_000_000),
        )
        .expect("valid range"),
        bar_interval: BarInterval::Minutes(1),
        initial_capital: 100_000.0,
        currency: "USD".to_string(),
        risk_free_rate: 0.025,
        execution: ExecutionConfig {
            signal_delay: 1,
            execution_price: ExecutionPrice::AtClose,
            max_position_pct: 1.0,
            allow_short: true,
            allow_fractional: true,
            skip_gap_bars: false,
            position_sizing_mode: PositionSizingMode::Units,
            fill_model: FillModel::default(),
            ..Default::default()
        },
        fees: FeeConfig {
            taker_fee_bps: 1.0,
            ..FeeConfig::default()
        },
        slippage: SlippageConfig::FixedBps { bps: 1.0 },
        data_version: None,
        rng_seed: Some(42),
        trading_days_per_year: 365.25,
        output_resolution: None,
        resample_to: None,
        feature_sets: vec![],
        ..Default::default()
    };

    let mut group = c.benchmark_group("runner");
    group.bench_function("single_asset_50k_rows_target_lt_500ms", |b| {
        b.iter(|| {
            let result = runner
                .run(&strategy, &config, &store)
                .expect("benchmark run");
            black_box(result.metrics.total_return)
        })
    });
    group.finish();
}

fn benchmark_sweep_1000_params(c: &mut Criterion) {
    let rows = 5_000;
    let store = BenchStore::new(rows);
    let runner = SimpleBacktestRunner::new();

    let signal = compile_expr(Expr::Literal(ScalarValue::Float64(1.0))).expect("compile signal");
    let sizing = compile_expr(Expr::Mul(
        Box::new(Expr::Column("signal".to_string())),
        Box::new(Expr::Parameter("size".to_string())),
    ))
    .expect("compile sizing");

    let mut signals = HashMap::new();
    signals.insert("signal".to_string(), signal);

    let strategy = StrategyPlan {
        name: "sweep_bench".to_string(),
        signals,
        position_sizing: sizing,
        required_columns: vec!["close".to_string()],
        required_features: Vec::new(),
        parameters: HashMap::from([("size".to_string(), ScalarValue::Float64(1.0))]),
        metadata: StrategyMetadata::default(),
        raw_signals: None,
        raw_position_sizing: None,
        has_dynamic_periods: false,
    };

    let config = BacktestConfig {
        universe: vec![SymbolId(1)],
        time_range: TimeRange::new(
            Timestamp::from_nanos(0),
            Timestamp::from_nanos(rows as i64 * 60_000_000_000),
        )
        .expect("valid range"),
        bar_interval: BarInterval::Minutes(1),
        initial_capital: 100_000.0,
        currency: "USD".to_string(),
        risk_free_rate: 0.025,
        execution: ExecutionConfig::default(),
        fees: FeeConfig::default(),
        slippage: SlippageConfig::FixedBps { bps: 0.0 },
        data_version: None,
        rng_seed: Some(42),
        trading_days_per_year: 365.25,
        output_resolution: None,
        resample_to: None,
        feature_sets: vec![],
        ..Default::default()
    };

    // 1000 combos: 10 values x 100 values (or just 1000 values for one param)
    let size_values: Vec<ScalarValue> = (1..=1000)
        .map(|i| ScalarValue::Float64(i as f64 * 0.001))
        .collect();
    let param_grid: bt_core::ParamGrid = HashMap::from([("size".to_string(), size_values)]);

    let max_threads = std::thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(4);

    let mut group = c.benchmark_group("sweep");
    group.sample_size(10);
    group.bench_function(
        format!("1000_params_{rows}_rows_{max_threads}_threads"),
        |b| {
            b.iter(|| {
                let results = runner
                    .run_sweep(&strategy, &param_grid, &config, &store, max_threads)
                    .expect("sweep run");
                black_box(results.len())
            })
        },
    );
    group.finish();
}

fn benchmark_strategy() -> StrategyPlan {
    let signal = compile_expr(Expr::Literal(ScalarValue::Float64(1.0))).expect("compile signal");
    let sizing = compile_expr(Expr::Mul(
        Box::new(Expr::Column("signal".to_string())),
        Box::new(Expr::Literal(ScalarValue::Float64(0.5))),
    ))
    .expect("compile sizing");

    let mut signals = HashMap::new();
    signals.insert("signal".to_string(), signal);

    StrategyPlan {
        name: "benchmark_buy_half".to_string(),
        signals,
        position_sizing: sizing,
        required_columns: vec!["close".to_string()],
        required_features: Vec::new(),
        parameters: HashMap::new(),
        metadata: StrategyMetadata::default(),
        raw_signals: None,
        raw_position_sizing: None,
        has_dynamic_periods: false,
    }
}

fn synthetic_bars(rows: usize) -> RecordBatch {
    let schema = bt_data::schema::canonical_bars_1m_schema();
    let mut timestamps = Vec::with_capacity(rows);
    let mut prices = Vec::with_capacity(rows);
    let mut volume = Vec::with_capacity(rows);
    for idx in 0..rows {
        timestamps.push(idx as i64 * 60_000_000_000);
        prices.push(100.0 + (idx as f64 * 0.0001));
        volume.push(100.0 + (idx % 25) as f64);
    }

    let arrays: Vec<ArrayRef> = vec![
        Arc::new(TimestampNanosecondArray::from(timestamps).with_timezone("UTC")),
        Arc::new(UInt32Array::from(vec![1_u32; rows])),
        Arc::new(Float64Array::from(prices.clone())),
        Arc::new(Float64Array::from(prices.clone())),
        Arc::new(Float64Array::from(prices.clone())),
        Arc::new(Float64Array::from(prices.clone())),
        Arc::new(Float64Array::from(prices)),
        Arc::new(Float64Array::from(volume)),
        Arc::new(Float64Array::from(vec![Some(50.0); rows])),
        Arc::new(Float64Array::from(vec![Some(50.0); rows])),
        Arc::new(UInt32Array::from(vec![1_u32; rows])),
        Arc::new(Float64Array::from(vec![Some(99.9); rows])),
        Arc::new(Float64Array::from(vec![Some(100.1); rows])),
        Arc::new(Float64Array::from(vec![Some(0.2); rows])),
        Arc::new(BooleanArray::from(vec![false; rows])),
        Arc::new(UInt8Array::from(vec![0_u8; rows])),
    ];

    RecordBatch::try_new(schema, arrays).expect("valid benchmark batch")
}

struct MultiBenchStore {
    bars: HashMap<SymbolId, RecordBatch>,
    version: DatasetVersion,
}

impl MultiBenchStore {
    fn new(num_symbols: usize, rows: usize) -> Self {
        let mut bars = HashMap::new();
        for sid in 1..=num_symbols {
            bars.insert(
                SymbolId(sid as u32),
                synthetic_bars_for_symbol(sid as u32, rows),
            );
        }
        Self {
            bars,
            version: DatasetVersion {
                id: "bench_multi_v1".to_string(),
                dataset: "bars_1m".to_string(),
                created_at: Timestamp::from_nanos(1),
                schema_hash: "bench".to_string(),
                row_count: rows as u64,
                time_range: TimeRange::new(
                    Timestamp::from_nanos(0),
                    Timestamp::from_nanos(rows as i64 * 60_000_000_000),
                )
                .expect("valid range"),
                symbols: (1..=num_symbols).map(|s| SymbolId(s as u32)).collect(),
                source_hash: "bench_multi".to_string(),
                metadata: serde_json::json!({}),
            },
        }
    }
}

impl DataStore for MultiBenchStore {
    fn load_bars(
        &self,
        symbol: SymbolId,
        _range: TimeRange,
        _version: Option<&str>,
    ) -> Result<RecordBatch> {
        self.bars
            .get(&symbol)
            .cloned()
            .ok_or_else(|| BtError::Data(format!("no bars for {:?}", symbol)))
    }

    fn load_features(
        &self,
        _feature_set: &str,
        _symbol: SymbolId,
        _range: TimeRange,
        _columns: Option<&[&str]>,
    ) -> Result<RecordBatch> {
        Err(BtError::Unsupported("benchmark store".to_string()))
    }

    fn load_exo(&self, _name: &str, _range: TimeRange) -> Result<RecordBatch> {
        Err(BtError::Unsupported("benchmark store".to_string()))
    }

    fn list_versions(&self, _dataset: &str) -> Result<Vec<DatasetVersion>> {
        Ok(vec![self.version.clone()])
    }

    fn active_version(&self, _dataset: &str) -> Result<DatasetVersion> {
        Ok(self.version.clone())
    }

    fn write_bars(
        &self,
        _symbol: SymbolId,
        _batch: &RecordBatch,
        _opts: WriteOptions,
    ) -> Result<DatasetVersion> {
        Err(BtError::Unsupported("benchmark store".to_string()))
    }
}

fn synthetic_bars_for_symbol(symbol_id: u32, rows: usize) -> RecordBatch {
    let schema = bt_data::schema::canonical_bars_1m_schema();
    let base_price = 100.0 * symbol_id as f64;
    let mut timestamps = Vec::with_capacity(rows);
    let mut prices = Vec::with_capacity(rows);
    let mut volume = Vec::with_capacity(rows);
    for idx in 0..rows {
        timestamps.push(idx as i64 * 60_000_000_000);
        prices.push(base_price + (idx as f64 * 0.0001));
        volume.push(100.0 + (idx % 25) as f64);
    }

    let arrays: Vec<ArrayRef> = vec![
        Arc::new(TimestampNanosecondArray::from(timestamps).with_timezone("UTC")),
        Arc::new(UInt32Array::from(vec![symbol_id; rows])),
        Arc::new(Float64Array::from(prices.clone())),
        Arc::new(Float64Array::from(prices.clone())),
        Arc::new(Float64Array::from(prices.clone())),
        Arc::new(Float64Array::from(prices.clone())),
        Arc::new(Float64Array::from(prices)),
        Arc::new(Float64Array::from(volume)),
        Arc::new(Float64Array::from(vec![Some(50.0); rows])),
        Arc::new(Float64Array::from(vec![Some(50.0); rows])),
        Arc::new(UInt32Array::from(vec![1_u32; rows])),
        Arc::new(Float64Array::from(vec![Some(99.9); rows])),
        Arc::new(Float64Array::from(vec![Some(100.1); rows])),
        Arc::new(Float64Array::from(vec![Some(0.2); rows])),
        Arc::new(BooleanArray::from(vec![false; rows])),
        Arc::new(UInt8Array::from(vec![0_u8; rows])),
    ];

    RecordBatch::try_new(schema, arrays).expect("valid benchmark batch")
}

fn benchmark_multi_asset_10_symbols(c: &mut Criterion) {
    let num_symbols = 10;
    let rows = 5_000;
    let store = MultiBenchStore::new(num_symbols, rows);
    let runner = SimpleBacktestRunner::new();
    let strategy = benchmark_strategy();

    let universe: Vec<SymbolId> = (1..=num_symbols).map(|s| SymbolId(s as u32)).collect();
    let config = BacktestConfig {
        universe,
        time_range: TimeRange::new(
            Timestamp::from_nanos(0),
            Timestamp::from_nanos(rows as i64 * 60_000_000_000),
        )
        .expect("valid range"),
        bar_interval: BarInterval::Minutes(1),
        initial_capital: 100_000.0,
        currency: "USD".to_string(),
        risk_free_rate: 0.025,
        execution: ExecutionConfig {
            signal_delay: 1,
            execution_price: ExecutionPrice::AtClose,
            max_position_pct: 1.0,
            allow_short: true,
            allow_fractional: true,
            skip_gap_bars: false,
            position_sizing_mode: PositionSizingMode::Units,
            fill_model: FillModel::default(),
            ..Default::default()
        },
        fees: FeeConfig {
            taker_fee_bps: 1.0,
            ..FeeConfig::default()
        },
        slippage: SlippageConfig::FixedBps { bps: 1.0 },
        data_version: None,
        rng_seed: Some(42),
        trading_days_per_year: 365.25,
        output_resolution: None,
        resample_to: None,
        feature_sets: vec![],
        ..Default::default()
    };

    let mut group = c.benchmark_group("runner");
    group.bench_function("multi_asset_10sym_5k_rows_target_lt_100ms", |b| {
        b.iter(|| {
            let result = runner
                .run(&strategy, &config, &store)
                .expect("benchmark run");
            black_box(result.metrics.total_return)
        })
    });
    group.finish();
}

criterion_group!(
    benches,
    benchmark_single_asset_runner,
    benchmark_sweep_1000_params,
    benchmark_multi_asset_10_symbols
);
criterion_main!(benches);
