//! SigSocket bindings for WASM - integrates sigsocket_client with vault system

use std::cell::RefCell;
use wasm_bindgen::prelude::*;
use serde::{Deserialize, Serialize};
use sigsocket_client::{SigSocketClient, SignRequest, SignRequestHandler, Result as SigSocketResult, SigSocketError};
use web_sys::console;
use base64::prelude::*;

use crate::vault_bindings::{get_workspace_default_public_key, get_current_keyspace_name, is_unlocked, sign_with_default_keypair};

// Global SigSocket client instance
thread_local! {
    static SIGSOCKET_CLIENT: RefCell<Option<SigSocketClient>> = RefCell::new(None);
}

// Helper macro for console logging
macro_rules! console_log {
    ($($t:tt)*) => (console::log_1(&format!($($t)*).into()))
}

/// Extension notification handler that forwards requests to JavaScript
pub struct ExtensionNotificationHandler {
    callback: js_sys::Function,
}

impl ExtensionNotificationHandler {
    pub fn new(callback: js_sys::Function) -> Self {
        Self { callback }
    }
}

impl SignRequestHandler for ExtensionNotificationHandler {
    fn handle_sign_request(&self, request: &SignRequest) -> SigSocketResult<Vec<u8>> {
        console_log!("📨 WASM: Handling sign request: {}", request.id);

        // First, store the request in the WASM client
        let store_result = SIGSOCKET_CLIENT.with(|c| {
            let mut client_opt = c.borrow_mut();
            if let Some(client) = client_opt.as_mut() {
                // Get the connected public key as the target
                if let Some(target_public_key) = client.connected_public_key() {
                    client.add_pending_request(request.clone(), target_public_key.to_string());
                    console_log!("✅ WASM: Stored sign request: {}", request.id);
                    Ok(())
                } else {
                    Err(SigSocketError::Other("No connected public key".to_string()))
                }
            } else {
                Err(SigSocketError::Other("No SigSocket client available".to_string()))
            }
        });

        // If storage failed, return error
        if let Err(e) = store_result {
            console_log!("❌ WASM: Failed to store request: {:?}", e);
            return Err(e);
        }

        // Create event object for JavaScript notification
        let event = js_sys::Object::new();
        js_sys::Reflect::set(&event, &"type".into(), &"sign_request".into())
            .map_err(|_| SigSocketError::Other("Failed to set event type".to_string()))?;
        js_sys::Reflect::set(&event, &"request_id".into(), &request.id.clone().into())
            .map_err(|_| SigSocketError::Other("Failed to set request_id".to_string()))?;
        js_sys::Reflect::set(&event, &"message".into(), &request.message.clone().into())
            .map_err(|_| SigSocketError::Other("Failed to set message".to_string()))?;

        // Notify the extension
        match self.callback.call1(&wasm_bindgen::JsValue::NULL, &event) {
            Ok(_) => {
                console_log!("✅ WASM: Notified extension about sign request: {}", request.id);
                // Return an error to indicate this request should not be auto-signed
                // The extension will handle the approval flow
                Err(SigSocketError::Other("Request forwarded to extension for approval".to_string()))
            }
            Err(e) => {
                console_log!("❌ WASM: Failed to notify extension: {:?}", e);
                Err(SigSocketError::Other("Extension notification failed".to_string()))
            }
        }
    }
}

/// Connection information for SigSocket
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SigSocketConnectionInfo {
    pub workspace: String,
    pub public_key: String,
    pub is_connected: bool,
    pub server_url: String,
}

/// SigSocket manager for high-level operations
#[wasm_bindgen]
pub struct SigSocketManager;

