//! Master-key rotation orchestration.
//!
//! Drives the multi-round server protocol that re-encrypts every device's key
//! envelope under a new master key, completes the rotation, and sweeps up
//! stragglers. The round count is bounded by `MAX_ROTATION_ROUNDS` so a
//! misbehaving server cannot spin the loop indefinitely.

use bytes::Bytes;
use tracing::instrument;
use uuid::Uuid;

use crate::{
    crypto,
    error::Result,
    ids::DeviceId,
    types::{
        BeginRotationRequest, BeginRotationResponse, CompleteRotationRequest, PendingKeyInfo,
        RotationBatchEntry, RotationBatchRequest, RotationBatchResponse, RotationEntriesRequest,
        RotationEntriesResponse,
    },
};

use super::SyncKitClient;
use super::helpers::{Idempotency, check_response};

/// Hard ceiling on re-encrypt / straggler rounds within a single rotation. Each
/// round re-encrypts one server batch and the set provably shrinks (re-encrypted
/// entries drop out of the server's `key_id != new_key_id` filter), so a real
/// rotation finishes in `total_entries / batch_size` rounds. The cap only fires
/// when the set does NOT shrink, a stalled or hostile server that keeps handing
/// back work, converting an unbounded spin into a bounded error. Generous enough
/// to cover any real sync log at any batch size.
const MAX_ROTATION_ROUNDS: u32 = 100_000;

impl SyncKitClient {
    /// Rotate the master encryption key.
    ///
    /// Generates a new 256-bit key, re-encrypts all sync log entries in batches,
    /// and commits the new key to the server.
    ///
    /// Genuinely idempotent: if a previous call was interrupted, the server still
    /// holds the rotation's `pending_key`. This call detects that and **resumes
    /// with the already-committed key** rather than minting a second new key,
    /// the latter would orphan any entries already re-encrypted under the first
    /// key. Only when there is no pending rotation is a fresh key generated.
    ///
    /// Requires the encryption password (to wrap the new key) and a device_id
    /// (the device performing the rotation). Only one device can rotate at a time.
    ///
    /// After completion, all devices will receive the new key on their next
    /// `GET /keys` call. During rotation, both old and new keys are available
    /// so pulls continue to work.
    #[instrument(skip(self, password))]
    pub async fn rotate_key(&self, device_id: DeviceId, password: &str) -> Result<()> {
        // 1. Verify password against the still-active key, and learn whether a
        //    rotation is already pending for this user.
        let key_state = self.get_server_key_full().await?;
        let key_version = key_state.key_version.unwrap_or(0);
        let old_key = crypto::verify_password_against_envelope(&key_state.encrypted_key, password)?;
        let old_key = crypto::ZeroizeOnDrop(old_key);

        // 2. Resume an interrupted rotation by adopting its committed pending key,
        //    or start fresh. Either way `new_envelope` is the key the server will
        //    promote on completion, so re-encryption uses the same key it commits.
        let (new_master_key, new_envelope) =
            Self::rotation_key_material(key_state.pending_key, password)?;
        // Wrap immediately so the new key is zeroized on every exit path, an
        // error in the re-encrypt or complete loops below would otherwise drop
        // the bare array without scrubbing it.
        let new_master_key = crypto::ZeroizeOnDrop(new_master_key);

        // 3. Begin (or resume) rotation. `begin_key_rotation` is idempotent for the
        //    same device: it returns the existing rotation, ignoring `new_envelope`
        //    on resume, which is why step 2 must adopt the committed key.
        let begin_resp = self
            .begin_rotation(device_id, &new_envelope, key_version)
            .await?;
        let rotation_id = begin_resp.rotation_id;
        let new_key_id = begin_resp.new_key_id;

        tracing::info!(
            rotation_id = %rotation_id,
            target_seq = begin_resp.target_seq,
            new_key_id = new_key_id,
            "Key rotation started",
        );

        // 4. Re-encrypt loop (bounded, see reencrypt_until_done).
        self.reencrypt_until_done(rotation_id, &old_key, &new_master_key)
            .await?;

        // 5. Complete, retry if stragglers arrived from concurrent pushes. Both
        //    the straggler retries and each re-encrypt pass are capped so a server
        //    that keeps returning 409 (a stall, or a hostile peer racing pushes)
        //    cannot spin this loop forever.
        let mut straggler_rounds = 0u32;
        loop {
            match self.complete_rotation(rotation_id).await {
                Ok(()) => break,
                Err(crate::error::SyncKitError::Server { status: 409, .. }) => {
                    straggler_rounds += 1;
                    if straggler_rounds >= MAX_ROTATION_ROUNDS {
                        return Err(crate::error::SyncKitError::Internal(
                            "key rotation did not converge: server kept reporting stragglers past the round cap".into(),
                        ));
                    }
                    tracing::debug!(
                        round = straggler_rounds,
                        "Stragglers detected, re-encrypting remaining entries"
                    );
                    self.reencrypt_until_done(rotation_id, &old_key, &new_master_key)
                        .await?;
                }
                Err(e) => return Err(e),
            }
        }

        // 6. Update local state
        let (app_id, user_id) = self.require_session_ids()?;
        crate::keystore::store_key(app_id, user_id, &new_master_key)?;
        *self.master_key.write() = Some(new_master_key);
        *self.master_key_id.write() = new_key_id;
        *self.pending_key.write() = None;

        tracing::info!("Key rotation completed");
        Ok(())
    }

