use crate::auth::{AuthToken, AuthResult, AuthError};
use aes_gcm::{
    aead::{Aead, KeyInit},
    Aes256Gcm, Nonce
};
use base64::{Engine as _, engine::general_purpose};
use gloo_storage::{LocalStorage, Storage};
use serde::{Deserialize, Serialize};
use wasm_bindgen::prelude::*;
use web_sys::window;

const TOKEN_KEY: &str = "cc_auth_token";
const ENCRYPTION_KEY: &str = "cc_encryption_key";

#[derive(Serialize, Deserialize)]
struct EncryptedData {
    data: String,
    nonce: String,
}

pub struct SecureTokenStorage {
    cipher: Aes256Gcm,
}

impl SecureTokenStorage {
    pub fn new() -> AuthResult<Self> {
        let key = Self::get_or_create_encryption_key()?;
        let cipher = Aes256Gcm::new_from_slice(&key)
            .map_err(|e| AuthError {
                message: format!("Failed to initialize cipher: {}", e),
                code: Some("CIPHER_INIT_ERROR".to_string()),
                field: None,
            })?;

        Ok(Self { cipher })
    }

    fn get_or_create_encryption_key() -> AuthResult<Vec<u8>> {
        // Try to get existing key from sessionStorage (more secure than localStorage for keys)
        if let Some(storage) = window().and_then(|w| w.session_storage().ok().flatten()) {
            if let Ok(encoded_key) = storage.get_item(ENCRYPTION_KEY) {
                if let Some(key_str) = encoded_key {
                    if let Ok(key_bytes) = general_purpose::STANDARD.decode(key_str) {
                        if key_bytes.len() == 32 {
                            return Ok(key_bytes);
                        }
                    }
                }
            }
        }

        // Generate new key if none exists
        let key = getrandom::getrandom::<[u8; 32]>().map_err(|_| AuthError {
            message: "Failed to generate encryption key".to_string(),
            code: Some("KEY_GENERATION_ERROR".to_string()),
            field: None,
        })?;
        let encoded_key = general_purpose::STANDARD.encode(&key);

        // Store in sessionStorage
        if let Some(storage) = window().and_then(|w| w.session_storage().ok().flatten()) {
            storage.set_item(ENCRYPTION_KEY, &encoded_key)
                .map_err(|_| AuthError {
                    message: "Failed to store encryption key".to_string(),
                    code: Some("KEY_STORAGE_ERROR".to_string()),
                    field: None,
                })?;
        }

        Ok(key)
    }

    pub fn store_token(&self, token: &AuthToken) -> AuthResult<()> {
        let serialized = serde_json::to_string(token)
            .map_err(|e| AuthError {
                message: format!("Failed to serialize token: {}", e),
                code: Some("SERIALIZATION_ERROR".to_string()),
                field: None,
            })?;

        let encrypted_data = self.encrypt_data(&serialized)?;
        let encrypted_json = serde_json::to_string(&encrypted_data)
            .map_err(|e| AuthError {
                message: format!("Failed to serialize encrypted data: {}", e),
                code: Some("SERIALIZATION_ERROR".to_string()),
                field: None,
            })?;

        LocalStorage::set(TOKEN_KEY, encrypted_json)
            .map_err(|_| AuthError {
                message: "Failed to store token in localStorage".to_string(),
                code: Some("STORAGE_ERROR".to_string()),
                field: None,
            })?;

        Ok(())
    }

    pub fn retrieve_token(&self) -> AuthResult<Option<AuthToken>> {
        let encrypted_json: String = match LocalStorage::get(TOKEN_KEY) {
            Ok(data) => data,
            Err(_) => return Ok(None),
        };

        let encrypted_data: EncryptedData = serde_json::from_str(&encrypted_json)
            .map_err(|e| AuthError {
                message: format!("Failed to deserialize encrypted data: {}", e),
                code: Some("DESERIALIZATION_ERROR".to_string()),
                field: None,
            })?;

        let decrypted_data = self.decrypt_data(&encrypted_data)?;
        let token: AuthToken = serde_json::from_str(&decrypted_data)
            .map_err(|e| AuthError {
                message: format!("Failed to deserialize token: {}", e),
                code: Some("DESERIALIZATION_ERROR".to_string()),
                field: None,
            })?;

        Ok(Some(token))
    }

    pub fn remove_token(&self) -> AuthResult<()> {
        LocalStorage::delete(TOKEN_KEY);

        // Also clear encryption key from sessionStorage
        if let Some(storage) = window().and_then(|w| w.session_storage().ok().flatten()) {
            let _ = storage.remove_item(ENCRYPTION_KEY);
        }

        Ok(())
    }

    fn encrypt_data(&self, data: &str) -> AuthResult<EncryptedData> {
        let mut nonce_bytes = [0u8; 12];
        getrandom::getrandom(&mut nonce_bytes).map_err(|_| AuthError {
            message: "Failed to generate nonce".to_string(),
            code: Some("NONCE_GENERATION_ERROR".to_string()),
            field: None,
        })?;
        let nonce = Nonce::from_slice(&nonce_bytes);

        let ciphertext = self.cipher.encrypt(nonce, data.as_bytes())
            .map_err(|e| AuthError {
                message: format!("Encryption failed: {}", e),
                code: Some("ENCRYPTION_ERROR".to_string()),
                field: None,
            })?;

        Ok(EncryptedData {
            data: general_purpose::STANDARD.encode(ciphertext),
            nonce: general_purpose::STANDARD.encode(nonce_bytes),
        })
    }

    fn decrypt_data(&self, encrypted_data: &EncryptedData) -> AuthResult<String> {
        let ciphertext = general_purpose::STANDARD.decode(&encrypted_data.data)
            .map_err(|e| AuthError {
                message: format!("Failed to decode ciphertext: {}", e),
                code: Some("DECODE_ERROR".to_string()),
                field: None,
            })?;

        let nonce_bytes = general_purpose::STANDARD.decode(&encrypted_data.nonce)
            .map_err(|e| AuthError {
                message: format!("Failed to decode nonce: {}", e),
                code: Some("DECODE_ERROR".to_string()),
                field: None,
            })?;

        let nonce = Nonce::from_slice(&nonce_bytes);
        let plaintext = self.cipher.decrypt(nonce, ciphertext.as_ref())
            .map_err(|e| AuthError {
                message: format!("Decryption failed: {}", e),
                code: Some("DECRYPTION_ERROR".to_string()),
                field: None,
            })?;

        String::from_utf8(plaintext)
            .map_err(|e| AuthError {
                message: format!("Failed to convert decrypted data to string: {}", e),
                code: Some("UTF8_ERROR".to_string()),
                field: None,
            })
    }

    pub fn clear_all_data(&self) -> AuthResult<()> {
        LocalStorage::clear();

        if let Some(storage) = window().and_then(|w| w.session_storage().ok().flatten()) {
            let _ = storage.clear();
        }

        Ok(())
    }
}

impl Default for SecureTokenStorage {
    fn default() -> Self {
        Self::new().expect("Failed to create SecureTokenStorage")
    }
}
