let vault = null;
let isInitialized = false;
let currentSession = 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() };
  await chrome.storage.session.set({ cryptoVaultSession: currentSession });
  console.log('Session saved:', currentSession);
}

async function loadSession() {
  try {
    const result = await chrome.storage.session.get(['cryptoVaultSession']);
    if (result.cryptoVaultSession) {
      currentSession = result.cryptoVaultSession;
      console.log('Session loaded:', currentSession);
      return currentSession;
    }
  } catch (error) {
    console.error('Failed to load session:', error);
  }
  return null;
}

async function clearSession() {
  currentSession = null;
  await chrome.storage.session.remove(['cryptoVaultSession']);
  console.log('Session cleared');
}

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) {
        console.log('Session restored successfully for keyspace:', session.keyspace);
        return session;
      } else {
        console.log('Session expired, clearing...');
        await clearSession();
      }
    } catch (error) {
      console.error('Error checking session validity:', error);
      await clearSession();
    }
  }
  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,
  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,
      lock_session
    };

    isInitialized = true;
    console.log('CryptoVault initialized successfully');

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

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

// 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 saveSession(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':
          console.log('Background: listing keypairs...');
          console.log('Background: vault object:', vault);
          console.log('Background: vault.list_keypairs function:', vault.list_keypairs);

          // Check if session is unlocked first
          const isUnlocked = vault.is_unlocked();
          console.log('Background: is session unlocked?', isUnlocked);

          if (!isUnlocked) {
            console.log('Background: Session is not unlocked, cannot list keypairs');
            return { success: false, error: 'Session is not unlocked' };
          }

          try {
            const keypairsRaw = await vault.list_keypairs();
            console.log('Background: keypairs raw result:', keypairsRaw);
            console.log('Background: keypairs type:', typeof keypairsRaw);

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

            console.log('Background: parsed keypairs:', keypairs);
            console.log('Background: parsed keypairs type:', typeof keypairs);
            console.log('Background: keypairs array length:', Array.isArray(keypairs) ? keypairs.length : 'not an array');

            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 'lockSession':
          vault.lock_session();
          await clearSession();
          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();
});