use pyo3::prelude::*;
use std::collections::HashMap;
use std::sync::RwLock;
use lazy_static::lazy_static;
use crate::tail::CandleData;
use serde::{Deserialize, Serialize};

// TickData 구조체 정의 (임시, 필요시 확장)
#[pyclass]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TickData {
    #[pyo3(get, set)]
    pub time_str: String,
    #[pyo3(get, set)]
    pub price: f64,
}

#[pymethods]
impl TickData {
    #[new]
    pub fn new(time_str: String, price: f64) -> Self {
        Self { time_str, price }
    }
}

pub struct BacktestSessionData {
    pub candles_by_code: HashMap<String, Vec<CandleData>>,
    pub ticks_by_code: HashMap<String, Vec<TickData>>,
}

lazy_static! {
    pub static ref SESSIONS: RwLock<HashMap<String, BacktestSessionData>> = RwLock::new(HashMap::new());
}

#[pyfunction]
pub fn init_backtest_session(
    session_id: String,
    candles_by_code: HashMap<String, Vec<CandleData>>,
    ticks_by_code: HashMap<String, Vec<TickData>>,
) -> PyResult<()> {
    let mut sessions = SESSIONS.write().unwrap();
    sessions.insert(
        session_id.clone(),
        BacktestSessionData {
            candles_by_code,
            ticks_by_code,
        },
    );
    println!("[Rust] Session {} initialized.", session_id);
    Ok(())
}

#[pyfunction]
pub fn clear_backtest_session(session_id: String) -> PyResult<()> {
    let mut sessions = SESSIONS.write().unwrap();
    sessions.remove(&session_id);
    println!("[Rust] Session {} cleared.", session_id);
    Ok(())
}

#[pyfunction]
pub fn init_backtest_session_json(
    session_id: String,
    candles_json: String,
) -> PyResult<()> {
    let candles_by_code: HashMap<String, Vec<CandleData>> = serde_json::from_str(&candles_json)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("JSON parsing error: {}", e)))?;
        
    let mut sessions = SESSIONS.write().unwrap();
    sessions.insert(
        session_id.clone(),
        BacktestSessionData {
            candles_by_code,
            ticks_by_code: HashMap::new(),
        },
    );
    println!("[Rust] Session {} initialized via JSON. ({} codes)", session_id, sessions.get(&session_id).unwrap().candles_by_code.len());
    Ok(())
}