//! Trace Serialization (ENT-108)
//!
//! Binary and JSON serialization for decision traces.

use super::path::DecisionPath;
use super::trace::DecisionTrace;
use serde::{Deserialize, Serialize};

/// Trace format for serialization
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum TraceFormat {
    /// APRT binary format (compact, fast)
    ///
    /// Header: 0x41505254 ("APRT"), version byte, path type
    Binary,

    /// Pretty-printed JSON
    Json,

    /// JSON Lines (one JSON per line)
    JsonLines,
}

/// Magic bytes for APRT binary format
pub const APRT_MAGIC: [u8; 4] = [0x41, 0x50, 0x52, 0x54]; // "APRT"

/// Current format version
pub const APRT_VERSION: u8 = 1;

/// Path type identifiers
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum PathType {
    Linear = 0,
    Tree = 1,
    Forest = 2,
    KNN = 3,
    Neural = 4,
    Custom = 255,
}

impl From<u8> for PathType {
    fn from(value: u8) -> Self {
        match value {
            0 => PathType::Linear,
            1 => PathType::Tree,
            2 => PathType::Forest,
            3 => PathType::KNN,
            4 => PathType::Neural,
            _ => PathType::Custom,
        }
    }
}

/// Trace serializer
pub struct TraceSerializer {
    format: TraceFormat,
}

impl TraceSerializer {
    /// Create a new serializer
    pub fn new(format: TraceFormat) -> Self {
        Self { format }
    }

    /// Get the format
    pub fn format(&self) -> TraceFormat {
        self.format
    }

    /// Serialize a trace to bytes
    pub fn serialize<P: DecisionPath + Serialize>(
        &self,
        trace: &DecisionTrace<P>,
        path_type: PathType,
    ) -> Result<Vec<u8>, SerializationError> {
        match self.format {
            TraceFormat::Binary => self.serialize_binary(trace, path_type),
            TraceFormat::Json => self.serialize_json(trace),
            TraceFormat::JsonLines => self.serialize_json_line(trace),
        }
    }

    /// Serialize to APRT binary format
    fn serialize_binary<P: DecisionPath + Serialize>(
        &self,
        trace: &DecisionTrace<P>,
        path_type: PathType,
    ) -> Result<Vec<u8>, SerializationError> {
        let trace_bytes = trace.to_bytes();

        let mut bytes = Vec::with_capacity(8 + trace_bytes.len());

        // Header
        bytes.extend_from_slice(&APRT_MAGIC);
        bytes.push(APRT_VERSION);
        bytes.push(path_type as u8);
        bytes.extend_from_slice(&[0, 0]); // Reserved (alignment)

        // Trace data
        bytes.extend_from_slice(&trace_bytes);

        Ok(bytes)
    }

    /// Serialize to JSON
    fn serialize_json<P: DecisionPath + Serialize>(
        &self,
        trace: &DecisionTrace<P>,
    ) -> Result<Vec<u8>, SerializationError> {
        serde_json::to_vec_pretty(trace).map_err(SerializationError::Json)
    }

    /// Serialize to JSON line (no trailing newline)
    fn serialize_json_line<P: DecisionPath + Serialize>(
        &self,
        trace: &DecisionTrace<P>,
    ) -> Result<Vec<u8>, SerializationError> {
        let mut bytes = serde_json::to_vec(trace).map_err(SerializationError::Json)?;
        bytes.push(b'\n');
        Ok(bytes)
    }

    /// Deserialize from bytes
    pub fn deserialize<P: DecisionPath + for<'de> Deserialize<'de>>(
        &self,
        bytes: &[u8],
    ) -> Result<DecisionTrace<P>, SerializationError> {
        match self.format {
            TraceFormat::Binary => self.deserialize_binary(bytes),
            TraceFormat::Json | TraceFormat::JsonLines => self.deserialize_json(bytes),
        }
    }

    /// Deserialize from APRT binary format
    fn deserialize_binary<P: DecisionPath>(
        &self,
        bytes: &[u8],
    ) -> Result<DecisionTrace<P>, SerializationError> {
        if bytes.len() < 8 {
            return Err(SerializationError::InvalidFormat(
                "Insufficient header bytes".to_string(),
            ));
        }

        // Check magic
        if bytes[0..4] != APRT_MAGIC {
            return Err(SerializationError::InvalidFormat(
                "Invalid APRT magic bytes".to_string(),
            ));
        }

        // Check version
        let version = bytes[4];
        if version != APRT_VERSION {
            return Err(SerializationError::VersionMismatch {
                expected: APRT_VERSION,
                actual: version,
            });
        }

        // Skip path type and reserved bytes (bytes[5], bytes[6], bytes[7])

        // Deserialize trace
        DecisionTrace::from_bytes(&bytes[8..]).map_err(|e| {
            SerializationError::InvalidFormat(format!("Path deserialization failed: {e}"))
        })
    }

    /// Deserialize from JSON
    fn deserialize_json<P: DecisionPath + for<'de> Deserialize<'de>>(
        &self,
        bytes: &[u8],
    ) -> Result<DecisionTrace<P>, SerializationError> {
        serde_json::from_slice(bytes).map_err(SerializationError::Json)
    }
}

impl Default for TraceSerializer {
    fn default() -> Self {
        Self::new(TraceFormat::Binary)
    }
}

/// Serialization error
#[derive(Debug)]
pub enum SerializationError {
    /// Invalid format
    InvalidFormat(String),
    /// Version mismatch
    VersionMismatch { expected: u8, actual: u8 },
    /// JSON error
    Json(serde_json::Error),
    /// IO error
    Io(std::io::Error),
}

