hyper-agent-core 0.1.0

use tokio::sync::Mutex;

use rusqlite::{params, Connection, Result as SqlResult};
use serde::{Deserialize, Serialize};

/// A trading position tracked in the local SQLite database.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Position {
    pub id: String,
    pub market: String,
    pub side: String,
    pub size: f64,
    pub entry_price: f64,
    pub current_price: Option<f64>,
    pub status: String,
    pub pnl: Option<f64>,
    pub mode: String,
    pub strategy: Option<String>,
    pub opened_at: String,
    pub closed_at: Option<String>,
    pub close_reason: Option<String>,
}

/// A single execution quality data point recorded after an order is submitted.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionMetric {
    pub order_id: String,
    pub market: String,
    pub side: String,
    pub requested_size: f64,
    pub filled_size: f64,
    pub status: String,
    pub mode: String,
}

/// Aggregated execution quality summary.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionSummary {
    pub total_orders: u64,
    pub full_fills: u64,
    pub partial_fills: u64,
    pub zero_fills: u64,
    pub avg_fill_rate_pct: f64,
}

/// Aggregated profit-and-loss summary over a time window.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PnlSummary {
    pub total_pnl: f64,
    pub win_count: u32,
    pub loss_count: u32,
    pub win_rate: f64,
    pub best_trade: Option<(String, f64)>,
    pub worst_trade: Option<(String, f64)>,
}

