//! HostIO event extraction and categorization.
//!
//! HostIO events represent calls from WASM to the Stylus VM runtime.
//! Common types: storage_read, storage_write, call, log, etc.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Type of HostIO operation
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum HostIoType {
    StorageLoad,
    StorageStore,
    Call,
    StaticCall,
    DelegateCall,
    Create,
    Log,
    SelfDestruct,
    AccountBalance,
    BlockHash,
    Other,
}

impl HostIoType {
    /// Parse HostIO type from string (from trace data)
    ///
    /// **Private** - only used internally during parsing
    fn from_str(s: &str) -> Self {
        match s.to_lowercase().as_str() {
            "storage_load" | "sload" => Self::StorageLoad,
            "storage_store" | "sstore" => Self::StorageStore,
            "call" => Self::Call,
            "staticcall" => Self::StaticCall,
            "delegatecall" => Self::DelegateCall,
            "create" | "create2" => Self::Create,
            "log" | "log0" | "log1" | "log2" | "log3" | "log4" => Self::Log,
            "selfdestruct" => Self::SelfDestruct,
            "balance" => Self::AccountBalance,
            "blockhash" => Self::BlockHash,
            _ => Self::Other,
        }
    }
}

/// A single HostIO event from the trace
#[derive(Debug, Clone)]
pub struct HostIoEvent {
    pub io_type: HostIoType,
    pub gas_cost: u64,
}

/// Aggregated HostIO statistics
#[derive(Debug, Clone)]
pub struct HostIoStats {
    counts: HashMap<HostIoType, u64>,
    total_gas: u64,
}

impl HostIoStats {
    /// Create new empty stats
    pub fn new() -> Self {
        Self {
            counts: HashMap::new(),
            total_gas: 0,
        }
    }

    /// Add a HostIO event to the statistics
    pub fn add_event(&mut self, event: HostIoEvent) {
        *self.counts.entry(event.io_type).or_insert(0) += 1;
        self.total_gas += event.gas_cost;
    }

    /// Get total number of HostIO calls
    pub fn total_calls(&self) -> u64 {
        self.counts.values().sum()
    }

    /// Get count for a specific HostIO type
    pub fn count_for_type(&self, io_type: HostIoType) -> u64 {
        self.counts.get(&io_type).copied().unwrap_or(0)
    }

    /// Get total gas consumed by HostIO
    pub fn total_gas(&self) -> u64 {
        self.total_gas
    }

    /// Convert to a map for JSON serialization
    pub fn to_map(&self) -> HashMap<String, u64> {
        self.counts
            .iter()
            .map(|(k, v)| (format!("{:?}", k), *v))
            .collect()
    }
}

impl Default for HostIoStats {
    fn default() -> Self {
        Self::new()
    }
}

/// Extract HostIO events from raw trace data
///
/// **Public** - used by the main parser to build statistics
///
/// # Arguments
/// * `trace_data` - Raw JSON from stylusTracer
///
/// # Returns
/// Parsed HostIO statistics
pub fn extract_hostio_events(trace_data: &serde_json::Value) -> HostIoStats {
    let mut stats = HostIoStats::new();
    
    // Try to extract HostIO array from trace
    // Actual field name depends on stylusTracer output format
    // This is a placeholder - adjust based on real trace format
    if let Some(hostio_array) = trace_data.get("hostio").and_then(|v| v.as_array()) {
        for event_json in hostio_array {
            if let Some(event) = parse_hostio_event(event_json) {
                stats.add_event(event);
            }
        }
    }
    
    stats
}

/// Parse a single HostIO event from JSON
///
/// **Private** - internal parsing logic
fn parse_hostio_event(event_json: &serde_json::Value) -> Option<HostIoEvent> {
    let io_type_str = event_json.get("type")?.as_str()?;
    let gas_cost = event_json.get("gas")?.as_u64()?;
    
    Some(HostIoEvent {
        io_type: HostIoType::from_str(io_type_str),
        gas_cost,
    })
}

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

    #[test]
    fn test_hostio_type_parsing() {
        assert_eq!(HostIoType::from_str("storage_load"), HostIoType::StorageLoad);
        assert_eq!(HostIoType::from_str("SSTORE"), HostIoType::StorageStore);
        assert_eq!(HostIoType::from_str("unknown"), HostIoType::Other);
    }

    #[test]
    fn test_hostio_stats() {
        let mut stats = HostIoStats::new();
        
        stats.add_event(HostIoEvent {
            io_type: HostIoType::StorageLoad,
            gas_cost: 100,
        });
        
        stats.add_event(HostIoEvent {
            io_type: HostIoType::StorageLoad,
            gas_cost: 200,
        });
        
        assert_eq!(stats.count_for_type(HostIoType::StorageLoad), 2);
        assert_eq!(stats.total_gas(), 300);
        assert_eq!(stats.total_calls(), 2);
    }
}