#[wasm_bindgen]
impl SigSocketManager {
    /// Connect to SigSocket server with smart connection management
    ///
    /// This handles all connection logic:
    /// - Reuses existing connection if same workspace
    /// - Switches connection if different workspace
    /// - Creates new connection if none exists
    ///
    /// # Arguments
    /// * `workspace` - The workspace name to connect with
    /// * `server_url` - The SigSocket server URL (e.g., "ws://localhost:8080/ws")
    /// * `event_callback` - JavaScript function to call when events occur
    ///
    /// # Returns
    /// * `Ok(connection_info)` - JSON string with connection details
    /// * `Err(error)` - If connection failed or workspace is invalid
    #[wasm_bindgen]
    pub async fn connect_workspace_with_events(workspace: &str, server_url: &str, event_callback: &js_sys::Function) -> Result<String, JsValue> {
        // 1. Validate workspace exists and get default public key from vault
        let public_key_js = get_workspace_default_public_key(workspace).await
            .map_err(|e| JsValue::from_str(&format!("Failed to get workspace public key: {:?}", e)))?;

        let public_key_hex = public_key_js.as_string()
            .ok_or_else(|| JsValue::from_str("Public key is not a string"))?;

        // 2. Decode public key
        let public_key_bytes = hex::decode(&public_key_hex)
            .map_err(|e| JsValue::from_str(&format!("Invalid public key format: {}", e)))?;
        
        // 3. Check if already connected to same workspace and handle disconnection
        let should_connect = SIGSOCKET_CLIENT.with(|c| {
            let mut client_opt = c.borrow_mut();

            // Check if we already have a client for this workspace
            if let Some(existing_client) = client_opt.as_ref() {
                if let Some(existing_key) = existing_client.connected_public_key() {
                    if existing_key == hex::encode(&public_key_bytes) && existing_client.is_connected() {
                        console_log!("🔄 WASM: Already connected to workspace: {}", workspace);
                        return false; // Reuse existing connection
                    } else {
                        console_log!("🔄 WASM: Switching workspace from {} to {}",
                                   existing_key, hex::encode(&public_key_bytes));

                        // Disconnect the old client
                        *client_opt = None; // This will drop the old client and close WebSocket
                        console_log!("🔌 WASM: Disconnected from old workspace");

                        return true; // Need new connection
                    }
                }
            }

            true // Need new connection, no old one to disconnect
        });

        // 4. Create and connect if needed
        if should_connect {
            console_log!("🔗 WASM: Creating new connection for workspace: {}", workspace);

            // Create new client
            let mut client = SigSocketClient::new(server_url, public_key_bytes.clone())
                .map_err(|e| JsValue::from_str(&format!("Failed to create client: {:?}", e)))?;

            // Set up extension notification handler
            let handler = ExtensionNotificationHandler::new(event_callback.clone());
            client.set_sign_handler(handler);

            // Connect to the WebSocket server
            client.connect().await
                .map_err(|e| JsValue::from_str(&format!("Connection failed: {:?}", e)))?;

            console_log!("✅ WASM: Connected to SigSocket server for workspace: {}", workspace);

            // Store the connected client
            SIGSOCKET_CLIENT.with(|c| {
                *c.borrow_mut() = Some(client);
            });
        }

        // 6. Return connection info
        let connection_info = SigSocketConnectionInfo {
            workspace: workspace.to_string(),
            public_key: public_key_hex.clone(),
            is_connected: true,
            server_url: server_url.to_string(),
        };
        
        // 7. Serialize and return connection info
        serde_json::to_string(&connection_info)
            .map_err(|e| JsValue::from_str(&format!("Serialization error: {}", e)))
    }

    /// Connect to SigSocket server with a specific workspace (backward compatibility)
    ///
    /// This is a simpler version that doesn't set up event callbacks.
    /// Use connect_workspace_with_events for full functionality.
    ///
    /// # Arguments
    /// * `workspace` - The workspace name to connect with
    /// * `server_url` - The SigSocket server URL (e.g., "ws://localhost:8080/ws")
    ///
    /// # Returns
    /// * `Ok(connection_info)` - JSON string with connection details
    /// * `Err(error)` - If connection failed or workspace is invalid
    #[wasm_bindgen]
    pub async fn connect_workspace(workspace: &str, server_url: &str) -> Result<String, JsValue> {
        // Create a dummy callback that just logs
        let dummy_callback = js_sys::Function::new_no_args("console.log('SigSocket event:', arguments[0]);");
        Self::connect_workspace_with_events(workspace, server_url, &dummy_callback).await
    }
    
