kaynine-core 0.1.0

use crate::ids::ToolCallId;
use serde::{Deserialize, Serialize};
use serde_json::Value;

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum ContentBlock {
    Text {
        text: String,
    },
    Image {
        media_type: String,
        data: BinaryRef,
    },
    Reasoning {
        text: String,
        provider_state: Option<Value>,
    },
    ToolCall {
        id: ToolCallId,
        name: String,
        arguments: Value,
    },
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BinaryRef {
    pub sha256: String,
    pub byte_len: u64,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum Message {
    User {
        blocks: Vec<ContentBlock>,
    },
    Assistant {
        blocks: Vec<ContentBlock>,
        finish_reason: FinishReason,
        truncated: bool,
    },
    ToolResult {
        results: Vec<ToolResultPayload>,
    },
    Summary {
        text: String,
    },
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolResultPayload {
    pub call_id: ToolCallId,
    pub is_error: bool,
    pub text: String,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum FinishReason {
    Stop,
    Length,
    ContentFilter,
    Other,
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct Usage {
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub cached_input_tokens: Option<u64>,
}

impl Message {
    pub fn tool_calls(&self) -> Vec<(ToolCallId, String, Value)> {
        match self {
            Message::Assistant { blocks, .. } => blocks
                .iter()
                .filter_map(|b| match b {
                    ContentBlock::ToolCall {
                        id,
                        name,
                        arguments,
                    } => Some((id.clone(), name.clone(), arguments.clone())),
                    _ => None,
                })
                .collect(),
            _ => Vec::new(),
        }
    }
}

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

    #[test]
    fn message_roundtrips_through_serde() {
        let msg = Message::Assistant {
            blocks: vec![
                ContentBlock::Reasoning {
                    text: "thinking".into(),
                    provider_state: None,
                },
                ContentBlock::ToolCall {
                    id: ToolCallId::from("call-1"),
                    name: "read_file".into(),
                    arguments: serde_json::json!({"path": "/tmp/a"}),
                },
            ],
            finish_reason: FinishReason::Stop,
            truncated: false,
        };
        let json = serde_json::to_string(&msg).unwrap();
        let back: Message = serde_json::from_str(&json).unwrap();
        assert_eq!(back, msg);
    }

    #[test]
    fn tool_call_extraction_preserves_order() {
        let msg = Message::Assistant {
            blocks: vec![
                ContentBlock::Text { text: "hi".into() },
                ContentBlock::ToolCall {
                    id: ToolCallId::from("c1"),
                    name: "a".into(),
                    arguments: serde_json::json!({}),
                },
                ContentBlock::ToolCall {
                    id: ToolCallId::from("c2"),
                    name: "b".into(),
                    arguments: serde_json::json!({}),
                },
            ],
            finish_reason: FinishReason::Stop,
            truncated: false,
        };
        let calls = msg.tool_calls();
        assert_eq!(calls.len(), 2);
        assert_eq!(calls[0].0, ToolCallId::from("c1"));
        assert_eq!(calls[1].1, "b");
    }
}