/// Error type for position manager operations.
#[derive(Debug, thiserror::Error)]
pub enum PositionError {
    #[error("Database error: {0}")]
    Db(#[from] rusqlite::Error),
    #[error("Position not found: {0}")]
    NotFound(String),
}

/// Manages position state in a local SQLite database.
///
/// Provides CRUD operations for positions, P&L tracking, and
/// query helpers for open/closed position lists.
///
/// The inner `Connection` is wrapped in a `tokio::sync::Mutex` so that
/// `PositionManager` is `Send + Sync` and can be shared via `Arc`
/// without risk of deadlocks when the lock is held across `.await` points.
pub struct PositionManager {
    db: Mutex<Connection>,
}

impl PositionManager {
    /// Open (or create) the SQLite database at `db_path` and ensure the
    /// positions table exists.
    pub fn new(db_path: &str) -> Result<Self, PositionError> {
        let db = Connection::open(db_path)?;
        db.execute_batch("PRAGMA journal_mode=WAL; PRAGMA foreign_keys=ON;")?;
        db.execute(
            "CREATE TABLE IF NOT EXISTS positions (
                id            TEXT PRIMARY KEY,
                market        TEXT NOT NULL,
                side          TEXT NOT NULL CHECK(side IN ('long','short')),
                size          REAL NOT NULL,
                entry_price   REAL NOT NULL,
                current_price REAL,
                status        TEXT NOT NULL DEFAULT 'open' CHECK(status IN ('open','closed')),
                pnl           REAL,
                mode          TEXT NOT NULL CHECK(mode IN ('live','paper')),
                strategy      TEXT,
                opened_at     TEXT NOT NULL,
                closed_at     TEXT,
                close_reason  TEXT
            )",
            [],
        )?;
        db.execute(
            "CREATE TABLE IF NOT EXISTS execution_metrics (
                id             INTEGER PRIMARY KEY AUTOINCREMENT,
                timestamp      TEXT NOT NULL,
                order_id       TEXT NOT NULL,
                market         TEXT NOT NULL,
                side           TEXT NOT NULL,
                requested_size REAL NOT NULL,
                filled_size    REAL NOT NULL,
                fill_rate      REAL NOT NULL,
                status         TEXT NOT NULL,
                mode           TEXT NOT NULL
            )",
            [],
        )?;
        Ok(Self { db: Mutex::new(db) })
    }

    /// Create an in-memory database (useful for tests).
    #[cfg(test)]
    pub fn in_memory() -> Result<Self, PositionError> {
        Self::new(":memory:")
    }

    /// Expose the database connection for test-only direct SQL manipulation.
    #[cfg(test)]
    pub async fn lock_db_for_test(&self) -> tokio::sync::MutexGuard<'_, Connection> {
        self.db.lock().await
    }

    /// Insert a new position into the database.
    pub async fn open_position(&self, pos: &Position) -> Result<(), PositionError> {
        let db = self.db.lock().await;
        db.execute(
            "INSERT INTO positions
                (id, market, side, size, entry_price, current_price,
                 status, pnl, mode, strategy, opened_at, closed_at, close_reason)
             VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11,?12,?13)",
            params![
                pos.id,
                pos.market,
                pos.side,
                pos.size,
                pos.entry_price,
                pos.current_price,
                pos.status,
                pos.pnl,
                pos.mode,
                pos.strategy,
                pos.opened_at,
                pos.closed_at,
                pos.close_reason,
            ],
        )?;
        Ok(())
    }

    /// Close a position: set status to "closed", record exit price,
    /// compute realised P&L, and store the close reason.
    pub async fn close_position(
        &self,
        id: &str,
        price: f64,
        reason: &str,
    ) -> Result<(), PositionError> {
        let db = self.db.lock().await;
        let pos = Self::get_position_inner(&db, id)?
            .ok_or_else(|| PositionError::NotFound(id.to_string()))?;

        let pnl = compute_pnl(&pos.side, pos.entry_price, price, pos.size);
        let now = chrono::Utc::now().to_rfc3339();

        let changed = db.execute(
            "UPDATE positions
                SET status = 'closed',
                    current_price = ?1,
                    pnl = ?2,
                    closed_at = ?3,
                    close_reason = ?4
              WHERE id = ?5 AND status = 'open'",
            params![price, pnl, now, reason, id],
        )?;

        if changed == 0 {
            return Err(PositionError::NotFound(id.to_string()));
        }
        Ok(())
    }

    /// Update the mark price of an open position and recompute unrealised P&L.
    pub async fn update_price(&self, id: &str, price: f64) -> Result<(), PositionError> {
        let db = self.db.lock().await;
        let pos = Self::get_position_inner(&db, id)?
            .ok_or_else(|| PositionError::NotFound(id.to_string()))?;

        let pnl = compute_pnl(&pos.side, pos.entry_price, price, pos.size);

        let changed = db.execute(
            "UPDATE positions SET current_price = ?1, pnl = ?2 WHERE id = ?3 AND status = 'open'",
            params![price, pnl, id],
        )?;

        if changed == 0 {
            return Err(PositionError::NotFound(id.to_string()));
        }
        Ok(())
    }

    /// Return all positions with status = "open".
    pub async fn list_open(&self) -> Result<Vec<Position>, PositionError> {
        let db = self.db.lock().await;
        let mut stmt =
            db.prepare("SELECT * FROM positions WHERE status = 'open' ORDER BY opened_at DESC")?;
        let rows = stmt.query_map([], row_to_position)?;
        rows.collect::<SqlResult<Vec<_>>>()
            .map_err(PositionError::from)
    }

    /// Return the most recent `limit` closed positions.
    pub async fn list_closed(&self, limit: usize) -> Result<Vec<Position>, PositionError> {
        let db = self.db.lock().await;
        let mut stmt = db.prepare(
            "SELECT * FROM positions WHERE status = 'closed' ORDER BY closed_at DESC LIMIT ?1",
        )?;
        let rows = stmt.query_map(params![limit as i64], row_to_position)?;
        rows.collect::<SqlResult<Vec<_>>>()
            .map_err(PositionError::from)
    }

    /// Compute an aggregate P&L summary for closed positions in the last
    /// `days` calendar days.
    pub async fn get_pnl_summary(&self, days: u32) -> Result<PnlSummary, PositionError> {
        let cutoff = chrono::Utc::now() - chrono::Duration::days(i64::from(days));
        let cutoff_str = cutoff.to_rfc3339();

        let db = self.db.lock().await;
        let mut stmt = db.prepare(
            "SELECT id, pnl FROM positions
              WHERE status = 'closed' AND closed_at >= ?1
              ORDER BY pnl DESC",
        )?;

        let trades: Vec<(String, f64)> = stmt
            .query_map(params![cutoff_str], |row| {
                Ok((row.get::<_, String>(0)?, row.get::<_, f64>(1)?))
            })?
            .collect::<SqlResult<Vec<_>>>()?;

        let mut total_pnl = 0.0;
        let mut win_count: u32 = 0;
        let mut loss_count: u32 = 0;
        let mut best: Option<(String, f64)> = None;
        let mut worst: Option<(String, f64)> = None;

        for (id, pnl) in &trades {
            total_pnl += pnl;
            if *pnl >= 0.0 {
                win_count += 1;
            } else {
                loss_count += 1;
            }
            if best.as_ref().map_or(true, |(_, b)| pnl > b) {
                best = Some((id.clone(), *pnl));
            }
            if worst.as_ref().map_or(true, |(_, w)| pnl < w) {
                worst = Some((id.clone(), *pnl));
            }
        }

        let total = win_count + loss_count;
        let win_rate = if total > 0 {
            f64::from(win_count) / f64::from(total)
        } else {
            0.0
        };

        Ok(PnlSummary {
            total_pnl,
            win_count,
            loss_count,
            win_rate,
            best_trade: best,
            worst_trade: worst,
        })
    }

    /// Record an execution quality metric after an order is submitted.
    pub async fn record_execution(&self, metric: ExecutionMetric) -> Result<(), PositionError> {
        let db = self.db.lock().await;
        let fill_rate = if metric.requested_size > 0.0 {
            (metric.filled_size / metric.requested_size) * 100.0
        } else {
            0.0
        };
        let now = chrono::Utc::now().to_rfc3339();
        db.execute(
            "INSERT INTO execution_metrics
                (timestamp, order_id, market, side, requested_size, filled_size,
                 fill_rate, status, mode)
             VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9)",
            params![
                now,
                metric.order_id,
                metric.market,
                metric.side,
                metric.requested_size,
                metric.filled_size,
                fill_rate,
                metric.status,
                metric.mode,
            ],
        )?;
        Ok(())
    }

    /// Return an aggregated execution quality summary across all recorded metrics.
    pub async fn get_execution_summary(&self) -> Result<ExecutionSummary, PositionError> {
        let db = self.db.lock().await;
        let mut stmt = db.prepare("SELECT fill_rate, status FROM execution_metrics")?;
        let rows: Vec<(f64, String)> = stmt
            .query_map([], |row| {
                Ok((row.get::<_, f64>(0)?, row.get::<_, String>(1)?))
            })?
            .collect::<SqlResult<Vec<_>>>()?;

        let total_orders = rows.len() as u64;
        let mut full_fills: u64 = 0;
        let mut partial_fills: u64 = 0;
        let mut zero_fills: u64 = 0;
        let mut fill_rate_sum: f64 = 0.0;

        for (fill_rate, status) in &rows {
            fill_rate_sum += fill_rate;
            match status.as_str() {
                "partial_fill" => partial_fills += 1,
                "resting" => zero_fills += 1,
                _ => {
                    // "filled", "simulated", "ok", etc. are full fills
                    if *fill_rate >= 100.0 - f64::EPSILON {
                        full_fills += 1;
                    } else if *fill_rate > 0.0 {
                        partial_fills += 1;
                    } else {
                        zero_fills += 1;
                    }
                }
            }
        }

        let avg_fill_rate_pct = if total_orders > 0 {
            fill_rate_sum / total_orders as f64
        } else {
            0.0
        };

        Ok(ExecutionSummary {
            total_orders,
            full_fills,
            partial_fills,
            zero_fills,
            avg_fill_rate_pct,
        })
    }

    /// Fetch a single position by ID, or `None` if it does not exist.
    pub async fn get_position(&self, id: &str) -> Result<Option<Position>, PositionError> {
        let db = self.db.lock().await;
        Self::get_position_inner(&db, id)
    }

    /// Inner helper that operates on a borrowed `Connection` (no locking).
    fn get_position_inner(db: &Connection, id: &str) -> Result<Option<Position>, PositionError> {
        let mut stmt = db.prepare("SELECT * FROM positions WHERE id = ?1")?;
        let mut rows = stmt.query_map(params![id], row_to_position)?;
        match rows.next() {
            Some(row) => Ok(Some(row?)),
            None => Ok(None),
        }
    }
}