    /// Disconnect from SigSocket server
    ///
    /// # Returns
    /// * `Ok(())` - Successfully disconnected
    /// * `Err(error)` - If disconnect failed
    #[wasm_bindgen]
    pub async fn disconnect() -> Result<(), JsValue> {
        SIGSOCKET_CLIENT.with(|c| {
            let mut client_opt = c.borrow_mut();
            if let Some(client) = client_opt.take() {
                let workspace_info = client.connected_public_key()
                    .map(|key| key[..16].to_string())
                    .unwrap_or_else(|| "unknown".to_string());

                // Dropping the client will close the WebSocket connection
                drop(client);
                console_log!("🔌 WASM: Disconnected SigSocket client (was: {}...)", workspace_info);
            } else {
                console_log!("🔌 WASM: No SigSocket client to disconnect");
            }
            Ok(())
        })
    }
    
    /// Check if we can approve a specific sign request
    /// 
    /// This validates that:
    /// 1. The request exists
    /// 2. The vault session is unlocked
    /// 3. The current workspace matches the request's target
    /// 
    /// # Arguments
    /// * `request_id` - The ID of the request to validate
    /// 
    /// # Returns
    /// * `Ok(true)` - Request can be approved
    /// * `Ok(false)` - Request cannot be approved
    /// * `Err(error)` - Validation error
    #[wasm_bindgen]
    pub async fn can_approve_request(request_id: &str) -> Result<bool, JsValue> {
        // 1. Check if vault session is unlocked
        if !is_unlocked() {
            return Ok(false);
        }

        // 2. Get current workspace and its public key
        let current_workspace = get_current_keyspace_name()
            .map_err(|e| JsValue::from_str(&format!("Failed to get current workspace: {:?}", e)))?;

        let current_public_key_js = get_workspace_default_public_key(&current_workspace).await
            .map_err(|e| JsValue::from_str(&format!("Failed to get current public key: {:?}", e)))?;

        let current_public_key = current_public_key_js.as_string()
            .ok_or_else(|| JsValue::from_str("Current public key is not a string"))?;

        // 3. Check the request
        SIGSOCKET_CLIENT.with(|c| {
            let client = c.borrow();
            let client = client.as_ref().ok_or_else(|| JsValue::from_str("Not connected to SigSocket"))?;

            // Get the request
            let request = client.get_pending_request(request_id)
                .ok_or_else(|| JsValue::from_str("Request not found"))?;

            // Check if request matches current session
            let can_approve = request.target_public_key == current_public_key;

            console_log!("Can approve request {}: {} (current: {}, target: {})",
                        request_id, can_approve, current_public_key, request.target_public_key);

            Ok(can_approve)
        })
    }
    
