foxhunt/testing/api-load/src/clients/authenticated_client.rs

use anyhow::{Context, Result};
use chrono::{Duration, Utc};
use jsonwebtoken::{encode, EncodingKey, Header};
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::time::Instant;

use crate::metrics::{RequestMetric, RequestStatus, ServiceType};

#[derive(Debug, Serialize, Deserialize)]
struct Claims {
    sub: String,
    username: String,
    exp: i64,
    iat: i64,
}

pub struct AuthenticatedClient {
    client: Client,
    gateway_url: String,
    jwt_token: String,
}

impl AuthenticatedClient {
    pub async fn new(
        gateway_url: String,
        jwt_secret: &str,
        user_id: &str,
        username: &str,
    ) -> Result<Self> {
        let client = Client::builder()
            .timeout(std::time::Duration::from_secs(30))
            .pool_max_idle_per_host(10)
            .build()
            .context("Failed to create HTTP client")?;

        // Generate JWT token
        let claims = Claims {
            sub: user_id.to_string(),
            username: username.to_string(),
            iat: Utc::now().timestamp(),
            exp: (Utc::now() + Duration::hours(24)).timestamp(),
        };

        let jwt_token = encode(
            &Header::default(),
            &claims,
            &EncodingKey::from_secret(jwt_secret.as_bytes()),
        )
        .context("Failed to generate JWT token")?;

        Ok(Self {
            client,
            gateway_url,
            jwt_token,
        })
    }

    pub async fn submit_order(&self, client_id: usize, order: TestOrder) -> Result<RequestMetric> {
        let start = Instant::now();

        let result = self
            .client
            .post(format!("{}/trading/orders", self.gateway_url))
            .header("Authorization", format!("Bearer {}", self.jwt_token))
            .json(&order)
            .send()
            .await;

        let latency = start.elapsed();

        let status = match result {
            Ok(response) => {
                if response.status().is_success() {
                    RequestStatus::Success
                } else if response.status().as_u16() == 429 {
                    RequestStatus::RateLimited
                } else if response.status().as_u16() == 503 {
                    RequestStatus::CircuitBreakerOpen
                } else {
                    RequestStatus::Error
                }
            },
            Err(e) => {
                if e.is_timeout() {
                    RequestStatus::Timeout
                } else {
                    RequestStatus::Error
                }
            },
        };

        Ok(RequestMetric {
            timestamp: chrono::Utc::now(),
            client_id,
            service: ServiceType::Trading,
            latency,
            status,
            error_type: if status != RequestStatus::Success {
                Some(format!("{:?}", status))
            } else {
                None
            },
        })
    }

    pub async fn get_positions(&self, client_id: usize) -> Result<RequestMetric> {
        let start = Instant::now();

        let result = self
            .client
            .get(format!("{}/trading/positions", self.gateway_url))
            .header("Authorization", format!("Bearer {}", self.jwt_token))
            .send()
            .await;

        let latency = start.elapsed();

        let status = match result {
            Ok(response) => {
                if response.status().is_success() {
                    RequestStatus::Success
                } else if response.status().as_u16() == 429 {
                    RequestStatus::RateLimited
                } else if response.status().as_u16() == 503 {
                    RequestStatus::CircuitBreakerOpen
                } else {
                    RequestStatus::Error
                }
            },
            Err(e) => {
                if e.is_timeout() {
                    RequestStatus::Timeout
                } else {
                    RequestStatus::Error
                }
            },
        };

        Ok(RequestMetric {
            timestamp: chrono::Utc::now(),
            client_id,
            service: ServiceType::Trading,
            latency,
            status,
            error_type: if status != RequestStatus::Success {
                Some(format!("{:?}", status))
            } else {
                None
            },
        })
    }

    pub async fn run_backtest(
        &self,
        client_id: usize,
        config: BacktestConfig,
    ) -> Result<RequestMetric> {
        let start = Instant::now();

        let result = self
            .client
            .post(format!("{}/backtesting/run", self.gateway_url))
            .header("Authorization", format!("Bearer {}", self.jwt_token))
            .json(&config)
            .send()
            .await;

        let latency = start.elapsed();

        let status = match result {
            Ok(response) => {
                if response.status().is_success() {
                    RequestStatus::Success
                } else if response.status().as_u16() == 429 {
                    RequestStatus::RateLimited
                } else if response.status().as_u16() == 503 {
                    RequestStatus::CircuitBreakerOpen
                } else {
                    RequestStatus::Error
                }
            },
            Err(e) => {
                if e.is_timeout() {
                    RequestStatus::Timeout
                } else {
                    RequestStatus::Error
                }
            },
        };

        Ok(RequestMetric {
            timestamp: chrono::Utc::now(),
            client_id,
            service: ServiceType::Backtesting,
            latency,
            status,
            error_type: if status != RequestStatus::Success {
                Some(format!("{:?}", status))
            } else {
                None
            },
        })
    }

    pub async fn train_model(
        &self,
        client_id: usize,
        config: TrainingConfig,
    ) -> Result<RequestMetric> {
        let start = Instant::now();

        let result = self
            .client
            .post(format!("{}/ml/train", self.gateway_url))
            .header("Authorization", format!("Bearer {}", self.jwt_token))
            .json(&config)
            .send()
            .await;

        let latency = start.elapsed();

        let status = match result {
            Ok(response) => {
                if response.status().is_success() {
                    RequestStatus::Success
                } else if response.status().as_u16() == 429 {
                    RequestStatus::RateLimited
                } else if response.status().as_u16() == 503 {
                    RequestStatus::CircuitBreakerOpen
                } else {
                    RequestStatus::Error
                }
            },
            Err(e) => {
                if e.is_timeout() {
                    RequestStatus::Timeout
                } else {
                    RequestStatus::Error
                }
            },
        };

        Ok(RequestMetric {
            timestamp: chrono::Utc::now(),
            client_id,
            service: ServiceType::MlTraining,
            latency,
            status,
            error_type: if status != RequestStatus::Success {
                Some(format!("{:?}", status))
            } else {
                None
            },
        })
    }
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct TestOrder {
    pub symbol: String,
    pub quantity: f64,
    pub side: String,
    pub order_type: String,
}

impl TestOrder {
    pub fn random() -> Self {
        use rand::Rng;
        let mut rng = rand::thread_rng();

        let symbols = ["AAPL", "GOOGL", "MSFT", "AMZN", "TSLA"];
        let sides = ["buy", "sell"];
        let order_types = ["market", "limit"];

        Self {
            symbol: symbols
                .get(rng.gen_range(0..symbols.len()))
                .map_or_else(|| "AAPL".to_owned(), |s| (*s).to_string()),
            quantity: rng.gen_range(1.0..101.0),
            side: sides
                .get(rng.gen_range(0..sides.len()))
                .map_or_else(|| "buy".to_owned(), |s| (*s).to_string()),
            order_type: order_types
                .get(rng.gen_range(0..order_types.len()))
                .map_or_else(|| "market".to_owned(), |s| (*s).to_string()),
        }
    }
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct BacktestConfig {
    pub strategy: String,
    pub start_date: String,
    pub end_date: String,
}

impl BacktestConfig {
    pub fn default() -> Self {
        Self {
            strategy: "momentum".to_owned(),
            start_date: "2024-01-01".to_owned(),
            end_date: "2024-12-31".to_owned(),
        }
    }
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct TrainingConfig {
    pub model_type: String,
    pub epochs: u32,
}

impl TrainingConfig {
    pub fn default() -> Self {
        Self {
            model_type: "mamba2".to_owned(),
            epochs: 10_u32,
        }
    }
}