    /// Decide the rotation's new key material: adopt the server's committed
    /// `pending_key` when resuming an interrupted rotation, otherwise mint a
    /// fresh key. Returns `(new_master_key, new_envelope)` where the envelope is
    /// what the server promotes on completion, so the re-encryption key always
    /// matches the committed envelope. Pure (no IO) so it is unit-testable.
    fn rotation_key_material(
        pending: Option<PendingKeyInfo>,
        password: &str,
    ) -> Result<([u8; 32], String)> {
        match pending {
            Some(pending) => {
                let key = crypto::unwrap_master_key(&pending.encrypted_key, password)?;
                Ok((key, pending.encrypted_key))
            }
            None => {
                // Wrap before the fallible `wrap_master_key` so a wrap error
                // scrubs the fresh key instead of dropping a bare array.
                let key = crypto::ZeroizeOnDrop(crypto::generate_master_key());
                let envelope = crypto::wrap_master_key(&key.0, password)?;
                Ok((key.0, envelope))
            }
        }
    }

    /// Run [`reencrypt_batch`](Self::reencrypt_batch) until the server reports no
    /// entries left, bounded by [`MAX_ROTATION_ROUNDS`]. The batch set shrinks
    /// every round under an honest server; the cap is the backstop against a
    /// server that never reports it done, so the client fails with an error
    /// instead of spinning forever.
    async fn reencrypt_until_done(
        &self,
        rotation_id: Uuid,
        old_key: &[u8; 32],
        new_key: &[u8; 32],
    ) -> Result<()> {
        for _ in 0..MAX_ROTATION_ROUNDS {
            if self.reencrypt_batch(rotation_id, old_key, new_key).await? {
                return Ok(());
            }
        }
        Err(crate::error::SyncKitError::Internal(
            "key rotation did not converge: re-encrypt exceeded the round cap without the entry set draining".into(),
        ))
    }

    /// Pull a batch of entries needing re-encryption, re-encrypt them, and push back.
    /// Returns true if there are no more entries to process.
    async fn reencrypt_batch(
        &self,
        rotation_id: Uuid,
        old_key: &[u8; 32],
        new_key: &[u8; 32],
    ) -> Result<bool> {
        // Pull entries needing re-encryption (starting from seq 0 each time,
        // since the server filters by key_id != new_key_id)
        let entries_resp = self.rotation_entries(rotation_id, 0).await?;

        if entries_resp.entries.is_empty() {
            return Ok(true);
        }

        let has_more = entries_resp.has_more;

        // Re-encrypt each entry
        let reencrypted: Vec<RotationBatchEntry> = entries_resp
            .entries
            .into_iter()
            .map(|entry| {
                let new_data = match entry.data {
                    Some(ref encrypted_value) => {
                        // Decrypt with old key, re-encrypt with new key, binding the
                        // same (table, row_id) AAD on both ends. A legacy untagged
                        // entry decrypts with empty AAD and re-emits as a v2 tagged,
                        // position-bound payload (opportunistic upgrade).
                        let ctx = crypto::AeadContext::entry(&entry.table, &entry.row_id);
                        let plaintext = crypto::decrypt_json_aad(encrypted_value, old_key, &ctx)?;
                        let reencrypted = crypto::encrypt_json_aad(&plaintext, new_key, &ctx)?;
                        Some(reencrypted)
                    }
                    None => None, // DELETE entries have no data
                };
                Ok(RotationBatchEntry {
                    seq: entry.seq,
                    data: new_data,
                })
            })
            .collect::<Result<Vec<_>>>()?;

        // Push re-encrypted batch
        let count = reencrypted.len();
        self.rotation_batch(rotation_id, reencrypted).await?;

        tracing::debug!(count, "Re-encrypted batch submitted");

        Ok(!has_more)
    }

    // ── HTTP helpers ──

