//! CRYPTO namespace (0x0001) — signing, verification, hashing.

use async_trait::async_trait;
use sha2::{Digest, Sha256};

use crate::shellstream::{crypto, ShellstreamResponse};
use super::NamespaceHandler;

pub struct CryptoHandler;

impl CryptoHandler {
    pub fn new() -> Self {
        Self
    }
}

#[async_trait]
impl NamespaceHandler for CryptoHandler {
    async fn handle(
        &self,
        function_id: u16,
        payload: &[u8],
        session_id: &[u8; 16],
    ) -> ShellstreamResponse {
        match function_id {
            crypto::SIGN => self.sign(payload, session_id).await,
            crypto::VERIFY => self.verify(payload, session_id).await,
            // ENCRYPT (canonical ID 6) — used for hash in soft mode
            crypto::ENCRYPT => self.hash(payload, session_id).await,
            _ => ShellstreamResponse::error(
                *session_id,
                0,
                &format!("Unknown CRYPTO function: {function_id:#06x}"),
            ),
        }
    }
}

/// Extract bytes from an rmpv map field — handles both Binary and String values.
fn extract_bytes_from_map(map: &rmpv::Value, key: &str) -> Vec<u8> {
    let key_val = rmpv::Value::String(key.into());
    if let rmpv::Value::Map(entries) = map {
        for (k, v) in entries {
            if k == &key_val {
                return match v {
                    rmpv::Value::Binary(b) => b.clone(),
                    rmpv::Value::String(s) => s.as_bytes().to_vec(),
                    _ => vec![],
                };
            }
        }
    }
    vec![]
}

/// Extract a UTF-8 string from an rmpv map field — handles Binary (as UTF-8), String, or missing.
fn extract_string_from_map(map: &rmpv::Value, key: &str) -> String {
    let key_val = rmpv::Value::String(key.into());
    if let rmpv::Value::Map(entries) = map {
        for (k, v) in entries {
            if k == &key_val {
                return match v {
                    rmpv::Value::String(s) => s.as_str().unwrap_or_default().to_string(),
                    rmpv::Value::Binary(b) => String::from_utf8_lossy(b).to_string(),
                    _ => String::new(),
                };
            }
        }
    }
    String::new()
}

impl CryptoHandler {
    async fn sign(&self, payload: &[u8], session_id: &[u8; 16]) -> ShellstreamResponse {
        // Soft mode: HMAC-SHA256 with a dev key.
        // Use rmpv::Value to preserve binary data from Python SDK.
        let payload_parsed: rmpv::Value =
            rmp_serde::from_slice(payload).unwrap_or(rmpv::Value::Nil);

        let data = extract_bytes_from_map(&payload_parsed, "data");

        let mut hasher = Sha256::new();
        hasher.update(b"soft-signing-key:");
        hasher.update(&data);
        let signature = hex::encode(hasher.finalize());

        let response = rmp_serde::to_vec(&serde_json::json!({
            "signature": signature,
            "algorithm": "sha256-hmac-soft",
        }))
        .unwrap_or_default();
        ShellstreamResponse::ok(*session_id, 0, response)
    }

    async fn verify(&self, payload: &[u8], session_id: &[u8; 16]) -> ShellstreamResponse {
        let payload_parsed: rmpv::Value =
            rmp_serde::from_slice(payload).unwrap_or(rmpv::Value::Nil);

        let data = extract_bytes_from_map(&payload_parsed, "data");
        // Signature may arrive as String or Binary (Python SDK sends bytes)
        let sig = extract_string_from_map(&payload_parsed, "signature");

        // Recompute
        let mut hasher = Sha256::new();
        hasher.update(b"soft-signing-key:");
        hasher.update(&data);
        let expected = hex::encode(hasher.finalize());

        let valid = sig == expected;

        let response = rmp_serde::to_vec(&serde_json::json!({
            "valid": valid,
        }))
        .unwrap_or_default();
        ShellstreamResponse::ok(*session_id, 0, response)
    }

    async fn hash(&self, payload: &[u8], session_id: &[u8; 16]) -> ShellstreamResponse {
        let hash = hex::encode(Sha256::digest(payload));
        let response = rmp_serde::to_vec(&serde_json::json!({
            "hash": hash,
            "algorithm": "sha256",
        }))
        .unwrap_or_default();
        ShellstreamResponse::ok(*session_id, 0, response)
    }
}