    /// Approve a sign request and send the signature to the server
    /// 
    /// This performs the complete approval flow:
    /// 1. Validates the request can be approved
    /// 2. Signs the message using the vault
    /// 3. Sends the signature to the SigSocket server
    /// 4. Removes the request from pending list
    /// 
    /// # Arguments
    /// * `request_id` - The ID of the request to approve
    /// 
    /// # Returns
    /// * `Ok(signature)` - Base64-encoded signature that was sent
    /// * `Err(error)` - If approval failed
    #[wasm_bindgen]
    pub async fn approve_request(request_id: &str) -> Result<String, JsValue> {
        // 1. Validate we can approve this request
        if !Self::can_approve_request(request_id).await? {
            return Err(JsValue::from_str("Cannot approve this request"));
        }
        
        // 2. Get request details and sign the message
        let (message_bytes, original_request) = SIGSOCKET_CLIENT.with(|c| {
            let client = c.borrow();
            let client = client.as_ref().ok_or_else(|| JsValue::from_str("Not connected"))?;
            
            let request = client.get_pending_request(request_id)
                .ok_or_else(|| JsValue::from_str("Request not found"))?;
            
            // Decode the message
            let message_bytes = request.message_bytes()
                .map_err(|e| JsValue::from_str(&format!("Invalid message format: {}", e)))?;
            
            Ok::<(Vec<u8>, SignRequest), JsValue>((message_bytes, request.request.clone()))
        })?;
        
        // 3. Sign with vault
        let signature_result = sign_with_default_keypair(&message_bytes).await?;
        let signature_hex = signature_result.as_string()
            .ok_or_else(|| JsValue::from_str("Signature result is not a string"))?;

        // Convert hex signature to base64 for SigSocket protocol
        let signature_bytes = hex::decode(&signature_hex)
            .map_err(|e| JsValue::from_str(&format!("Invalid hex signature: {}", e)))?;
        let signature_base64 = base64::prelude::BASE64_STANDARD.encode(&signature_bytes);
        
        // 4. Get original message for response
        let original_message = SIGSOCKET_CLIENT.with(|c| {
            let client = c.borrow();
            let client = client.as_ref().ok_or_else(|| JsValue::from_str("Not connected"))?;

            let request = client.get_pending_request(request_id)
                .ok_or_else(|| JsValue::from_str("Request not found"))?;

            Ok::<String, JsValue>(request.request.message.clone())
        })?;

        // 5. Send response to server (create a new scope to avoid borrowing issues)
        {
            let client_ref = SIGSOCKET_CLIENT.with(|c| {
                c.borrow().as_ref().map(|client| client as *const SigSocketClient)
            }).ok_or_else(|| JsValue::from_str("Not connected"))?;

            // SAFETY: We know the client exists and we're using it synchronously
            let client = unsafe { &*client_ref };

            client.send_response(request_id, &original_message, &signature_base64).await
                .map_err(|e| JsValue::from_str(&format!("Failed to send response: {:?}", e)))?;

            console_log!("✅ WASM: Sent signature response to server for request: {}", request_id);
        }

        // 6. Remove the request after successful send
        SIGSOCKET_CLIENT.with(|c| {
            let mut client = c.borrow_mut();
            if let Some(client) = client.as_mut() {
                client.remove_pending_request(request_id);
                console_log!("✅ WASM: Removed request from pending list: {}", request_id);
            }
        });

        console_log!("🎉 WASM: Successfully approved and sent signature for request: {}", request_id);
        Ok(signature_base64)
    }
    
    /// Reject a sign request
    ///
    /// # Arguments
    /// * `request_id` - The ID of the request to reject
    /// * `reason` - The reason for rejection
    ///
    /// # Returns
    /// * `Ok(())` - Request rejected successfully
    /// * `Err(error)` - If rejection failed
    #[wasm_bindgen]
    pub async fn reject_request(request_id: &str, reason: &str) -> Result<(), JsValue> {
        // Send rejection to server first
        {
            let client_ref = SIGSOCKET_CLIENT.with(|c| {
                c.borrow().as_ref().map(|client| client as *const SigSocketClient)
            }).ok_or_else(|| JsValue::from_str("Not connected"))?;

            // SAFETY: We know the client exists and we're using it synchronously
            let client = unsafe { &*client_ref };

            client.send_rejection(request_id, reason).await
                .map_err(|e| JsValue::from_str(&format!("Failed to send rejection: {:?}", e)))?;

            console_log!("✅ WASM: Sent rejection to server for request: {}", request_id);
        }

        // Remove the request after successful send
        SIGSOCKET_CLIENT.with(|c| {
            let mut client = c.borrow_mut();
            if let Some(client) = client.as_mut() {
                client.remove_pending_request(request_id);
                console_log!("✅ WASM: Removed rejected request from pending list: {}", request_id);
            }
        });

        console_log!("🚫 WASM: Successfully rejected request: {} (reason: {})", request_id, reason);
        Ok(())
    }