/// Compute P&L given direction, entry, exit, and size.
fn compute_pnl(side: &str, entry: f64, exit: f64, size: f64) -> f64 {
    match side {
        "long" => (exit - entry) * size,
        "short" => (entry - exit) * size,
        _ => 0.0,
    }
}

/// Map a rusqlite row to a `Position`.
fn row_to_position(row: &rusqlite::Row) -> SqlResult<Position> {
    Ok(Position {
        id: row.get("id")?,
        market: row.get("market")?,
        side: row.get("side")?,
        size: row.get("size")?,
        entry_price: row.get("entry_price")?,
        current_price: row.get("current_price")?,
        status: row.get("status")?,
        pnl: row.get("pnl")?,
        mode: row.get("mode")?,
        strategy: row.get("strategy")?,
        opened_at: row.get("opened_at")?,
        closed_at: row.get("closed_at")?,
        close_reason: row.get("close_reason")?,
    })
}

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

    fn make_position(id: &str, market: &str, side: &str, size: f64, entry: f64) -> Position {
        Position {
            id: id.to_string(),
            market: market.to_string(),
            side: side.to_string(),
            size,
            entry_price: entry,
            current_price: None,
            status: "open".to_string(),
            pnl: None,
            mode: "paper".to_string(),
            strategy: Some("test-strat".to_string()),
            opened_at: chrono::Utc::now().to_rfc3339(),
            closed_at: None,
            close_reason: None,
        }
    }

    #[tokio::test]
    async fn test_open_and_get_position() {
        let pm = PositionManager::in_memory().unwrap();
        let pos = make_position("pos-1", "BTC-PERP", "long", 0.5, 60000.0);
        pm.open_position(&pos).await.unwrap();

        let fetched = pm.get_position("pos-1").await.unwrap().unwrap();
        assert_eq!(fetched.market, "BTC-PERP");
        assert_eq!(fetched.side, "long");
        assert_eq!(fetched.size, 0.5);
        assert_eq!(fetched.entry_price, 60000.0);
        assert_eq!(fetched.status, "open");
    }

    #[tokio::test]
    async fn test_get_nonexistent_returns_none() {
        let pm = PositionManager::in_memory().unwrap();
        assert!(pm.get_position("nonexistent").await.unwrap().is_none());
    }

    #[tokio::test]
    async fn test_close_position_long_profit() {
        let pm = PositionManager::in_memory().unwrap();
        let pos = make_position("pos-long", "ETH-PERP", "long", 2.0, 3000.0);
        pm.open_position(&pos).await.unwrap();

        pm.close_position("pos-long", 3500.0, "take-profit")
            .await
            .unwrap();

        let closed = pm.get_position("pos-long").await.unwrap().unwrap();
        assert_eq!(closed.status, "closed");
        assert!((closed.pnl.unwrap() - 1000.0).abs() < 1e-6);
        assert_eq!(closed.close_reason.as_deref(), Some("take-profit"));
        assert!(closed.closed_at.is_some());
    }

    #[tokio::test]
    async fn test_close_position_short_profit() {
        let pm = PositionManager::in_memory().unwrap();
        let pos = make_position("pos-short", "SOL-PERP", "short", 10.0, 100.0);
        pm.open_position(&pos).await.unwrap();

        pm.close_position("pos-short", 90.0, "take-profit")
            .await
            .unwrap();

        let closed = pm.get_position("pos-short").await.unwrap().unwrap();
        assert!((closed.pnl.unwrap() - 100.0).abs() < 1e-6);
    }

    #[tokio::test]
    async fn test_close_position_long_loss() {
        let pm = PositionManager::in_memory().unwrap();
        let pos = make_position("pos-loss", "BTC-PERP", "long", 1.0, 60000.0);
        pm.open_position(&pos).await.unwrap();

        pm.close_position("pos-loss", 59000.0, "stop-loss")
            .await
            .unwrap();

        let closed = pm.get_position("pos-loss").await.unwrap().unwrap();
        assert!((closed.pnl.unwrap() - (-1000.0)).abs() < 1e-6);
    }

    #[tokio::test]
    async fn test_close_nonexistent_position() {
        let pm = PositionManager::in_memory().unwrap();
        let err = pm
            .close_position("ghost", 100.0, "reason")
            .await
            .unwrap_err();
        assert!(matches!(err, PositionError::NotFound(_)));
    }

    #[tokio::test]
    async fn test_update_price() {
        let pm = PositionManager::in_memory().unwrap();
        let pos = make_position("pos-upd", "BTC-PERP", "long", 1.0, 60000.0);
        pm.open_position(&pos).await.unwrap();

        pm.update_price("pos-upd", 61000.0).await.unwrap();

        let fetched = pm.get_position("pos-upd").await.unwrap().unwrap();
        assert_eq!(fetched.current_price, Some(61000.0));
        assert!((fetched.pnl.unwrap() - 1000.0).abs() < 1e-6);
    }

    #[tokio::test]
    async fn test_list_open_and_closed() {
        let pm = PositionManager::in_memory().unwrap();
        pm.open_position(&make_position("a", "BTC-PERP", "long", 1.0, 60000.0))
            .await
            .unwrap();
        pm.open_position(&make_position("b", "ETH-PERP", "short", 5.0, 3000.0))
            .await
            .unwrap();
        pm.open_position(&make_position("c", "SOL-PERP", "long", 10.0, 100.0))
            .await
            .unwrap();

        assert_eq!(pm.list_open().await.unwrap().len(), 3);
        assert_eq!(pm.list_closed(10).await.unwrap().len(), 0);

        pm.close_position("a", 61000.0, "tp").await.unwrap();

        assert_eq!(pm.list_open().await.unwrap().len(), 2);
        assert_eq!(pm.list_closed(10).await.unwrap().len(), 1);
    }

    #[tokio::test]
    async fn test_pnl_summary() {
        let pm = PositionManager::in_memory().unwrap();

        // 3 trades: 2 wins, 1 loss
        pm.open_position(&make_position("w1", "BTC-PERP", "long", 1.0, 60000.0))
            .await
            .unwrap();
        pm.open_position(&make_position("w2", "ETH-PERP", "short", 10.0, 3000.0))
            .await
            .unwrap();
        pm.open_position(&make_position("l1", "SOL-PERP", "long", 100.0, 100.0))
            .await
            .unwrap();

        pm.close_position("w1", 62000.0, "tp").await.unwrap(); // +2000
        pm.close_position("w2", 2800.0, "tp").await.unwrap(); // +2000
        pm.close_position("l1", 95.0, "sl").await.unwrap(); // -500

        let summary = pm.get_pnl_summary(30).await.unwrap();
        assert!((summary.total_pnl - 3500.0).abs() < 1e-6);
        assert_eq!(summary.win_count, 2);
        assert_eq!(summary.loss_count, 1);
        assert!((summary.win_rate - 2.0 / 3.0).abs() < 1e-6);

        let (best_id, best_pnl) = summary.best_trade.unwrap();
        assert!((best_pnl - 2000.0).abs() < 1e-6);
        // best could be w1 or w2 (both +2000), just check it's one of them
        assert!(best_id == "w1" || best_id == "w2");

        let (_, worst_pnl) = summary.worst_trade.unwrap();
        assert!((worst_pnl - (-500.0)).abs() < 1e-6);
    }

    #[tokio::test]
    async fn test_pnl_summary_empty() {
        let pm = PositionManager::in_memory().unwrap();
        let summary = pm.get_pnl_summary(30).await.unwrap();
        assert_eq!(summary.total_pnl, 0.0);
        assert_eq!(summary.win_count, 0);
        assert_eq!(summary.loss_count, 0);
        assert_eq!(summary.win_rate, 0.0);
        assert!(summary.best_trade.is_none());
        assert!(summary.worst_trade.is_none());
    }

    #[tokio::test]
    async fn test_duplicate_id_fails() {
        let pm = PositionManager::in_memory().unwrap();
        let pos = make_position("dup", "BTC-PERP", "long", 1.0, 60000.0);
        pm.open_position(&pos).await.unwrap();
        assert!(pm.open_position(&pos).await.is_err());
    }

    #[tokio::test]
    async fn test_invalid_side_rejected() {
        let pm = PositionManager::in_memory().unwrap();
        let mut pos = make_position("bad", "BTC-PERP", "long", 1.0, 60000.0);
        pos.side = "sideways".to_string();
        assert!(pm.open_position(&pos).await.is_err());
    }

    // -- Execution metrics tests --

    fn make_metric(order_id: &str, status: &str, requested: f64, filled: f64) -> ExecutionMetric {
        ExecutionMetric {
            order_id: order_id.to_string(),
            market: "BTC-PERP".to_string(),
            side: "buy".to_string(),
            requested_size: requested,
            filled_size: filled,
            status: status.to_string(),
            mode: "paper".to_string(),
        }
    }

    #[tokio::test]
    async fn test_execution_summary_empty() {
        let pm = PositionManager::in_memory().unwrap();
        let summary = pm.get_execution_summary().await.unwrap();
        assert_eq!(summary.total_orders, 0);
        assert_eq!(summary.full_fills, 0);
        assert_eq!(summary.partial_fills, 0);
        assert_eq!(summary.zero_fills, 0);
        assert_eq!(summary.avg_fill_rate_pct, 0.0);
    }

    #[tokio::test]
    async fn test_record_and_summarize_executions() {
        let pm = PositionManager::in_memory().unwrap();

        // Full fill
        pm.record_execution(make_metric("o1", "filled", 1.0, 1.0))
            .await
            .unwrap();
        // Partial fill
        pm.record_execution(make_metric("o2", "partial_fill", 1.0, 0.5))
            .await
            .unwrap();
        // Zero fill (resting)
        pm.record_execution(make_metric("o3", "resting", 1.0, 0.0))
            .await
            .unwrap();

        let summary = pm.get_execution_summary().await.unwrap();
        assert_eq!(summary.total_orders, 3);
        assert_eq!(summary.full_fills, 1);
        assert_eq!(summary.partial_fills, 1);
        assert_eq!(summary.zero_fills, 1);
        // avg: (100 + 50 + 0) / 3 = 50.0
        assert!((summary.avg_fill_rate_pct - 50.0).abs() < 1e-6);
    }

    #[tokio::test]
    async fn test_all_full_fills() {
        let pm = PositionManager::in_memory().unwrap();
        pm.record_execution(make_metric("a", "filled", 2.0, 2.0))
            .await
            .unwrap();
        pm.record_execution(make_metric("b", "ok", 0.5, 0.5))
            .await
            .unwrap();

        let summary = pm.get_execution_summary().await.unwrap();
        assert_eq!(summary.total_orders, 2);
        assert_eq!(summary.full_fills, 2);
        assert_eq!(summary.partial_fills, 0);
        assert_eq!(summary.zero_fills, 0);
        assert!((summary.avg_fill_rate_pct - 100.0).abs() < 1e-6);
    }
}