use crate::PROTOCOL;
use crate::message::Envelope;
use anyhow::{Context, Result, anyhow};
use chacha20poly1305::aead::{Aead, KeyInit};
use chacha20poly1305::{ChaCha20Poly1305, Key, Nonce};
use ed25519_dalek::{Signature, Signer, SigningKey, VerifyingKey};
use hkdf::Hkdf;
use rand::RngCore;
use rand::rngs::OsRng;
use serde_json::Value;
use sha2::Sha256;
use x25519_dalek::{EphemeralSecret, PublicKey as X25519PublicKey, StaticSecret};

pub const ENC_ALG: &str = "x25519-hkdf-sha256-chacha20poly1305";

pub struct Keypair {
    signing: SigningKey,
}

impl Keypair {
    pub fn generate() -> Self {
        let mut bytes = [0u8; 32];
        rand::thread_rng().fill_bytes(&mut bytes);
        Self {
            signing: SigningKey::from_bytes(&bytes),
        }
    }

    pub fn from_hex(s: &str) -> Result<Self> {
        let raw = hex::decode(s.trim()).context("key is not hex")?;
        let bytes: [u8; 32] = raw
            .as_slice()
            .try_into()
            .map_err(|_| anyhow!("key must be 32 bytes"))?;
        Ok(Self {
            signing: SigningKey::from_bytes(&bytes),
        })
    }

    pub fn to_hex(&self) -> String {
        hex::encode(self.signing.to_bytes())
    }

    pub fn public_hex(&self) -> String {
        hex::encode(self.signing.verifying_key().to_bytes())
    }

    pub fn verifying_key(&self) -> VerifyingKey {
        self.signing.verifying_key()
    }

    pub fn sign(&self, bytes: &[u8]) -> String {
        hex::encode(self.signing.sign(bytes).to_bytes())
    }
}

pub fn random_nonce() -> String {
    let mut bytes = [0u8; 16];
    rand::thread_rng().fill_bytes(&mut bytes);
    hex::encode(bytes)
}

pub fn random_id() -> String {
    let mut bytes = [0u8; 8];
    rand::thread_rng().fill_bytes(&mut bytes);
    hex::encode(bytes)
}

pub fn canonical_json(value: &Value) -> String {
    let mut out = String::new();
    write_canonical(value, &mut out);
    out
}

fn write_canonical(value: &Value, out: &mut String) {
    match value {
        Value::Null => out.push_str("null"),
        Value::Bool(b) => out.push_str(if *b { "true" } else { "false" }),
        Value::Number(n) => out.push_str(&n.to_string()),
        Value::String(s) => out.push_str(&serde_json::to_string(s).expect("string serializes")),
        Value::Array(items) => {
            out.push('[');
            for (i, item) in items.iter().enumerate() {
                if i > 0 {
                    out.push(',');
                }
                write_canonical(item, out);
            }
            out.push(']');
        }
        Value::Object(map) => {
            let mut keys: Vec<&String> = map.keys().collect();
            keys.sort();
            out.push('{');
            for (i, key) in keys.iter().enumerate() {
                if i > 0 {
                    out.push(',');
                }
                out.push_str(&serde_json::to_string(key).expect("key serializes"));
                out.push(':');
                write_canonical(&map[*key], out);
            }
            out.push('}');
        }
    }
}

pub fn poll_signing_bytes(member: &str, nonce: &str) -> Vec<u8> {
    format!("{}\npoll\n{}\n{}", PROTOCOL, member, nonce).into_bytes()
}