    async fn begin_rotation(
        &self,
        device_id: DeviceId,
        new_encrypted_key: &str,
        expected_key_version: i32,
    ) -> Result<BeginRotationResponse> {
        let token = self.require_token()?;
        let url = format!(
            "{}/api/v1/sync/keys/rotate",
            self.config.server_url.trim_end_matches('/')
        );

        let body = Bytes::from(serde_json::to_vec(&BeginRotationRequest {
            device_id,
            new_encrypted_key: new_encrypted_key.to_string(),
            expected_key_version,
        })?);

        self.retry_request_json(
            Idempotency::IdempotentWrite {
                on: "one rotation per (app_id, user_id), migration 089",
            },
            || {
                let req = self
                    .http
                    .post(&url)
                    .bearer_auth(&token)
                    .header("content-type", "application/json")
                    .body(body.clone());
                async move { check_response(req.send().await?).await }
            },
        )
        .await
    }

    async fn rotation_entries(
        &self,
        rotation_id: Uuid,
        after_seq: i64,
    ) -> Result<RotationEntriesResponse> {
        let token = self.require_token()?;
        let url = format!(
            "{}/api/v1/sync/keys/rotate/entries",
            self.config.server_url.trim_end_matches('/')
        );

        let body = Bytes::from(serde_json::to_vec(&RotationEntriesRequest {
            rotation_id,
            after_seq,
        })?);

        self.retry_request_json(Idempotency::ReadOnly, || {
            let req = self
                .http
                .post(&url)
                .bearer_auth(&token)
                .header("content-type", "application/json")
                .body(body.clone());
            async move { check_response(req.send().await?).await }
        })
        .await
    }

    async fn rotation_batch(
        &self,
        rotation_id: Uuid,
        entries: Vec<RotationBatchEntry>,
    ) -> Result<RotationBatchResponse> {
        let token = self.require_token()?;
        let url = format!(
            "{}/api/v1/sync/keys/rotate/batch",
            self.config.server_url.trim_end_matches('/')
        );

        let body = Bytes::from(serde_json::to_vec(&RotationBatchRequest {
            rotation_id,
            entries,
        })?);

        self.retry_request_json(
            Idempotency::IdempotentWrite {
                on: "seq (re-encrypt overwrites the same row)",
            },
            || {
                let req = self
                    .http
                    .post(&url)
                    .bearer_auth(&token)
                    .header("content-type", "application/json")
                    .body(body.clone());
                async move { check_response(req.send().await?).await }
            },
        )
        .await
    }

