//! Multi-leg atomic execution for prediction markets.
//!
//! Types and logic for executing multiple legs atomically with
//! various fill policies (AllOrNone, BestEffort, DeltaNeutral).

use pyo3::prelude::*;

/// Fill policy for multi-leg execution.
#[pyclass(eq, eq_int, frozen, hash)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum FillPolicy {
    AllOrNone = 0,
    BestEffort = 1,
    DeltaNeutral = 2,
}

#[pymethods]
impl FillPolicy {
    fn __repr__(&self) -> String {
        match self {
            Self::AllOrNone => "FillPolicy.AllOrNone",
            Self::BestEffort => "FillPolicy.BestEffort",
            Self::DeltaNeutral => "FillPolicy.DeltaNeutral",
        }
        .to_string()
    }
    fn __str__(&self) -> String { self.__repr__() }
}

/// Status of an individual leg.
#[pyclass(eq, eq_int, frozen, hash)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum LegStatus {
    Pending = 0,
    Submitted = 1,
    Filled = 2,
    PartiallyFilled = 3,
    Canceled = 4,
    Failed = 5,
}

#[pymethods]
impl LegStatus {
    fn __repr__(&self) -> String {
        match self {
            Self::Pending => "LegStatus.Pending",
            Self::Submitted => "LegStatus.Submitted",
            Self::Filled => "LegStatus.Filled",
            Self::PartiallyFilled => "LegStatus.PartiallyFilled",
            Self::Canceled => "LegStatus.Canceled",
            Self::Failed => "LegStatus.Failed",
        }
        .to_string()
    }
    fn __str__(&self) -> String { self.__repr__() }
}

/// A single leg in a multi-leg order.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct Leg {
    pub market_id: String,
    pub side: String,
    pub size: f64,
    pub price: f64,
    pub exchange: String,
}

#[pymethods]
impl Leg {
    #[new]
    #[pyo3(signature = (market_id, side, size, price, exchange="paper"))]
    fn new(market_id: String, side: String, size: f64, price: f64, exchange: &str) -> Self {
        Self { market_id, side, size, price, exchange: exchange.to_string() }
    }

    fn __repr__(&self) -> String {
        format!("Leg(market='{}', side='{}', size={:.2}, price={:.4})",
            self.market_id, self.side, self.size, self.price)
    }

    fn notional(&self) -> f64 {
        self.size * self.price
    }
}

/// Result of a single leg execution.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct LegResult {
    pub leg_index: usize,
    pub market_id: String,
    pub status: LegStatus,
    pub order_id: String,
    pub filled_size: f64,
    pub filled_price: f64,
    pub error: String,
}

#[pymethods]
impl LegResult {
    fn __repr__(&self) -> String {
        format!("LegResult(idx={}, market='{}', status={:?}, filled={:.2})",
            self.leg_index, self.market_id, self.status, self.filled_size)
    }
}

/// Result of multi-leg execution.
#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct MultiLegResult {
    pub policy: FillPolicy,
    pub legs: Vec<LegResult>,
    pub total_filled: usize,
    pub total_legs: usize,
    pub total_notional: f64,
    pub success: bool,
    pub error: String,
}

#[pymethods]
impl MultiLegResult {
    fn __repr__(&self) -> String {
        format!("MultiLegResult(policy={:?}, filled={}/{}, success={})",
            self.policy, self.total_filled, self.total_legs, self.success)
    }

    fn fill_rate(&self) -> f64 {
        if self.total_legs == 0 { return 0.0; }
        self.total_filled as f64 / self.total_legs as f64
    }
}

/// Validate legs before execution. Returns error message or empty string.
#[pyfunction]
pub fn validate_legs(legs: Vec<Leg>) -> String {
    if legs.is_empty() {
        return "no legs provided".to_string();
    }
    for (i, leg) in legs.iter().enumerate() {
        if leg.size <= 0.0 {
            return format!("leg {} has invalid size: {}", i, leg.size);
        }
        if leg.price <= 0.0 || leg.price >= 1.0 {
            return format!("leg {} has invalid price: {}", i, leg.price);
        }
        if leg.market_id.is_empty() {
            return format!("leg {} has empty market_id", i);
        }
    }
    String::new()
}

/// Compute aggregate risk metrics for a set of legs.
///
/// Returns (total_notional, max_single_leg, num_exchanges).
#[pyfunction]
pub fn aggregate_leg_risk(legs: Vec<Leg>) -> (f64, f64, usize) {
    let mut total_notional = 0.0;
    let mut max_leg = 0.0f64;
    let mut exchanges = std::collections::HashSet::new();

    for leg in &legs {
        let notional = leg.size * leg.price;
        total_notional += notional;
        max_leg = max_leg.max(notional);
        exchanges.insert(leg.exchange.clone());
    }

    (total_notional, max_leg, exchanges.len())
}

