//! Master-key setup and password lifecycle.
//!
//! Establishes the master key on a first or subsequent device by wrapping and
//! unwrapping it against the server-held key envelope, caches it in the OS
//! keychain when available, and re-wraps it on a password change.

use bytes::Bytes;
use std::sync::Arc;
use tracing::instrument;

use crate::{
    crypto,
    error::Result,
    keystore,
    types::{GetKeyResponse, PutKeyRequest},
};

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

impl SyncKitClient {
    /// Check if the server has an encrypted master key for this user.
    #[instrument(skip(self))]
    pub async fn has_server_key(&self) -> Result<bool> {
        let (url, token) = self.key_url_and_token()?;

        let result = self
            .retry_request(Idempotency::ReadOnly, || {
                let req = self.http.get(url).bearer_auth(&token);
                async move {
                    let resp = req.send().await?;
                    match resp.status().as_u16() {
                        200 | 404 => Ok(resp),
                        status => {
                            let message = crate::client::helpers::read_text_capped(
                                resp,
                                crate::client::helpers::MAX_CONTROL_BODY_BYTES,
                            )
                            .await;
                            Err(crate::error::SyncKitError::Server {
                                status,
                                message,
                                retry_after_secs: None,
                            })
                        }
                    }
                }
            })
            .await?;

        Ok(result.status().as_u16() == 200)
    }

    /// First device: generate a new master key, encrypt it, push to server, cache in keychain.
    #[instrument(skip(self, password))]
    pub async fn setup_encryption_new(&self, password: &str) -> Result<()> {
        let (app_id, user_id) = self.require_session_ids()?;

        // Wrap the fresh key immediately so it is scrubbed on every early-return
        // path below (wrap/put/store can each fail), not left as a bare array on
        // the stack until the final move into the lock.
        let master_key = crypto::ZeroizeOnDrop(crypto::generate_master_key());
        let envelope = crypto::wrap_master_key(&master_key.0, password)?;

        // Push to server (expected_version 0 = first key)
        self.put_server_key(&envelope, 0).await?;

        // Cache in OS keychain
        keystore::store_key(app_id, user_id, &master_key.0)?;

        // Store in memory
        *self.master_key.write() = Some(master_key);

        tracing::info!("New master key generated and stored");
        Ok(())
    }

    /// Second device: pull encrypted master key from server, decrypt with password, cache.
    #[instrument(skip(self, password))]
    pub async fn setup_encryption_existing(&self, password: &str) -> Result<()> {
        let (app_id, user_id) = self.require_session_ids()?;

        let (envelope_json, _key_version) = self.get_server_key().await?;
        // Wrap immediately so a failed keychain store scrubs the key rather than
        // dropping a bare array.
        let master_key =
            crypto::ZeroizeOnDrop(crypto::unwrap_master_key(&envelope_json, password)?);

        // Cache in OS keychain
        keystore::store_key(app_id, user_id, &master_key.0)?;

        // Store in memory
        *self.master_key.write() = Some(master_key);

        tracing::info!("Master key recovered from server");
        Ok(())
    }

    /// Subsequent launches: try to load the master key from the OS keychain.
    /// Returns true if a key was found.
    pub fn try_load_key_from_keychain(&self) -> Result<bool> {
        let (app_id, user_id) = self.require_session_ids()?;

        if let Some(key) = keystore::load_key(app_id, user_id)? {
            *self.master_key.write() = Some(key);
            tracing::info!("Master key loaded from keychain");
            Ok(true)
        } else {
            Ok(false)
        }
    }

    /// Change the encryption password. Always validates the old password
    /// against the server envelope before re-encrypting with the new password.
    ///
    /// Even when the master key is cached in memory (normal case, user is
    /// logged in), the old password is verified by attempting to unwrap the
    /// server envelope. This prevents an attacker with session access from
    /// changing the password without knowing the current one.
    #[instrument(skip(self, old_password, new_password))]
    pub async fn change_password(&self, old_password: &str, new_password: &str) -> Result<()> {
        // Always fetch the envelope from the server and verify old_password
        // can decrypt it, regardless of whether we have a cached key.
        let (envelope_json, key_version) = self.get_server_key().await?;
        let verified_key = crypto::verify_password_against_envelope(&envelope_json, old_password)?;

        let master_key = crypto::ZeroizeOnDrop(verified_key);

        // If we hold the master key in memory, the envelope the server just handed
        // us MUST wrap that same key. A hostile server could otherwise substitute a
        // *different* envelope that also unwraps under `old_password` (e.g. one it
        // captured), making us re-wrap the wrong key under the new password and
        // lock the user out of their own (original-key-encrypted) data. Refuse.
        if let Some(cached) = self.master_key.read().as_ref()
            && cached.0 != master_key.0
        {
            return Err(crate::error::SyncKitError::Crypto(
                "server key envelope does not match the in-memory master key; refusing to change password".into(),
            ));
        }

        // Re-wrap with new password (generates fresh random salt)
        let new_envelope = crypto::wrap_master_key(&master_key.0, new_password)?;

        // Optimistic lock: reject if another device changed the password
        // between our GET and this PUT.
        self.put_server_key(&new_envelope, key_version).await?;

        tracing::info!("Encryption password changed");
        Ok(())
    }

