//! QR-based mobile wallet pairing using the Polkadot App protocol.
//!
//! Crypto: P-256 ECDH, AES-128-GCM, HKDF-SHA256, HMAC-SHA256 topic derivation.
//! Identity: sr25519 keypair at //wallet//sso derivation path.
//! Wire: SCALE-encoded handshake in QR, encrypted JSON over Statement Store.
//!
//! # Pairing flow
//! 1. Host generates P-256 ephemeral keypair.
//! 2. Host constructs QR URI with SCALE-encoded handshake.
//! 3. Mobile scans, performs P-256 ECDH, derives AES key, sends encrypted response.
//! 4. Host decrypts response with its ECDH-derived AES key, stores session data.
//!
//! # Channel derivation
//! All channels are hex-encoded first 32 bytes of HMAC-SHA256:
//! - Pairing topic:   HMAC-SHA256(key=account_id, msg=p256_pubkey || b"topic")
//! - Request channel: HMAC-SHA256(key=account_id, msg=p256_pubkey || b"request")
//! - Response channel: HMAC-SHA256(key=account_id, msg=p256_pubkey || b"response")

use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc;
use std::sync::Arc;
use std::time::{Duration, Instant};

#[cfg(test)]
use crate::crypto::scale_compact_encode;
use crate::crypto::{
    decrypt_aes_gcm, derive_aes_key, derive_topic, encode_handshake, encrypt_aes_gcm,
};
use p256::ecdh::EphemeralSecret;
use p256::EncodedPoint;
use rand_core::OsRng;

/// Timeout waiting for the mobile to complete the pairing handshake.
const PAIRING_TIMEOUT: Duration = Duration::from_secs(120);

/// Poll interval when waiting for a response on a channel.
const POLL_INTERVAL: Duration = Duration::from_millis(200);

/// Pairing handshake state, observed by the UI.
#[derive(Debug, Clone)]
pub enum PairingState {
    /// Waiting for the user to scan the QR code on their mobile device.
    AwaitingScan { qr_uri: String },
    /// Pairing completed - mobile confirmed, session established.
    Established {
        address: String,
        display_name: String,
    },
    /// Pairing failed with the given error message.
    Failed(String),
}

/// Handle returned by `start_pairing`. The caller polls `state_rx` for updates
/// and may set `cancel` to abort the pairing attempt.
pub struct PairingSession {
    /// Receives `PairingState` updates from the background thread.
    pub state_rx: mpsc::Receiver<PairingState>,
    /// Set to `true` to cancel the pairing in progress.
    pub cancel: Arc<AtomicBool>,
}

/// Boxed one-shot closure that subscribes to a Statement Store topic.
///
/// Returns `(subscription_id, message_rx)` where each message is `(key, value)`.
type SubscribeFn =
    Box<dyn FnOnce(&str) -> Result<(u64, mpsc::Receiver<(String, String)>), String> + Send>;

/// Boxed one-shot closure that writes a key-value pair to a Statement Store topic.
type WriteFn = Box<dyn FnOnce(&str, &str) -> Result<(), String> + Send>;

/// Configuration for a pairing session - injected Statement Store callbacks.
/// No direct dependency on host-core; all I/O is done via boxed closures.
pub struct PairingConfig {
    /// Subscribe to a Statement Store topic. Returns (subscription_id, message_rx).
    /// Each message is (key: String, value: String).
    pub subscribe: SubscribeFn,
    /// Unsubscribe from a previously subscribed topic using its subscription ID.
    pub unsubscribe: Box<dyn FnOnce(u64) + Send>,
    /// sr25519 account ID (32 bytes) from the //wallet//sso derivation path.
    /// Used as the HMAC key for topic derivation.
    pub identity_pubkey: [u8; 32],
    /// Metadata URL to embed in the QR handshake (may be empty string).
    pub metadata_url: String,
}

/// Configuration for sending a sign request to a paired mobile wallet.
pub struct PairedSignConfig {
    /// Write a value to a Statement Store topic.
    pub write: WriteFn,
    /// Subscribe to a Statement Store topic.
    pub subscribe: SubscribeFn,
    /// Unsubscribe from a topic.
    pub unsubscribe: Box<dyn FnOnce(u64) + Send>,
}