impl std::fmt::Display for SerializationError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SerializationError::InvalidFormat(msg) => write!(f, "Invalid format: {msg}"),
            SerializationError::VersionMismatch { expected, actual } => {
                write!(f, "Version mismatch: expected {expected}, got {actual}")
            }
            SerializationError::Json(e) => write!(f, "JSON error: {e}"),
            SerializationError::Io(e) => write!(f, "IO error: {e}"),
        }
    }
}

impl std::error::Error for SerializationError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            SerializationError::Json(e) => Some(e),
            SerializationError::Io(e) => Some(e),
            _ => None,
        }
    }
}

impl From<serde_json::Error> for SerializationError {
    fn from(e: serde_json::Error) -> Self {
        SerializationError::Json(e)
    }
}

impl From<std::io::Error> for SerializationError {
    fn from(e: std::io::Error) -> Self {
        SerializationError::Io(e)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::monitor::inference::path::LinearPath;

    fn make_test_trace() -> DecisionTrace<LinearPath> {
        let path = LinearPath::new(vec![0.5, -0.3, 0.2], 0.1, 0.5, 0.87).with_probability(0.87);
        DecisionTrace::new(1_000_000, 42, 0xdeadbeef, path, 0.87, 500)
    }

    #[test]
    fn test_path_type_from_u8() {
        assert_eq!(PathType::from(0), PathType::Linear);
        assert_eq!(PathType::from(1), PathType::Tree);
        assert_eq!(PathType::from(4), PathType::Neural);
        assert_eq!(PathType::from(100), PathType::Custom);
    }

    #[test]
    fn test_serialize_binary() {
        let serializer = TraceSerializer::new(TraceFormat::Binary);
        let trace = make_test_trace();

        let bytes = serializer
            .serialize(&trace, PathType::Linear)
            .expect("Serialization failed");

        // Check header
        assert_eq!(&bytes[0..4], &APRT_MAGIC);
        assert_eq!(bytes[4], APRT_VERSION);
        assert_eq!(bytes[5], PathType::Linear as u8);
    }

    #[test]
    fn test_binary_roundtrip() {
        let serializer = TraceSerializer::new(TraceFormat::Binary);
        let trace = make_test_trace();

        let bytes = serializer
            .serialize(&trace, PathType::Linear)
            .expect("Serialization failed");

        let restored: DecisionTrace<LinearPath> = serializer
            .deserialize(&bytes)
            .expect("Deserialization failed");

        assert_eq!(trace.timestamp_ns, restored.timestamp_ns);
        assert_eq!(trace.sequence, restored.sequence);
        assert_eq!(trace.input_hash, restored.input_hash);
        assert!((trace.output - restored.output).abs() < 1e-6);
    }

    #[test]
    fn test_serialize_json() {
        let serializer = TraceSerializer::new(TraceFormat::Json);
        let trace = make_test_trace();

        let bytes = serializer
            .serialize(&trace, PathType::Linear)
            .expect("Serialization failed");

        let json = String::from_utf8(bytes).expect("Invalid UTF-8");
        assert!(json.contains("timestamp_ns"));
        assert!(json.contains("sequence"));
        assert!(json.contains("path"));
    }

    #[test]
    fn test_json_roundtrip() {
        let serializer = TraceSerializer::new(TraceFormat::Json);
        let trace = make_test_trace();

        let bytes = serializer
            .serialize(&trace, PathType::Linear)
            .expect("Serialization failed");

        let restored: DecisionTrace<LinearPath> = serializer
            .deserialize(&bytes)
            .expect("Deserialization failed");

        assert_eq!(trace.sequence, restored.sequence);
    }

    #[test]
    fn test_serialize_json_lines() {
        let serializer = TraceSerializer::new(TraceFormat::JsonLines);
        let trace = make_test_trace();

        let bytes = serializer
            .serialize(&trace, PathType::Linear)
            .expect("Serialization failed");

        // Should end with newline
        assert_eq!(bytes.last(), Some(&b'\n'));

        // Should be single line
        let json = String::from_utf8(bytes).expect("Invalid UTF-8");
        assert_eq!(json.lines().count(), 1);
    }

    #[test]
    fn test_deserialize_invalid_magic() {
        let serializer = TraceSerializer::new(TraceFormat::Binary);
        let bytes = vec![0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00];

        let result: Result<DecisionTrace<LinearPath>, _> = serializer.deserialize(&bytes);
        assert!(matches!(result, Err(SerializationError::InvalidFormat(_))));
    }

    #[test]
    fn test_deserialize_version_mismatch() {
        let serializer = TraceSerializer::new(TraceFormat::Binary);
        let mut bytes = Vec::new();
        bytes.extend_from_slice(&APRT_MAGIC);
        bytes.push(99); // Invalid version
        bytes.extend_from_slice(&[0, 0, 0]);

        let result: Result<DecisionTrace<LinearPath>, _> = serializer.deserialize(&bytes);
        assert!(matches!(
            result,
            Err(SerializationError::VersionMismatch { .. })
        ));
    }

    #[test]
    fn test_deserialize_insufficient_data() {
        let serializer = TraceSerializer::new(TraceFormat::Binary);
        let bytes = vec![0x41, 0x50]; // Only 2 bytes

        let result: Result<DecisionTrace<LinearPath>, _> = serializer.deserialize(&bytes);
        assert!(matches!(result, Err(SerializationError::InvalidFormat(_))));
    }

    #[test]
    fn test_serializer_default() {
        let serializer = TraceSerializer::default();
        assert_eq!(serializer.format(), TraceFormat::Binary);
    }

    #[test]
    fn test_error_display() {
        let err = SerializationError::InvalidFormat("test".to_string());
        assert!(err.to_string().contains("Invalid format"));

        let err = SerializationError::VersionMismatch {
            expected: 1,
            actual: 2,
        };
        assert!(err.to_string().contains("Version mismatch"));
    }
}