pub fn verify_signature(public_hex: &str, message: &[u8], sig_hex: &str) -> Result<()> {
    let key_bytes = hex::decode(public_hex).context("public key is not hex")?;
    let key_bytes: [u8; 32] = key_bytes
        .as_slice()
        .try_into()
        .map_err(|_| anyhow!("public key must be 32 bytes"))?;
    let verifying_key =
        VerifyingKey::from_bytes(&key_bytes).map_err(|e| anyhow!("bad public key: {e}"))?;
    let sig_bytes = hex::decode(sig_hex).context("signature is not hex")?;
    let sig_bytes: [u8; 64] = sig_bytes
        .as_slice()
        .try_into()
        .map_err(|_| anyhow!("signature must be 64 bytes"))?;
    let signature = Signature::from_bytes(&sig_bytes);
    verifying_key
        .verify_strict(message, &signature)
        .map_err(|_| anyhow!("signature verification failed"))
}

pub fn signing_bytes(envelope: &Envelope) -> Vec<u8> {
    let unsigned = serde_json::json!({
        "type": envelope.msg_type,
        "from": envelope.from,
        "key": envelope.key,
        "ts": envelope.ts,
        "nonce": envelope.nonce,
        "body": envelope.body,
    });
    format!("{}\n{}", PROTOCOL, canonical_json(&unsigned)).into_bytes()
}

pub fn verify_envelope(envelope: &Envelope) -> Result<()> {
    verify_signature(&envelope.key, &signing_bytes(envelope), &envelope.sig)
}

pub fn generate_enc_keypair() -> (String, String) {
    let secret = StaticSecret::random_from_rng(OsRng);
    let public = X25519PublicKey::from(&secret);
    (
        hex::encode(secret.to_bytes()),
        hex::encode(public.to_bytes()),
    )
}

pub fn enc_public_from_secret(secret_hex: &str) -> Result<String> {
    let secret = enc_secret_from_hex(secret_hex)?;
    Ok(hex::encode(X25519PublicKey::from(&secret).to_bytes()))
}

fn enc_secret_from_hex(secret_hex: &str) -> Result<StaticSecret> {
    let bytes = hex::decode(secret_hex.trim()).context("enc key is not hex")?;
    let bytes: [u8; 32] = bytes
        .as_slice()
        .try_into()
        .map_err(|_| anyhow!("enc key must be 32 bytes"))?;
    Ok(StaticSecret::from(bytes))
}

fn enc_public_from_hex(public_hex: &str) -> Result<X25519PublicKey> {
    let bytes = hex::decode(public_hex.trim()).context("enc public key is not hex")?;
    let bytes: [u8; 32] = bytes
        .as_slice()
        .try_into()
        .map_err(|_| anyhow!("enc public key must be 32 bytes"))?;
    Ok(X25519PublicKey::from(bytes))
}

pub fn unicast_context(from: &str, recipient_enc_public: &str) -> String {
    format!("{} unicast {} {}", PROTOCOL, from, recipient_enc_public)
}

pub fn encrypt_unicast(recipient_enc_hex: &str, context: &str, plaintext: &[u8]) -> Result<Value> {
    let recipient = enc_public_from_hex(recipient_enc_hex)?;
    let ephemeral = EphemeralSecret::random_from_rng(OsRng);
    let epk = X25519PublicKey::from(&ephemeral);
    let shared = ephemeral.diffie_hellman(&recipient);
    let mut key = [0u8; 32];
    Hkdf::<Sha256>::new(None, shared.as_bytes())
        .expand(context.as_bytes(), &mut key)
        .map_err(|_| anyhow!("hkdf expand failed"))?;
    let cipher = ChaCha20Poly1305::new(Key::from_slice(&key));
    let mut nonce_bytes = [0u8; 12];
    rand::thread_rng().fill_bytes(&mut nonce_bytes);
    let ciphertext = cipher
        .encrypt(Nonce::from_slice(&nonce_bytes), plaintext)
        .map_err(|_| anyhow!("encryption failed"))?;
    let mut payload = nonce_bytes.to_vec();
    payload.extend_from_slice(&ciphertext);
    Ok(serde_json::json!({
        "enc": {
            "alg": ENC_ALG,
            "epk": hex::encode(epk.to_bytes()),
            "ct": hex::encode(payload),
        }
    }))
}

