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

const MAX_ENVELOPE_BYTES: usize = 2_000_000;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentIdentityV1 {
    pub agent_id: String,
    pub agent_type: String,
    pub daemon_fingerprint: String,
    pub capabilities: Vec<String>,
}

impl AgentIdentityV1 {
    pub fn from_current(agent_id: &str, agent_type: &str) -> Self {
        Self {
            agent_id: agent_id.to_string(),
            agent_type: agent_type.to_string(),
            daemon_fingerprint: compute_daemon_fingerprint(),
            capabilities: default_capabilities(),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TransportEnvelopeV1 {
    pub format_version: u32,
    pub sent_at: DateTime<Utc>,
    pub sender: AgentIdentityV1,
    pub recipient: Option<String>,
    pub content_type: TransportContentType,
    pub payload_json: String,
    pub signature: Option<String>,
    pub metadata: HashMap<String, String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum TransportContentType {
    HandoffBundle,
    ContextPackage,
    A2AMessage,
    A2ATask,
}

impl TransportEnvelopeV1 {
    pub fn new(
        sender: AgentIdentityV1,
        recipient: Option<&str>,
        content_type: TransportContentType,
        payload_json: String,
    ) -> Self {
        Self {
            format_version: 1,
            sent_at: Utc::now(),
            sender,
            recipient: recipient.map(std::string::ToString::to_string),
            content_type,
            payload_json,
            signature: None,
            metadata: HashMap::new(),
        }
    }

    pub fn sign(&mut self, secret: &[u8]) {
        use hmac::{Hmac, Mac};
        use sha2::Sha256;

        let header = format!(
            "{}:{}:{}:{}",
            self.format_version,
            self.sender.agent_id,
            self.content_type_str(),
            self.payload_json.len()
        );
        let mut mac = Hmac::<Sha256>::new_from_slice(secret).expect("HMAC accepts any key length");
        mac.update(header.as_bytes());
        mac.update(b"\0");
        mac.update(self.payload_json.as_bytes());
        let result = mac.finalize().into_bytes();
        let mut sig = String::with_capacity(result.len() * 2);
        for b in &result {
            use std::fmt::Write;
            let _ = write!(sig, "{b:02x}");
        }
        self.signature = Some(sig);
    }

    pub fn verify_signature(&self, secret: &[u8]) -> bool {
        let Some(ref sig) = self.signature else {
            return false;
        };
        use hmac::{Hmac, Mac};
        use sha2::Sha256;

        let header = format!(
            "{}:{}:{}:{}",
            self.format_version,
            self.sender.agent_id,
            self.content_type_str(),
            self.payload_json.len()
        );
        let mut mac = Hmac::<Sha256>::new_from_slice(secret).expect("HMAC accepts any key length");
        mac.update(header.as_bytes());
        mac.update(b"\0");
        mac.update(self.payload_json.as_bytes());

        let expected: Vec<u8> = (0..sig.len())
            .step_by(2)
            .filter_map(|i| u8::from_str_radix(&sig[i..i + 2], 16).ok())
            .collect();
        if expected.len() != sig.len() / 2 {
            return false;
        }
        mac.verify_slice(&expected).is_ok()
    }

    fn content_type_str(&self) -> &str {
        match self.content_type {
            TransportContentType::HandoffBundle => "handoff_bundle",
            TransportContentType::ContextPackage => "context_package",
            TransportContentType::A2AMessage => "a2a_message",
            TransportContentType::A2ATask => "a2a_task",
        }
    }
}

pub fn serialize_envelope(envelope: &TransportEnvelopeV1) -> Result<String, String> {
    let json = serde_json::to_string_pretty(envelope).map_err(|e| e.to_string())?;
    if json.len() > MAX_ENVELOPE_BYTES {
        return Err(format!(
            "envelope too large ({} bytes, max {})",
            json.len(),
            MAX_ENVELOPE_BYTES
        ));
    }
    Ok(json)
}

pub fn parse_envelope(json: &str) -> Result<TransportEnvelopeV1, String> {
    if json.len() > MAX_ENVELOPE_BYTES {
        return Err(format!(
            "envelope too large ({} bytes, max {})",
            json.len(),
            MAX_ENVELOPE_BYTES
        ));
    }
    let env: TransportEnvelopeV1 = serde_json::from_str(json).map_err(|e| e.to_string())?;
    if env.format_version != 1 {
        return Err(format!(
            "unsupported format_version {} (expected 1)",
            env.format_version
        ));
    }
    Ok(env)
}

fn compute_daemon_fingerprint() -> String {
    use sha2::{Digest, Sha256};
    let mut hasher = Sha256::new();
    hasher.update(env!("CARGO_PKG_VERSION").as_bytes());
    if let Ok(exe) = std::env::current_exe() {
        hasher.update(exe.to_string_lossy().as_bytes());
    }
    format!("{:x}", hasher.finalize())[..16].to_string()
}

fn default_capabilities() -> Vec<String> {
    vec![
        "context_compression".to_string(),
        "knowledge_graph".to_string(),
        "shared_sessions".to_string(),
        "a2a_messaging".to_string(),
    ]
}

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

    fn test_sender() -> AgentIdentityV1 {
        AgentIdentityV1 {
            agent_id: "test-agent".to_string(),
            agent_type: "cursor".to_string(),
            daemon_fingerprint: "abcd1234".to_string(),
            capabilities: vec!["context_compression".to_string()],
        }
    }

    #[test]
    fn envelope_roundtrip() {
        let env = TransportEnvelopeV1::new(
            test_sender(),
            Some("target-agent"),
            TransportContentType::A2AMessage,
            r#"{"hello":"world"}"#.to_string(),
        );
        let json = serialize_envelope(&env).unwrap();
        let parsed = parse_envelope(&json).unwrap();
        assert_eq!(parsed.format_version, 1);
        assert_eq!(parsed.sender.agent_id, "test-agent");
        assert_eq!(parsed.recipient, Some("target-agent".to_string()));
        assert_eq!(parsed.content_type, TransportContentType::A2AMessage);
    }

    #[test]
    fn hmac_sign_verify() {
        let secret = b"test-secret-key";
        let mut env = TransportEnvelopeV1::new(
            test_sender(),
            None,
            TransportContentType::HandoffBundle,
            "payload".to_string(),
        );
        assert!(!env.verify_signature(secret));

        env.sign(secret);
        assert!(env.signature.is_some());
        assert!(env.verify_signature(secret));
        assert!(!env.verify_signature(b"wrong-key"));
    }

    #[test]
    fn rejects_oversized_envelope() {
        let big = "x".repeat(MAX_ENVELOPE_BYTES + 1);
        assert!(parse_envelope(&big).is_err());
    }

    #[test]
    fn rejects_wrong_version() {
        let json = r#"{"format_version":99,"sent_at":"2026-01-01T00:00:00Z","sender":{"agent_id":"a","agent_type":"b","daemon_fingerprint":"c","capabilities":[]},"recipient":null,"content_type":"a2a_message","payload_json":"{}","signature":null,"metadata":{}}"#;
        assert!(parse_envelope(json).is_err());
    }
}