    /// Build the key endpoint URL and extract the bearer token.
    pub(super) fn key_url_and_token(&self) -> Result<(&str, Arc<SecretToken>)> {
        let token = self.require_token()?;
        Ok((&self.endpoints.keys, token))
    }

    /// Upload the encrypted master key envelope to the server (PUT /api/sync/keys).
    ///
    /// `expected_version` is the key version the client expects. The server
    /// rejects with 409 if the current version doesn't match (another device
    /// changed the password). Use 0 for the initial key upload.
    pub(super) async fn put_server_key(
        &self,
        envelope_json: &str,
        expected_version: i32,
    ) -> Result<()> {
        let (url, token) = self.key_url_and_token()?;

        let body = Bytes::from(serde_json::to_vec(&PutKeyRequest {
            encrypted_key: envelope_json.to_string(),
            expected_version,
        })?);

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

    /// Download the encrypted master key envelope from the server (GET /api/sync/keys).
    /// Returns `(envelope_json, key_version)`.
    pub(super) async fn get_server_key(&self) -> Result<(String, i32)> {
        let resp = self.get_server_key_full().await?;
        Ok((resp.encrypted_key, resp.key_version.unwrap_or(0)))
    }

    /// Download the full key state, including any in-progress rotation's
    /// `pending_key`. Used by `rotate_key` to resume an interrupted rotation
    /// with the already-committed key rather than minting a fresh one.
    pub(super) async fn get_server_key_full(&self) -> Result<GetKeyResponse> {
        let (url, token) = self.key_url_and_token()?;

        self.retry_request_json(Idempotency::ReadOnly, || {
            let req = self.http.get(url).bearer_auth(&token);
            async move { check_response(req.send().await?).await }
        })
        .await
    }
}

#[cfg(test)]
mod tests {
    use crate::error::SyncKitError;

    use super::*;

    fn test_config() -> super::super::SyncKitConfig {
        super::super::SyncKitConfig {
            server_url: "https://example.com".to_string(),
            api_key: "test-api-key-123".to_string(),
        }
    }

    fn test_ids() -> (crate::ids::AppId, crate::ids::UserId) {
        (
            crate::ids::AppId::new(
                uuid::Uuid::parse_str("550e8400-e29b-41d4-a716-446655440000").unwrap(),
            ),
            crate::ids::UserId::new(
                uuid::Uuid::parse_str("6ba7b810-9dad-11d1-80b4-00c04fd430c8").unwrap(),
            ),
        )
    }

    #[test]
    fn key_url_and_token_builds_correct_url() {
        let client = SyncKitClient::new(test_config());
        let (app_id, user_id) = test_ids();
        client.restore_session("bearer-token", user_id, app_id);

        let (url, token) = client.key_url_and_token().unwrap();
        assert_eq!(url, "https://example.com/api/v1/sync/keys");
        assert_eq!(token.as_str(), "bearer-token");
    }

    #[test]
    fn key_url_and_token_fails_without_session() {
        let client = SyncKitClient::new(test_config());
        let err = client.key_url_and_token().unwrap_err();
        assert!(matches!(err, SyncKitError::NotAuthenticated));
    }

    // ── Key types ──

    #[test]
    fn put_key_request_serialization() {
        let req = PutKeyRequest {
            encrypted_key: "envelope-json-here".to_string(),
            expected_version: 1,
        };

        let json = serde_json::to_string(&req).unwrap();
        let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed["encrypted_key"], "envelope-json-here");
        assert_eq!(parsed["expected_version"], 1);
    }

    #[test]
    fn get_key_response_deserialization() {
        let json = r#"{"encrypted_key": "{\"v\":1,\"salt\":\"...\",\"nonce\":\"...\",\"ciphertext\":\"...\"}"}"#;
        let resp: GetKeyResponse = serde_json::from_str(json).unwrap();
        assert!(resp.encrypted_key.contains("\"v\":1"));
    }
}