let vault = null;
let isInitialized = false;
let currentSession = null;
let keepAliveInterval = null;

// Utility function to convert Uint8Array to hex
function toHex(uint8Array) {
  return Array.from(uint8Array)
    .map(b => b.toString(16).padStart(2, '0'))
    .join('');
}



// Session persistence functions
async function saveSession(keyspace) {
  currentSession = { keyspace, timestamp: Date.now() };

  // Save to both session and local storage for better persistence
  try {
    await chrome.storage.session.set({ cryptoVaultSession: currentSession });
    await chrome.storage.local.set({ cryptoVaultSessionBackup: currentSession });
  } catch (error) {
    console.error('Failed to save session:', error);
  }
}

async function loadSession() {
  try {
    // Try session storage first
    let result = await chrome.storage.session.get(['cryptoVaultSession']);
    if (result.cryptoVaultSession) {
      currentSession = result.cryptoVaultSession;
      return currentSession;
    }

    // Fallback to local storage
    result = await chrome.storage.local.get(['cryptoVaultSessionBackup']);
    if (result.cryptoVaultSessionBackup) {
      currentSession = result.cryptoVaultSessionBackup;
      // Restore to session storage
      await chrome.storage.session.set({ cryptoVaultSession: currentSession });
      return currentSession;
    }
  } catch (error) {
    console.error('Failed to load session:', error);
  }
  return null;
}

async function clearSession() {
  currentSession = null;
  try {
    await chrome.storage.session.remove(['cryptoVaultSession']);
    await chrome.storage.local.remove(['cryptoVaultSessionBackup']);
  } catch (error) {
    console.error('Failed to clear session:', error);
  }
}

// Keep service worker alive
function startKeepAlive() {
  if (keepAliveInterval) {
    clearInterval(keepAliveInterval);
  }

  // Ping every 20 seconds to keep service worker alive
  keepAliveInterval = setInterval(() => {
    // Simple operation to keep service worker active
    chrome.storage.session.get(['keepAlive']).catch(() => {
      // Ignore errors
    });
  }, 20000);
}

function stopKeepAlive() {
  if (keepAliveInterval) {
    clearInterval(keepAliveInterval);
    keepAliveInterval = null;
  }
}

// Enhanced session management with keep-alive
async function saveSessionWithKeepAlive(keyspace) {
  await saveSession(keyspace);
  startKeepAlive();
}

async function clearSessionWithKeepAlive() {
  await clearSession();
  stopKeepAlive();
}

async function restoreSession() {
  const session = await loadSession();
  if (session && vault) {
    try {
      // Check if the session is still valid by testing if vault is unlocked
      const isUnlocked = vault.is_unlocked();
      if (isUnlocked) {
        // Restart keep-alive for restored session
        startKeepAlive();
        return session;
      } else {
        await clearSessionWithKeepAlive();
      }
    } catch (error) {
      console.error('Error checking session validity:', error);
      await clearSessionWithKeepAlive();
    }
  }
  return null;
}

// Import WASM module functions
import init, {
  create_keyspace,
  init_session,
  is_unlocked,
  add_keypair,
  list_keypairs,
  select_keypair,
  current_keypair_metadata,
  current_keypair_public_key,
  sign,
  verify,
  encrypt_data,
  decrypt_data,
  lock_session
} from './wasm/wasm_app.js';

// Initialize WASM module
async function initVault() {
  try {
    if (vault && isInitialized) return vault;

    // Initialize with the WASM file
    const wasmUrl = chrome.runtime.getURL('wasm/wasm_app_bg.wasm');
    await init(wasmUrl);

    // Create a vault object with all the imported functions
    vault = {
      create_keyspace,
      init_session,
      is_unlocked,
      add_keypair,
      list_keypairs,
      select_keypair,
      current_keypair_metadata,
      current_keypair_public_key,
      sign,
      verify,
      encrypt_data,
      decrypt_data,
      lock_session
    };

    isInitialized = true;

    // Try to restore previous session
    await restoreSession();

    return vault;
  } catch (error) {
    console.error('Failed to initialize CryptoVault:', error);
    throw error;
  }
}

// Handle popup connection/disconnection
chrome.runtime.onConnect.addListener((port) => {
  if (port.name === 'popup') {
    // If we have an active session, ensure keep-alive is running
    if (currentSession) {
      startKeepAlive();
    }

    port.onDisconnect.addListener(() => {
      // Keep the keep-alive running even after popup disconnects
      // This ensures session persistence across popup closes
    });
  }
});