pub fn decrypt_unicast(secret_hex: &str, context: &str, enc: &Value) -> Result<Vec<u8>> {
    let secret = enc_secret_from_hex(secret_hex)?;
    let alg = enc
        .get("alg")
        .and_then(Value::as_str)
        .ok_or_else(|| anyhow!("missing algorithm"))?;
    if alg != ENC_ALG {
        return Err(anyhow!("unsupported encryption algorithm"));
    }
    let epk = enc
        .get("epk")
        .and_then(Value::as_str)
        .ok_or_else(|| anyhow!("missing ephemeral key"))?;
    let epk = enc_public_from_hex(epk)?;
    let ct = enc
        .get("ct")
        .and_then(Value::as_str)
        .ok_or_else(|| anyhow!("missing ciphertext"))?;
    let payload = hex::decode(ct).context("ciphertext is not hex")?;
    if payload.len() < 12 {
        return Err(anyhow!("ciphertext too short"));
    }
    let (nonce_bytes, ciphertext) = payload.split_at(12);
    let shared = secret.diffie_hellman(&epk);
    let mut key = [0u8; 32];
    Hkdf::<Sha256>::new(None, shared.as_bytes())
        .expand(context.as_bytes(), &mut key)
        .map_err(|_| anyhow!("hkdf expand failed"))?;
    let cipher = ChaCha20Poly1305::new(Key::from_slice(&key));
    cipher
        .decrypt(Nonce::from_slice(nonce_bytes), ciphertext)
        .map_err(|_| anyhow!("decryption failed"))
}

pub fn now_ts() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0)
}

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

    #[test]
    fn canonical_json_sorts_keys_recursively() {
        let a = json!({"b": 1, "a": {"d": [3, 2], "c": "x"}});
        let b = json!({"a": {"c": "x", "d": [3, 2]}, "b": 1});
        assert_eq!(canonical_json(&a), canonical_json(&b));
        assert_eq!(canonical_json(&a), r#"{"a":{"c":"x","d":[3,2]},"b":1}"#);
    }

    #[test]
    fn envelope_sign_verify_roundtrip() {
        let key = Keypair::generate();
        let env = Envelope::new(
            &key,
            "alice.example",
            crate::message::TYPE_QUERY,
            json!({"text": "hello"}),
        );
        verify_envelope(&env).unwrap();
    }

    #[test]
    fn unicast_encryption_roundtrips_and_rejects_tampering() {
        let (secret, public) = generate_enc_keypair();
        let context = unicast_context("alice.example", &public);
        let plaintext = b"{\"qid\":\"q1\",\"results\":[]}";
        let envelope_body = encrypt_unicast(&public, &context, plaintext).unwrap();
        assert_eq!(
            envelope_body.pointer("/enc/alg").and_then(Value::as_str),
            Some(ENC_ALG)
        );
        let enc = envelope_body.get("enc").unwrap();
        let decrypted = decrypt_unicast(&secret, &context, enc).unwrap();
        assert_eq!(decrypted, plaintext);

        let (other_secret, _) = generate_enc_keypair();
        assert!(decrypt_unicast(&other_secret, &context, enc).is_err());

        let mut tampered = enc.clone();
        let ct = tampered.get("ct").and_then(Value::as_str).unwrap();
        let mut bytes = hex::decode(ct).unwrap();
        bytes[13] ^= 0xff;
        tampered["ct"] = json!(hex::encode(bytes));
        assert!(decrypt_unicast(&secret, &context, &tampered).is_err());
    }

    #[test]
    fn tampered_body_fails_verification() {
        let key = Keypair::generate();
        let mut env = Envelope::new(
            &key,
            "alice.example",
            crate::message::TYPE_QUERY,
            json!({"text": "hello"}),
        );
        env.body = json!({"text": "hello", "extra": true});
        assert!(verify_envelope(&env).is_err());
    }
}