/// Start a pairing session in a background thread.
///
/// Returns a `PairingSession` immediately. The caller should call
/// `state_rx.try_recv()` periodically to check for state updates.
/// The first state emitted is always `PairingState::AwaitingScan`.
pub fn start_pairing(config: PairingConfig) -> PairingSession {
    let (state_tx, state_rx) = mpsc::sync_channel::<PairingState>(4);
    let cancel = Arc::new(AtomicBool::new(false));
    let cancel_clone = cancel.clone();

    std::thread::spawn(move || {
        run_pairing(config, state_tx, cancel_clone);
    });

    PairingSession { state_rx, cancel }
}

/// Send a sign request payload to a paired mobile wallet and wait for the
/// signature response. Runs on the calling thread (caller should use a
/// background thread).
///
/// `payload` is the raw bytes to sign.
/// `session_key` is the AES-128 key negotiated during pairing.
/// `identity_pubkey` is the sr25519 account ID from //wallet//sso.
/// `p256_pubkey` is the 65-byte uncompressed P-256 public key from pairing.
/// `result_tx` receives `Ok(signature_bytes)` or `Err(message)`.
pub fn sign_via_paired(
    payload: Vec<u8>,
    session_key: [u8; 16],
    identity_pubkey: [u8; 32],
    p256_pubkey: [u8; 65],
    config: PairedSignConfig,
    result_tx: mpsc::SyncSender<Result<Vec<u8>, String>>,
) {
    // Derive request and response channels
    let req_channel = derive_topic(&identity_pubkey, &p256_pubkey, b"request");
    let res_channel = derive_topic(&identity_pubkey, &p256_pubkey, b"response");
    let req_hex = hex::encode(req_channel);
    let res_hex = hex::encode(res_channel);

    // Encrypt the payload for the request channel
    let encrypted = encrypt_aes_gcm(&session_key, &payload);
    let encrypted_hex = hex::encode(&encrypted);

    // Subscribe to response channel before writing the request (avoid race)
    let (sub_id, msg_rx) = match (config.subscribe)(&res_hex) {
        Ok(v) => v,
        Err(e) => {
            let _ = result_tx.send(Err(format!("Subscribe failed: {e}")));
            return;
        }
    };

    // Publish the encrypted sign request
    if let Err(e) = (config.write)(&req_hex, &encrypted_hex) {
        (config.unsubscribe)(sub_id);
        let _ = result_tx.send(Err(format!("Write failed: {e}")));
        return;
    }

    // Wait for the encrypted response
    match wait_for_response(&msg_rx, PAIRING_TIMEOUT, &AtomicBool::new(false)) {
        Some((_key, value)) => {
            (config.unsubscribe)(sub_id);
            // Decode and decrypt the response
            let bytes = match hex::decode(&value) {
                Ok(b) => b,
                Err(e) => {
                    let _ = result_tx.send(Err(format!("Response hex decode failed: {e}")));
                    return;
                }
            };
            match decrypt_aes_gcm(&session_key, &bytes) {
                Ok(plaintext) => {
                    let _ = result_tx.send(Ok(plaintext));
                }
                Err(e) => {
                    let _ = result_tx.send(Err(format!("Response decrypt failed: {e}")));
                }
            }
        }
        None => {
            (config.unsubscribe)(sub_id);
            let _ = result_tx.send(Err("Sign request timed out".to_string()));
        }
    }
}

// ---------------------------------------------------------------------------
// Background pairing logic
// ---------------------------------------------------------------------------