    async fn complete_rotation(&self, rotation_id: Uuid) -> Result<()> {
        let token = self.require_token()?;
        let url = format!(
            "{}/api/v1/sync/keys/rotate/complete",
            self.config.server_url.trim_end_matches('/')
        );

        let body = Bytes::from(serde_json::to_vec(&CompleteRotationRequest {
            rotation_id,
        })?);

        self.retry_request(Idempotency::IdempotentWrite { on: "rotation_id" }, || {
            let req = self
                .http
                .post(&url)
                .bearer_auth(&token)
                .header("content-type", "application/json")
                .body(body.clone());
            async move { check_response(req.send().await?).await }
        })
        .await?;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{GetKeyResponse, PullChangeEntry};

    #[test]
    fn rotation_resume_reuses_committed_pending_key() {
        // Simulate an interrupted rotation: the server holds a pending envelope
        // wrapping key K. Resuming must recover K (not mint a fresh key), or
        // entries already re-encrypted under K become undecryptable.
        let committed_key = crypto::generate_master_key();
        let pending_envelope = crypto::wrap_master_key(&committed_key, "pw").unwrap();
        let pending = Some(PendingKeyInfo {
            encrypted_key: pending_envelope.clone(),
            key_id: 2,
        });

        let (resumed_key, envelope) = SyncKitClient::rotation_key_material(pending, "pw").unwrap();
        assert_eq!(
            resumed_key, committed_key,
            "resume must adopt the committed key"
        );
        assert_eq!(
            envelope, pending_envelope,
            "resume must keep the committed envelope"
        );
    }

    #[test]
    fn rotation_fresh_start_mints_unwrappable_key() {
        // No pending rotation: a fresh key is generated and its envelope unwraps
        // back to that same key under the password.
        let (fresh_key, envelope) = SyncKitClient::rotation_key_material(None, "pw").unwrap();
        let recovered = crypto::unwrap_master_key(&envelope, "pw").unwrap();
        assert_eq!(recovered, fresh_key);
    }

    #[test]
    fn reencrypt_preserves_plaintext() {
        // Simulate re-encryption: encrypt with old key, decrypt, re-encrypt with new key
        let old_key = crypto::generate_master_key();
        let new_key = crypto::generate_master_key();
        let original = serde_json::json!({"title": "Test task", "priority": 3});

        // Encrypt with old key (simulates existing sync_log entry)
        let encrypted_old = crypto::encrypt_json(&original, &old_key).unwrap();

        // Re-encrypt: decrypt with old, encrypt with new
        let plaintext = crypto::decrypt_json(&encrypted_old, &old_key).unwrap();
        let encrypted_new = crypto::encrypt_json(&plaintext, &new_key).unwrap();

        // Verify: new key can decrypt to original data
        let recovered = crypto::decrypt_json(&encrypted_new, &new_key).unwrap();
        assert_eq!(recovered, original);

        // Verify: old key cannot decrypt re-encrypted data
        assert!(crypto::decrypt_json(&encrypted_new, &old_key).is_err());
    }

    #[test]
    fn reencrypt_null_data_stays_null() {
        // DELETE entries have no data, rotation should preserve this
        let old_key = crypto::generate_master_key();
        let new_key = crypto::generate_master_key();

        // No data to re-encrypt
        let data: Option<serde_json::Value> = None;
        let reencrypted = match data {
            Some(ref val) => {
                let plaintext = crypto::decrypt_json(val, &old_key).unwrap();
                Some(crypto::encrypt_json(&plaintext, &new_key).unwrap())
            }
            None => None,
        };

        assert!(reencrypted.is_none());
    }

    #[test]
    fn rotation_request_types_serialize() {
        let req = BeginRotationRequest {
            device_id: DeviceId::new(Uuid::new_v4()),
            new_encrypted_key: "envelope-json".to_string(),
            expected_key_version: 1,
        };
        let json = serde_json::to_string(&req).unwrap();
        assert!(json.contains("expected_key_version"));

        let req = RotationBatchRequest {
            rotation_id: Uuid::new_v4(),
            entries: vec![
                RotationBatchEntry {
                    seq: 1,
                    data: Some(serde_json::json!("encrypted")),
                },
                RotationBatchEntry { seq: 2, data: None },
            ],
        };
        let json = serde_json::to_string(&req).unwrap();
        assert!(json.contains("rotation_id"));
        assert!(json.contains("\"seq\":1"));
    }

    #[test]
    fn rotation_response_types_deserialize() {
        let json = r#"{"rotation_id": "550e8400-e29b-41d4-a716-446655440000", "target_seq": 100, "new_key_id": 2}"#;
        let resp: BeginRotationResponse = serde_json::from_str(json).unwrap();
        assert_eq!(resp.target_seq, 100);
        assert_eq!(resp.new_key_id, 2);

        let json = r#"{"entries": [{"seq": 1, "table": "tasks", "row_id": "r1", "data": "encrypted"}, {"seq": 2, "table": "tasks", "row_id": "r2", "data": null}], "has_more": true}"#;
        let resp: RotationEntriesResponse = serde_json::from_str(json).unwrap();
        assert_eq!(resp.entries.len(), 2);
        assert!(resp.has_more);
        assert!(resp.entries[0].data.is_some());
        assert_eq!(resp.entries[0].table, "tasks");
        assert_eq!(resp.entries[0].row_id, "r1");
        assert!(resp.entries[1].data.is_none());
    }

    #[test]
    fn pending_key_info_deserializes() {
        let json = r#"{
            "encrypted_key": "envelope",
            "key_version": 2,
            "key_id": 1,
            "pending_key": {"encrypted_key": "new-envelope", "key_id": 2}
        }"#;
        let resp: GetKeyResponse = serde_json::from_str(json).unwrap();
        assert!(resp.pending_key.is_some());
        let pending = resp.pending_key.unwrap();
        assert_eq!(pending.key_id, 2);
        assert_eq!(pending.encrypted_key, "new-envelope");
    }

    #[test]
    fn get_key_response_without_pending_key() {
        let json = r#"{"encrypted_key": "envelope", "key_version": 1}"#;
        let resp: GetKeyResponse = serde_json::from_str(json).unwrap();
        assert!(resp.pending_key.is_none());
        assert_eq!(resp.key_id, None); // backward compat
    }

    #[test]
    fn pull_change_entry_with_key_id() {
        let device_id = Uuid::new_v4();
        let json = format!(
            r#"{{"seq": 1, "device_id": "{device_id}", "table": "t", "op": "INSERT", "row_id": "r", "timestamp": "2025-06-01T12:00:00Z", "data": null, "key_id": 2}}"#
        );
        let entry: PullChangeEntry = serde_json::from_str(&json).unwrap();
        assert_eq!(entry.key_id, Some(2));
    }

    #[test]
    fn pull_change_entry_without_key_id_defaults_none() {
        let device_id = Uuid::new_v4();
        let json = format!(
            r#"{{"seq": 1, "device_id": "{device_id}", "table": "t", "op": "INSERT", "row_id": "r", "timestamp": "2025-06-01T12:00:00Z", "data": null}}"#
        );
        let entry: PullChangeEntry = serde_json::from_str(&json).unwrap();
        assert_eq!(entry.key_id, None);
    }
}