/// Build a mock MultiLegResult (for testing / paper execution).
#[pyfunction]
pub fn build_mock_result(
    legs: Vec<Leg>,
    policy: FillPolicy,
    success: bool,
) -> MultiLegResult {
    let total_legs = legs.len();
    let mut leg_results = Vec::with_capacity(total_legs);
    let mut total_notional = 0.0;
    let mut filled = 0;

    for (i, leg) in legs.iter().enumerate() {
        let status = if success { LegStatus::Filled } else { LegStatus::Failed };
        let filled_size = if success { leg.size } else { 0.0 };
        let filled_price = if success { leg.price } else { 0.0 };
        if success { filled += 1; }
        total_notional += leg.size * leg.price;

        leg_results.push(LegResult {
            leg_index: i,
            market_id: leg.market_id.clone(),
            status,
            order_id: format!("mock_{}", i),
            filled_size,
            filled_price,
            error: String::new(),
        });
    }

    MultiLegResult {
        policy,
        legs: leg_results,
        total_filled: filled,
        total_legs,
        total_notional,
        success,
        error: String::new(),
    }
}

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<FillPolicy>()?;
    m.add_class::<LegStatus>()?;
    m.add_class::<Leg>()?;
    m.add_class::<LegResult>()?;
    m.add_class::<MultiLegResult>()?;
    m.add_function(wrap_pyfunction!(validate_legs, m)?)?;
    m.add_function(wrap_pyfunction!(aggregate_leg_risk, m)?)?;
    m.add_function(wrap_pyfunction!(build_mock_result, m)?)?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_leg_notional() {
        let leg = Leg {
            market_id: "M1".into(), side: "buy".into(),
            size: 100.0, price: 0.50, exchange: "paper".into(),
        };
        assert!((leg.notional() - 50.0).abs() < 1e-10);
    }

    #[test]
    fn test_validate_legs_valid() {
        let legs = vec![
            Leg { market_id: "M1".into(), side: "buy".into(), size: 10.0, price: 0.5, exchange: "paper".into() },
        ];
        assert!(validate_legs(legs).is_empty());
    }

    #[test]
    fn test_validate_legs_empty() {
        assert!(!validate_legs(vec![]).is_empty());
    }

    #[test]
    fn test_validate_legs_invalid_size() {
        let legs = vec![
            Leg { market_id: "M1".into(), side: "buy".into(), size: 0.0, price: 0.5, exchange: "paper".into() },
        ];
        assert!(validate_legs(legs).contains("invalid size"));
    }

    #[test]
    fn test_validate_legs_invalid_price() {
        let legs = vec![
            Leg { market_id: "M1".into(), side: "buy".into(), size: 10.0, price: 1.0, exchange: "paper".into() },
        ];
        assert!(validate_legs(legs).contains("invalid price"));
    }

    #[test]
    fn test_aggregate_leg_risk() {
        let legs = vec![
            Leg { market_id: "M1".into(), side: "buy".into(), size: 100.0, price: 0.5, exchange: "poly".into() },
            Leg { market_id: "M2".into(), side: "sell".into(), size: 200.0, price: 0.3, exchange: "kalshi".into() },
        ];
        let (total, max, n_exch) = aggregate_leg_risk(legs);
        assert!((total - 110.0).abs() < 1e-10);
        assert!((max - 60.0).abs() < 1e-10);
        assert_eq!(n_exch, 2);
    }

    #[test]
    fn test_build_mock_result_success() {
        let legs = vec![
            Leg { market_id: "M1".into(), side: "buy".into(), size: 10.0, price: 0.5, exchange: "paper".into() },
            Leg { market_id: "M2".into(), side: "sell".into(), size: 20.0, price: 0.3, exchange: "paper".into() },
        ];
        let result = build_mock_result(legs, FillPolicy::AllOrNone, true);
        assert!(result.success);
        assert_eq!(result.total_filled, 2);
        assert_eq!(result.total_legs, 2);
        assert!((result.fill_rate() - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_build_mock_result_failure() {
        let legs = vec![
            Leg { market_id: "M1".into(), side: "buy".into(), size: 10.0, price: 0.5, exchange: "paper".into() },
        ];
        let result = build_mock_result(legs, FillPolicy::BestEffort, false);
        assert!(!result.success);
        assert_eq!(result.total_filled, 0);
    }
}
