use serde::{Deserialize, Serialize};
use serde_json::Value;

// ------------------------------------------------------------------
// AxonMethod
// ------------------------------------------------------------------

/// Typed method discriminant for Axon protocol messages.
///
/// Replaces the raw `method: String` field — callers that need a string
/// representation can use [`AxonMethod::as_str`] or the `Display` impl.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AxonMethod {
    /// Invoke a named tool.
    ToolCall,
    /// Return value from a tool invocation.
    ToolResult,
    /// Deliver a typed Axon event envelope.
    Event,
    /// Liveness check.
    Ping,
    /// Liveness reply.
    Pong,
    /// Subscribe to an event stream.
    Subscribe,
    /// Cancel a subscription.
    Unsubscribe,
}

impl AxonMethod {
    /// Return the canonical lowercase string for this method.
    pub fn as_str(self) -> &'static str {
        match self {
            Self::ToolCall => "tool_call",
            Self::ToolResult => "tool_result",
            Self::Event => "event",
            Self::Ping => "ping",
            Self::Pong => "pong",
            Self::Subscribe => "subscribe",
            Self::Unsubscribe => "unsubscribe",
        }
    }

    /// Parse a string into an `AxonMethod`, returning `None` on mismatch.
    pub fn from_str(s: &str) -> Option<Self> {
        match s {
            "tool_call" => Some(Self::ToolCall),
            "tool_result" => Some(Self::ToolResult),
            "event" => Some(Self::Event),
            "ping" => Some(Self::Ping),
            "pong" => Some(Self::Pong),
            "subscribe" => Some(Self::Subscribe),
            "unsubscribe" => Some(Self::Unsubscribe),
            _ => None,
        }
    }
}

impl std::fmt::Display for AxonMethod {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.as_str())
    }
}

// ------------------------------------------------------------------
// AxonMessage
// ------------------------------------------------------------------

/// An Axon protocol message — wraps a tool call, event delivery, or
/// other method invocation.
///
/// The `method` field is now typed as [`AxonMethod`].  For backward
/// compatibility the raw string representation is available via
/// [`AxonMessage::method_str`].
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AxonMessage {
    /// Unique message ID.
    pub id: String,
    /// Typed method discriminant.
    pub method: AxonMethod,
    /// Method-specific parameters.
    pub params: Value,
}

impl AxonMessage {
    /// Return the canonical string representation of the method (e.g.
    /// `"tool_call"`, `"ping"`).  Useful when serialising to a transport
    /// that expects a raw string rather than a tagged enum.
    pub fn method_str(&self) -> &'static str {
        self.method.as_str()
    }

    /// Construct a ping message.
    pub fn ping(id: impl Into<String>) -> Self {
        Self {
            id: id.into(),
            method: AxonMethod::Ping,
            params: Value::Object(Default::default()),
        }
    }
}

// ------------------------------------------------------------------
// AxonResponse
// ------------------------------------------------------------------

/// An Axon protocol response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AxonResponse {
    /// Echo of the request ID.
    pub id: String,
    /// Whether the call succeeded.
    pub success: bool,
    /// The result payload (on success) or error message (on failure).
    pub result: Value,
}

impl AxonResponse {
    pub fn ok(id: impl Into<String>, result: Value) -> Self {
        Self {
            id: id.into(),
            success: true,
            result,
        }
    }

    pub fn err(id: impl Into<String>, message: impl Into<String>) -> Self {
        Self {
            id: id.into(),
            success: false,
            result: serde_json::json!({ "error": message.into() }),
        }
    }
}

// ------------------------------------------------------------------
// AxonEvent  (legacy SSE event — kept for backward compat)
// ------------------------------------------------------------------

/// An Axon SSE event — sent over the streaming channel.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AxonEvent {
    pub event_type: String,
    pub payload: Value,
    pub timestamp: i64,
}

impl AxonEvent {
    pub fn new(event_type: impl Into<String>, payload: Value) -> Self {
        Self {
            event_type: event_type.into(),
            payload,
            timestamp: neuron_domain::now_ms(),
        }
    }
}

// ------------------------------------------------------------------
// Tests
// ------------------------------------------------------------------

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

    #[test]
    fn axon_response_ok() {
        let r = AxonResponse::ok("req-1", serde_json::json!({"id": "abc"}));
        assert!(r.success);
        assert_eq!(r.id, "req-1");
    }

    #[test]
    fn axon_response_err() {
        let r = AxonResponse::err("req-2", "not found");
        assert!(!r.success);
        assert_eq!(r.result["error"], "not found");
    }

    #[test]
    fn axon_message_roundtrip_with_method_enum() {
        let msg = AxonMessage {
            id: "m1".into(),
            method: AxonMethod::ToolCall,
            params: serde_json::json!({ "tool": "remember", "content": "hello" }),
        };
        let json = serde_json::to_string(&msg).unwrap();
        let back: AxonMessage = serde_json::from_str(&json).unwrap();
        assert_eq!(back.method, AxonMethod::ToolCall);
        assert_eq!(back.method_str(), "tool_call");
    }

    #[test]
    fn axon_method_as_str_all_variants() {
        let cases = [
            (AxonMethod::ToolCall, "tool_call"),
            (AxonMethod::ToolResult, "tool_result"),
            (AxonMethod::Event, "event"),
            (AxonMethod::Ping, "ping"),
            (AxonMethod::Pong, "pong"),
            (AxonMethod::Subscribe, "subscribe"),
            (AxonMethod::Unsubscribe, "unsubscribe"),
        ];
        for (method, expected) in cases {
            assert_eq!(method.as_str(), expected);
            assert_eq!(method.to_string(), expected);
        }
    }

    #[test]
    fn axon_method_from_str_roundtrip() {
        let variants = [
            AxonMethod::ToolCall,
            AxonMethod::ToolResult,
            AxonMethod::Event,
            AxonMethod::Ping,
            AxonMethod::Pong,
            AxonMethod::Subscribe,
            AxonMethod::Unsubscribe,
        ];
        for v in variants {
            let s = v.as_str();
            assert_eq!(AxonMethod::from_str(s), Some(v));
        }
    }

    #[test]
    fn axon_method_from_str_unknown() {
        assert_eq!(AxonMethod::from_str("mystery"), None);
        assert_eq!(AxonMethod::from_str(""), None);
    }

    #[test]
    fn axon_method_serde_roundtrip() {
        let m = AxonMethod::Subscribe;
        let json = serde_json::to_string(&m).unwrap();
        let back: AxonMethod = serde_json::from_str(&json).unwrap();
        assert_eq!(back, AxonMethod::Subscribe);
    }

    #[test]
    fn axon_message_ping_helper() {
        let msg = AxonMessage::ping("p-1");
        assert_eq!(msg.method, AxonMethod::Ping);
        assert_eq!(msg.id, "p-1");
    }

    #[test]
    fn axon_method_all_variants_have_unique_str() {
        use std::collections::HashSet;
        let strs: HashSet<&str> = [
            AxonMethod::ToolCall,
            AxonMethod::ToolResult,
            AxonMethod::Event,
            AxonMethod::Ping,
            AxonMethod::Pong,
            AxonMethod::Subscribe,
            AxonMethod::Unsubscribe,
        ]
        .iter()
        .map(|m| m.as_str())
        .collect();
        assert_eq!(strs.len(), 7);
    }
}