fn run_pairing(
    config: PairingConfig,
    state_tx: mpsc::SyncSender<PairingState>,
    cancel: Arc<AtomicBool>,
) {
    // Generate ephemeral P-256 keypair
    let ephemeral_secret = EphemeralSecret::random(&mut OsRng);
    let p256_pubkey_point = EncodedPoint::from(ephemeral_secret.public_key());
    // Uncompressed form: 0x04 || x (32) || y (32) = 65 bytes
    let p256_pubkey_bytes: [u8; 65] = match p256_pubkey_point.as_bytes().try_into() {
        Ok(b) => b,
        Err(_) => {
            let _ = state_tx.send(PairingState::Failed(
                "P-256 public key encoding error".to_string(),
            ));
            return;
        }
    };

    // Derive the pairing topic for the mobile to send its response to
    let topic = derive_topic(&config.identity_pubkey, &p256_pubkey_bytes, b"topic");
    let topic_hex = hex::encode(topic);

    // Encode the QR handshake payload
    let handshake = encode_handshake(
        &config.identity_pubkey,
        &p256_pubkey_bytes,
        &config.metadata_url,
    );
    let handshake_hex = hex::encode(&handshake);
    let qr_uri = format!("polkadotapp://pair?handshake={handshake_hex}");

    // Emit the QR URI state immediately so the UI can display it
    if state_tx
        .send(PairingState::AwaitingScan { qr_uri })
        .is_err()
    {
        return;
    }

    // Subscribe to the pairing topic
    let (sub_id, msg_rx) = match (config.subscribe)(&topic_hex) {
        Ok(v) => v,
        Err(e) => {
            let _ = state_tx.send(PairingState::Failed(format!(
                "Failed to subscribe to pairing topic: {e}"
            )));
            return;
        }
    };

    // Wait for the mobile's encrypted pairing response
    let response = wait_for_response(&msg_rx, PAIRING_TIMEOUT, &cancel);
    (config.unsubscribe)(sub_id);

    if cancel.load(Ordering::Relaxed) {
        let _ = state_tx.send(PairingState::Failed("Pairing cancelled".to_string()));
        return;
    }

    let (_key, value) = match response {
        Some(msg) => msg,
        None => {
            let _ = state_tx.send(PairingState::Failed(
                "Pairing timed out waiting for mobile response".to_string(),
            ));
            return;
        }
    };

    // Decode the response hex
    let response_bytes = match hex::decode(&value) {
        Ok(b) => b,
        Err(e) => {
            let _ = state_tx.send(PairingState::Failed(format!(
                "Invalid response encoding: {e}"
            )));
            return;
        }
    };

    // Perform P-256 ECDH with the mobile's public key embedded in the response.
    // The mobile sends: p256_mobile_pubkey (65B) || nonce (12B) || ciphertext
    // The ciphertext decrypts to a JSON object with {address, displayName}.
    if response_bytes.len() < 65 + 12 + 1 {
        let _ = state_tx.send(PairingState::Failed(
            "Pairing response too short".to_string(),
        ));
        return;
    }

    let mobile_p256_bytes = &response_bytes[..65];
    let encrypted_payload = &response_bytes[65..];

    // Parse the mobile's P-256 public key
    let mobile_pubkey = match p256::PublicKey::from_sec1_bytes(mobile_p256_bytes) {
        Ok(k) => k,
        Err(e) => {
            let _ = state_tx.send(PairingState::Failed(format!(
                "Invalid mobile P-256 key: {e}"
            )));
            return;
        }
    };

    // Compute shared secret and derive AES key
    let shared_secret = ephemeral_secret.diffie_hellman(&mobile_pubkey);
    let aes_key = derive_aes_key(shared_secret.raw_secret_bytes().as_slice());

    // Decrypt the pairing payload
    let plaintext = match decrypt_aes_gcm(&aes_key, encrypted_payload) {
        Ok(p) => p,
        Err(e) => {
            let _ = state_tx.send(PairingState::Failed(format!(
                "Failed to decrypt pairing response: {e}"
            )));
            return;
        }
    };

    // Parse the JSON response: {address: string, displayName: string}
    let json: serde_json::Value = match serde_json::from_slice(&plaintext) {
        Ok(v) => v,
        Err(e) => {
            let _ = state_tx.send(PairingState::Failed(format!(
                "Invalid pairing response JSON: {e}"
            )));
            return;
        }
    };

    let address = match json.get("address").and_then(|v| v.as_str()) {
        Some(s) => s.to_string(),
        None => {
            let _ = state_tx.send(PairingState::Failed(
                "Pairing response missing 'address' field".to_string(),
            ));
            return;
        }
    };

    let display_name = json
        .get("displayName")
        .and_then(|v| v.as_str())
        .unwrap_or("")
        .to_string();

    let _ = state_tx.send(PairingState::Established {
        address,
        display_name,
    });
}

