use std::collections::HashMap;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use arrow::array::{
    Array, ArrayRef, BooleanArray, Float64Array, TimestampNanosecondArray, UInt32Array, UInt8Array,
};
use arrow::record_batch::RecordBatch;
use bt_core::{
    BacktestConfig, BacktestRunner, BarInterval, ExecutionConfig, ExecutionPrice, FeeConfig,
    FillModel, OrderConfig, 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 parquet::arrow::ArrowWriter;
use parquet::basic::Compression;
use parquet::file::properties::WriterProperties;

fn main() -> Result<()> {
    let fixture_root = fixture_root();
    std::fs::create_dir_all(&fixture_root)?;

    let bars_s1 = golden_bars(1, &[100.0, 101.0, 102.0, 103.0]);
    let bars_s2 = golden_bars(2, &[200.0, 198.0, 202.0, 205.0]);

    write_bars_parquet(&fixture_root.join("bars_symbol_1.parquet"), &bars_s1)?;
    write_bars_parquet(&fixture_root.join("bars_symbol_2.parquet"), &bars_s2)?;

    // Strategy: half-allocation per symbol (target = 0.5)
    let signal = compile_expr(Expr::Literal(ScalarValue::Float64(1.0)))
        .map_err(|err| BtError::Data(format!("compile signal: {err}")))?;
    let sizing = compile_expr(Expr::Mul(
        Box::new(Expr::Column("signal".to_string())),
        Box::new(Expr::Literal(ScalarValue::Float64(0.5))),
    ))
    .map_err(|err| BtError::Data(format!("compile sizing: {err}")))?;

    let strategy = StrategyPlan {
        name: "golden_multi_asset_half".to_string(),
        signals: HashMap::from([("signal".to_string(), signal)]),
        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,
    };

    let config = golden_config();
    let store = MultiAssetGoldenStore::new(bars_s1, bars_s2);
    let result = SimpleBacktestRunner::new().run(&strategy, &config, &store)?;

    // Write expected equity.
    let equity = result
        .equity_curve
        .as_any()
        .downcast_ref::<Float64Array>()
        .ok_or_else(|| BtError::Data("equity curve must be Float64".to_string()))?;
    let equity_values: Vec<f64> = (0..equity.len()).map(|i| equity.value(i)).collect();
    write_json(&fixture_root.join("expected_equity.json"), &equity_values)?;

    // Write expected trades.
    let trades_json = trade_snapshot(&result.trades)?;
    write_json(&fixture_root.join("expected_trades.json"), &trades_json)?;

    println!(
        "Multi-asset golden fixtures generated in {}",
        fixture_root.display()
    );
    println!("Equity: {:?}", equity_values);
    println!("Trade count: {}", result.trades.num_rows());
    Ok(())
}

struct MultiAssetGoldenStore {
    bars_s1: RecordBatch,
    bars_s2: RecordBatch,
}

impl MultiAssetGoldenStore {
    fn new(bars_s1: RecordBatch, bars_s2: RecordBatch) -> Self {
        Self { bars_s1, bars_s2 }
    }
}

impl DataStore for MultiAssetGoldenStore {
    fn load_bars(
        &self,
        symbol: SymbolId,
        _range: TimeRange,
        _version: Option<&str>,
    ) -> Result<RecordBatch> {
        match symbol.0 {
            1 => Ok(self.bars_s1.clone()),
            2 => Ok(self.bars_s2.clone()),
            _ => Err(BtError::Data(format!("unknown symbol {:?}", symbol))),
        }
    }

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

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

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

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

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

fn trade_snapshot(trades: &RecordBatch) -> Result<serde_json::Value> {
    let quantity = f64_column(trades, "quantity")?;
    let fill_price = f64_column(trades, "fill_price")?;
    let side = u8_column(trades, "side")?;
    let symbol_id = u32_column(trades, "symbol_id")?;

    let mut out = Vec::with_capacity(trades.num_rows());
    for idx in 0..trades.num_rows() {
        out.push(serde_json::json!({
            "symbol_id": symbol_id.value(idx),
            "side": side.value(idx),
            "quantity": quantity.value(idx),
            "fill_price": fill_price.value(idx),
        }));
    }

    Ok(serde_json::json!(out))
}

fn write_bars_parquet(path: &Path, batch: &RecordBatch) -> Result<()> {
    let file = File::create(path)?;
    let props = WriterProperties::builder()
        .set_compression(Compression::ZSTD(Default::default()))
        .build();

    let mut writer = ArrowWriter::try_new(file, batch.schema(), Some(props))?;
    writer.write(batch)?;
    writer.close()?;
    Ok(())
}

fn write_json<T: serde::Serialize>(path: &Path, value: &T) -> Result<()> {
    let payload = serde_json::to_string_pretty(value)?;
    std::fs::write(path, payload)?;
    Ok(())
}

fn golden_bars(symbol_id: u32, closes: &[f64]) -> RecordBatch {
    let schema = bt_data::schema::canonical_bars_1m_schema();
    let n = closes.len();
    let timestamps: Vec<i64> = (0..n as i64).map(|i| i * 60_000_000_000).collect();

    let arrays: Vec<ArrayRef> = vec![
        Arc::new(TimestampNanosecondArray::from(timestamps).with_timezone("UTC")),
        Arc::new(UInt32Array::from(vec![symbol_id; n])),
        Arc::new(Float64Array::from(closes.to_vec())), // open
        Arc::new(Float64Array::from(closes.to_vec())), // high
        Arc::new(Float64Array::from(closes.to_vec())), // low
        Arc::new(Float64Array::from(closes.to_vec())), // close
        Arc::new(Float64Array::from(closes.to_vec())), // vwap
        Arc::new(Float64Array::from(vec![10.0; n])),   // volume
        Arc::new(Float64Array::from(
            closes.iter().map(|_| Some(5.0)).collect::<Vec<_>>(),
        )), // buy_volume
        Arc::new(Float64Array::from(
            closes.iter().map(|_| Some(5.0)).collect::<Vec<_>>(),
        )), // sell_volume
        Arc::new(UInt32Array::from(vec![1_u32; n])),   // trade_count
        Arc::new(Float64Array::from(
            closes.iter().map(|c| Some(c - 0.1)).collect::<Vec<_>>(),
        )), // bid
        Arc::new(Float64Array::from(
            closes.iter().map(|c| Some(c + 0.1)).collect::<Vec<_>>(),
        )), // ask
        Arc::new(Float64Array::from(
            closes.iter().map(|_| Some(0.2)).collect::<Vec<_>>(),
        )), // spread
        Arc::new(BooleanArray::from(vec![false; n])),  // is_gap
        Arc::new(UInt8Array::from(vec![0_u8; n])),     // gap_fill_method
    ];

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

fn golden_config() -> BacktestConfig {
    BacktestConfig {
        universe: vec![SymbolId(1), SymbolId(2)],
        time_range: TimeRange::new(
            Timestamp::from_nanos(0),
            Timestamp::from_nanos(240_000_000_000),
        )
        .expect("range"),
        bar_interval: BarInterval::Minutes(1),
        initial_capital: 10000.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: false,
            allow_fractional: true,
            skip_gap_bars: false,
            position_sizing_mode: PositionSizingMode::Units,
            fill_model: FillModel::default(),
            orders: OrderConfig::default(),
            pyramiding: false,
        },
        fees: FeeConfig::default(),
        slippage: SlippageConfig::FixedBps { bps: 0.0 },
        data_version: Some("golden_multi_v1".to_string()),
        rng_seed: Some(7),
        trading_days_per_year: 365.25,
        output_resolution: Some(BarInterval::Minutes(1)),
        resample_to: None,
        feature_sets: vec![],
        symbol_names: HashMap::new(),
        warmup_bars: 0,
    }
}

fn golden_version() -> DatasetVersion {
    DatasetVersion {
        id: "golden_multi_v1".to_string(),
        dataset: "bars_1m".to_string(),
        created_at: Timestamp::from_nanos(1),
        schema_hash: "schema".to_string(),
        row_count: 4,
        time_range: TimeRange::new(
            Timestamp::from_nanos(0),
            Timestamp::from_nanos(240_000_000_000),
        )
        .expect("range"),
        symbols: vec![SymbolId(1), SymbolId(2)],
        source_hash: "golden_multi".to_string(),
        metadata: serde_json::json!({}),
    }
}

fn fixture_root() -> PathBuf {
    Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("tests")
        .join("fixtures")
        .join("golden")
        .join("multi_asset")
        .join("v1")
}

fn f64_column<'a>(batch: &'a RecordBatch, name: &str) -> Result<&'a Float64Array> {
    batch
        .column_by_name(name)
        .ok_or_else(|| BtError::Data(format!("missing column '{name}'")))?
        .as_any()
        .downcast_ref::<Float64Array>()
        .ok_or_else(|| BtError::Data(format!("column '{name}' must be Float64")))
}

fn u8_column<'a>(batch: &'a RecordBatch, name: &str) -> Result<&'a UInt8Array> {
    batch
        .column_by_name(name)
        .ok_or_else(|| BtError::Data(format!("missing column '{name}'")))?
        .as_any()
        .downcast_ref::<UInt8Array>()
        .ok_or_else(|| BtError::Data(format!("column '{name}' must be UInt8")))
}

fn u32_column<'a>(batch: &'a RecordBatch, name: &str) -> Result<&'a UInt32Array> {
    batch
        .column_by_name(name)
        .ok_or_else(|| BtError::Data(format!("missing column '{name}'")))?
        .as_any()
        .downcast_ref::<UInt32Array>()
        .ok_or_else(|| BtError::Data(format!("column '{name}' must be UInt32")))
}