    /// Get pending requests filtered by current workspace
    ///
    /// This returns only the requests that the current vault session can handle,
    /// based on the unlocked workspace and its public key.
    ///
    /// # Returns
    /// * `Ok(requests_json)` - JSON array of filtered requests
    /// * `Err(error)` - If filtering failed
    #[wasm_bindgen]
    pub async fn get_filtered_requests() -> Result<String, JsValue> {
        // If vault is locked, return empty array
        if !is_unlocked() {
            return Ok("[]".to_string());
        }

        // Get current workspace public key
        let current_workspace = get_current_keyspace_name()
            .map_err(|e| JsValue::from_str(&format!("Failed to get current workspace: {:?}", e)))?;

        let current_public_key_js = get_workspace_default_public_key(&current_workspace).await
            .map_err(|e| JsValue::from_str(&format!("Failed to get current public key: {:?}", e)))?;

        let current_public_key = current_public_key_js.as_string()
            .ok_or_else(|| JsValue::from_str("Current public key is not a string"))?;

        // Filter requests for current workspace
        SIGSOCKET_CLIENT.with(|c| {
            let client = c.borrow();
            let client = client.as_ref().ok_or_else(|| JsValue::from_str("Not connected to SigSocket"))?;

            let filtered_requests: Vec<_> = client.get_requests_for_public_key(&current_public_key);

            console_log!("Filtered requests: {} total, {} for current workspace",
                        client.pending_request_count(), filtered_requests.len());

            // Serialize and return
            serde_json::to_string(&filtered_requests)
                .map_err(|e| JsValue::from_str(&format!("Serialization error: {}", e)))
        })
    }

    /// Add a pending sign request (called when request arrives from server)
    ///
    /// # Arguments
    /// * `request_json` - JSON string containing the sign request
    ///
    /// # Returns
    /// * `Ok(())` - Request added successfully
    /// * `Err(error)` - If adding failed
    #[wasm_bindgen]
    pub fn add_pending_request(request_json: &str) -> Result<(), JsValue> {
        // Parse the request
        let request: SignRequest = serde_json::from_str(request_json)
            .map_err(|e| JsValue::from_str(&format!("Invalid request JSON: {}", e)))?;

        SIGSOCKET_CLIENT.with(|c| {
            let mut client = c.borrow_mut();
            let client = client.as_mut().ok_or_else(|| JsValue::from_str("Not connected to SigSocket"))?;

            // Get the connected public key as the target
            let target_public_key = client.connected_public_key()
                .ok_or_else(|| JsValue::from_str("No connected public key"))?
                .to_string();

            // Add the request
            client.add_pending_request(request, target_public_key);

            console_log!("Added pending request: {}", request_json);

            Ok(())
        })
    }

    /// Get connection status
    ///
    /// # Returns
    /// * `Ok(status_json)` - JSON object with connection status
    /// * `Err(error)` - If getting status failed
    #[wasm_bindgen]
    pub fn get_connection_status() -> Result<String, JsValue> {
        SIGSOCKET_CLIENT.with(|c| {
            let client = c.borrow();

            if let Some(client) = client.as_ref() {
                let status = serde_json::json!({
                    "is_connected": client.is_connected(),
                    "connected_public_key": client.connected_public_key(),
                    "pending_request_count": client.pending_request_count(),
                    "server_url": client.url()
                });

                Ok(status.to_string())
            } else {
                let status = serde_json::json!({
                    "is_connected": false,
                    "connected_public_key": null,
                    "pending_request_count": 0,
                    "server_url": null
                });

                Ok(status.to_string())
            }
        })
    }

    /// Clear all pending requests
    ///
    /// # Returns
    /// * `Ok(())` - Requests cleared successfully
    #[wasm_bindgen]
    pub fn clear_pending_requests() -> Result<(), JsValue> {
        SIGSOCKET_CLIENT.with(|c| {
            let mut client = c.borrow_mut();
            if let Some(client) = client.as_mut() {
                client.clear_pending_requests();
                console_log!("Cleared all pending requests");
            }
            Ok(())
        })
    }
}