/// Poll a message channel until a message arrives, timeout elapses, or the
/// cancel flag is set. Returns `None` on timeout or cancellation.
fn wait_for_response(
    rx: &mpsc::Receiver<(String, String)>,
    timeout: Duration,
    cancel: &AtomicBool,
) -> Option<(String, String)> {
    let deadline = Instant::now() + timeout;
    loop {
        if cancel.load(Ordering::Relaxed) {
            return None;
        }
        match rx.try_recv() {
            Ok(msg) => return Some(msg),
            Err(mpsc::TryRecvError::Disconnected) => return None,
            Err(mpsc::TryRecvError::Empty) => {}
        }
        if Instant::now() >= deadline {
            return None;
        }
        std::thread::sleep(POLL_INTERVAL);
    }
}

// Crypto primitives (AES-GCM, HKDF, HMAC, SCALE encoding) live in
// crate::crypto and are imported at the top of this file.

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

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

    #[test]
    fn aes_gcm_roundtrip() {
        let key = [0x42u8; 16];
        let plaintext = b"hello pairing world";
        let encrypted = encrypt_aes_gcm(&key, plaintext);
        assert!(encrypted.len() >= 12 + 16 + plaintext.len()); // nonce + tag overhead
        let decrypted = decrypt_aes_gcm(&key, &encrypted).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn aes_gcm_wrong_key_fails() {
        let key1 = [0x42u8; 16];
        let key2 = [0x43u8; 16];
        let encrypted = encrypt_aes_gcm(&key1, b"test");
        assert!(decrypt_aes_gcm(&key2, &encrypted).is_err());
    }

    #[test]
    fn aes_gcm_too_short_fails() {
        let key = [0u8; 16];
        assert!(decrypt_aes_gcm(&key, &[0u8; 11]).is_err());
    }

    #[test]
    fn hkdf_aes_key_length() {
        let key = derive_aes_key(&[0u8; 32]);
        assert_eq!(key.len(), 16);
    }

    #[test]
    fn hkdf_aes_key_deterministic() {
        let secret = [0xABu8; 32];
        let k1 = derive_aes_key(&secret);
        let k2 = derive_aes_key(&secret);
        assert_eq!(k1, k2);
    }

    #[test]
    fn topic_derivation_deterministic() {
        let account_id = [0x01u8; 32];
        let p256_pubkey = [0x02u8; 65];
        let t1 = derive_topic(&account_id, &p256_pubkey, b"topic");
        let t2 = derive_topic(&account_id, &p256_pubkey, b"topic");
        assert_eq!(t1, t2);
    }

    #[test]
    fn topic_derivation_suffix_changes_result() {
        let account_id = [0x01u8; 32];
        let p256_pubkey = [0x02u8; 65];
        let pairing = derive_topic(&account_id, &p256_pubkey, b"topic");
        let request = derive_topic(&account_id, &p256_pubkey, b"request");
        let response = derive_topic(&account_id, &p256_pubkey, b"response");
        assert_ne!(pairing, request);
        assert_ne!(pairing, response);
        assert_ne!(request, response);
    }

    #[test]
    fn encode_handshake_length() {
        let sr_key = [0x01u8; 32];
        let p256_key = [0x04u8; 65];
        let url = "https://example.com/metadata";
        let encoded = encode_handshake(&sr_key, &p256_key, url);
        // 32 + 65 + 1 (compact prefix for 29-char url) + 29
        assert_eq!(encoded.len(), 32 + 65 + 1 + url.len());
        // First 32 bytes are the sr25519 key
        assert_eq!(&encoded[..32], &sr_key);
        // Next 65 bytes are the P-256 key
        assert_eq!(&encoded[32..97], &p256_key);
        // Compact length for 29-byte string: 29 << 2 = 116 = 0x74
        assert_eq!(encoded[97], (url.len() << 2) as u8);
    }

    #[test]
    fn encode_handshake_empty_url() {
        let sr_key = [0x00u8; 32];
        let p256_key = [0x04u8; 65];
        let encoded = encode_handshake(&sr_key, &p256_key, "");
        // 32 + 65 + 1 (compact 0) = 98
        assert_eq!(encoded.len(), 98);
        assert_eq!(encoded[97], 0); // SCALE compact for 0
    }

    #[test]
    fn scale_compact_encode_boundaries() {
        // Single-byte mode (0..=63)
        assert_eq!(scale_compact_encode(0), vec![0x00]);
        assert_eq!(scale_compact_encode(1), vec![0x04]);
        assert_eq!(scale_compact_encode(63), vec![0xfc]);
        // Two-byte mode (64..=16383)
        let enc64 = scale_compact_encode(64);
        assert_eq!(enc64.len(), 2);
        let val = u16::from_le_bytes([enc64[0], enc64[1]]);
        assert_eq!((val >> 2) as usize, 64);
        assert_eq!(val & 0b11, 0b01);
    }

    #[test]
    fn p256_ecdh_produces_aes_key() {
        use p256::ecdh::EphemeralSecret;
        use p256::EncodedPoint;

        let secret_a = EphemeralSecret::random(&mut OsRng);
        let pubkey_a = EncodedPoint::from(secret_a.public_key());

        let secret_b = EphemeralSecret::random(&mut OsRng);
        let pubkey_b = EncodedPoint::from(secret_b.public_key());

        let pubkey_a_parsed = p256::PublicKey::from_sec1_bytes(pubkey_a.as_bytes()).unwrap();
        let pubkey_b_parsed = p256::PublicKey::from_sec1_bytes(pubkey_b.as_bytes()).unwrap();

        let shared_a = secret_a.diffie_hellman(&pubkey_b_parsed);
        let shared_b = secret_b.diffie_hellman(&pubkey_a_parsed);

        let aes_a = derive_aes_key(shared_a.raw_secret_bytes().as_slice());
        let aes_b = derive_aes_key(shared_b.raw_secret_bytes().as_slice());

        assert_eq!(aes_a, aes_b, "Both sides must derive the same AES key");
    }

    /// Deterministic P-256 ECDH → HKDF test vector with hardcoded private keys.
    ///
    /// Key A scalar: 0x01 (big-endian, padded to 32 bytes).
    /// Key B scalar: 0x02 (big-endian, padded to 32 bytes).
    /// Expected shared secret x-coordinate:
    ///   7cf27b188d034f7e8a52380304b51ac3c08969e277f21b35a60b48fc47669978
    /// Expected AES key (HKDF-SHA256, info=b"polkadot-pairing"):
    ///   020e7545a89f35c553474f69b918e6d7
    #[test]
    fn p256_ecdh_pinned_vector() {
        use p256::SecretKey;

        let mut scalar_a = [0u8; 32];
        scalar_a[31] = 0x01;
        let mut scalar_b = [0u8; 32];
        scalar_b[31] = 0x02;

        let sk_a = SecretKey::from_bytes((&scalar_a).into()).unwrap();
        let sk_b = SecretKey::from_bytes((&scalar_b).into()).unwrap();

        let pk_a = sk_a.public_key();
        let pk_b = sk_b.public_key();

        // ECDH is commutative: A×PK_B and B×PK_A must yield the same shared secret.
        let shared_a = p256::ecdh::diffie_hellman(sk_a.to_nonzero_scalar(), pk_b.as_affine());
        let shared_b = p256::ecdh::diffie_hellman(sk_b.to_nonzero_scalar(), pk_a.as_affine());

        assert_eq!(
            shared_a.raw_secret_bytes().as_slice(),
            shared_b.raw_secret_bytes().as_slice(),
            "ECDH must be commutative"
        );

        let aes_a = derive_aes_key(shared_a.raw_secret_bytes().as_slice());
        let aes_b = derive_aes_key(shared_b.raw_secret_bytes().as_slice());

        assert_eq!(aes_a, aes_b, "Both sides must derive the same AES key");

        // Pinned expected value — computed offline and hardcoded for regression detection.
        let expected: [u8; 16] = [
            0x02, 0x0e, 0x75, 0x45, 0xa8, 0x9f, 0x35, 0xc5, 0x53, 0x47, 0x4f, 0x69, 0xb9, 0x18,
            0xe6, 0xd7,
        ];
        assert_eq!(aes_a, expected, "AES key must match pinned vector");
    }
}