// Handle messages from popup and content scripts
chrome.runtime.onMessage.addListener((request, _sender, sendResponse) => {
  const handleRequest = async () => {
    try {
      if (!vault) {
        await initVault();
      }

      switch (request.action) {
        case 'createKeyspace':
          await vault.create_keyspace(request.keyspace, request.password);
          return { success: true };

        case 'initSession':
          await vault.init_session(request.keyspace, request.password);
          await saveSessionWithKeepAlive(request.keyspace);
          return { success: true };

        case 'isUnlocked':
          const unlocked = vault.is_unlocked();
          return { success: true, unlocked };

        case 'addKeypair':
          const result = await vault.add_keypair(request.keyType, request.metadata);
          return { success: true, result };

        case 'listKeypairs':
          // Check if session is unlocked first
          const isUnlocked = vault.is_unlocked();

          if (!isUnlocked) {
            return { success: false, error: 'Session is not unlocked' };
          }

          try {
            const keypairsRaw = await vault.list_keypairs();

            // Parse JSON string if needed
            let keypairs;
            if (typeof keypairsRaw === 'string') {
              keypairs = JSON.parse(keypairsRaw);
            } else {
              keypairs = keypairsRaw;
            }

            return { success: true, keypairs };
          } catch (listError) {
            console.error('Background: Error calling list_keypairs:', listError);
            throw listError;
          }

        case 'selectKeypair':
          vault.select_keypair(request.keyId);
          return { success: true };

        case 'getCurrentKeypairMetadata':
          const metadata = vault.current_keypair_metadata();
          return { success: true, metadata };

        case 'getCurrentKeypairPublicKey':
          const publicKey = vault.current_keypair_public_key();
          const hexKey = toHex(publicKey);
          return { success: true, publicKey: hexKey };

        case 'sign':
          const signature = await vault.sign(new Uint8Array(request.message));
          return { success: true, signature };

        case 'encrypt':
          // Check if session is unlocked
          if (!vault.is_unlocked()) {
            return { success: false, error: 'Session is not unlocked' };
          }

          try {
            // Convert message to Uint8Array for WASM
            const messageBytes = new TextEncoder().encode(request.message);

            // Use WASM encrypt_data function with ChaCha20-Poly1305
            const encryptedData = await vault.encrypt_data(messageBytes);

            // Convert result to base64 for easy handling
            const encryptedMessage = btoa(String.fromCharCode(...new Uint8Array(encryptedData)));
            return { success: true, encryptedMessage };
          } catch (error) {
            console.error('Encryption error:', error);
            return { success: false, error: error.message };
          }

        case 'decrypt':
          // Check if session is unlocked
          if (!vault.is_unlocked()) {
            return { success: false, error: 'Session is not unlocked' };
          }

          try {
            // Convert base64 back to Uint8Array
            const encryptedBytes = new Uint8Array(atob(request.encryptedMessage).split('').map(c => c.charCodeAt(0)));

            // Use WASM decrypt_data function with ChaCha20-Poly1305
            const decryptedData = await vault.decrypt_data(encryptedBytes);

            // Convert result back to string
            const decryptedMessage = new TextDecoder().decode(new Uint8Array(decryptedData));
            return { success: true, decryptedMessage };
          } catch (error) {
            console.error('Decryption error:', error);
            return { success: false, error: error.message };
          }

        case 'verify':
          // Check if a keypair is selected
          try {
            const metadata = vault.current_keypair_metadata();
            if (!metadata) {
              return { success: false, error: 'No keypair selected' };
            }

            // Use WASM verify function
            const isValid = await vault.verify(new Uint8Array(request.message), request.signature);
            return { success: true, isValid };
          } catch (error) {
            console.error('Verification error:', error);
            return { success: false, error: error.message };
          }

        case 'lockSession':
          vault.lock_session();
          await clearSessionWithKeepAlive();
          return { success: true };

        case 'getStatus':
          const status = vault ? vault.is_unlocked() : false;
          const session = await loadSession();
          return {
            success: true,
            status,
            session: session ? { keyspace: session.keyspace } : null
          };

        default:
          throw new Error('Unknown action: ' + request.action);
      }
    } catch (error) {
      console.error('Background script error:', error);
      return { success: false, error: error.message };
    }
  };

  handleRequest().then(sendResponse);
  return true; // Keep the message channel open for async response
});

// Initialize vault when extension starts
chrome.runtime.onStartup.addListener(() => {
  initVault();
});

chrome.runtime.onInstalled.addListener(() => {
  initVault();
});