import { create } from 'zustand';
import { getWasmModule, stringToUint8Array } from '../wasm/wasmHelper';
import { getChromeApi } from '../utils/chromeApi';

// Import Chrome types
/// <reference types="chrome" />

interface KeypairMetadata {
  id: string;
  type: string;
  name?: string;
  description?: string;
  createdAt: number;
}

interface SessionState {
  isSessionUnlocked: boolean;
  currentKeyspace: string | null;
  currentKeypair: KeypairMetadata | null;
  availableKeypairs: KeypairMetadata[];
  isWebSocketConnected: boolean;
  webSocketUrl: string | null;
  isWasmLoaded: boolean;
  
  // Actions
  initWasm: () => Promise<boolean>;
  checkSessionStatus: () => Promise<boolean>;
  unlockSession: (keyspace: string, password: string) => Promise<boolean>;
  lockSession: () => Promise<boolean>;
  createKeyspace: (keyspace: string, password: string) => Promise<boolean>;
  listKeypairs: () => Promise<KeypairMetadata[]>;
  selectKeypair: (keypairId: string) => Promise<boolean>;
  createKeypair: (type: string, metadata?: Record<string, any>) => Promise<string>;
  connectWebSocket: (url: string) => Promise<boolean>;
  disconnectWebSocket: () => Promise<boolean>;
  executeScript: (script: string) => Promise<string>;
  signMessage: (message: string) => Promise<string>;
}

// Create the store
export const useSessionStore = create<SessionState>((set: any, get: any) => ({
  isSessionUnlocked: false,
  currentKeyspace: null,
  currentKeypair: null,
  availableKeypairs: [],
  isWebSocketConnected: false,
  webSocketUrl: null,
  isWasmLoaded: false,
  
  // Initialize WASM module
  initWasm: async () => {
    try {
      set({ isWasmLoaded: true });
      return true;
    } catch (error) {
      console.error('Failed to initialize WASM module:', error);
      return false;
    }
  },
  
  // Check if a session is currently active
  checkSessionStatus: async () => {
    try {
      // First check with the background service worker
      const chromeApi = getChromeApi();
      const response = await chromeApi.runtime.sendMessage({ type: 'SESSION_STATUS' });
      
      if (response && response.active) {
        // If session is active in the background, check with WASM
        try {
          const wasmModule = await getWasmModule();
          const isUnlocked = wasmModule.is_unlocked();
          
          if (isUnlocked) {
            // Get current keypair metadata if available
            try {
              const keypairMetadata = await wasmModule.current_keypair_metadata();
              const parsedMetadata = JSON.parse(keypairMetadata);
              
              set({ 
                isSessionUnlocked: true,
                currentKeypair: parsedMetadata
              });
              
              // Load keypairs
              await get().listKeypairs();
            } catch (e) {
              // No keypair selected, but session is unlocked
              set({ isSessionUnlocked: true });
            }
            return true;
          }
        } catch (wasmError) {
          console.error('WASM error checking session status:', wasmError);
        }
      }
      
      set({ isSessionUnlocked: false });
      return false;
    } catch (error) {
      console.error('Failed to check session status:', error);
      set({ isSessionUnlocked: false });
      return false;
    }
  },
  
  // Unlock a session with keyspace and password
  unlockSession: async (keyspace: string, password: string) => {
    try {
      const wasmModule = await getWasmModule();
      
      // Call the WASM init_session function
      await wasmModule.init_session(keyspace, password);
      
      // Initialize Rhai environment
      wasmModule.init_rhai_env();
      
      // Notify background service worker
      const chromeApi = getChromeApi();
      await chromeApi.runtime.sendMessage({ type: 'SESSION_UNLOCK' });
      
      set({ 
        isSessionUnlocked: true,
        currentKeyspace: keyspace,
        currentKeypair: null
      });
      
      // Load keypairs after unlocking
      const keypairs = await get().listKeypairs();
      set({ availableKeypairs: keypairs });
      
      return true;
    } catch (error) {
      console.error('Failed to unlock session:', error);
      return false;
    }
  },
  
  // Lock the current session
  lockSession: async () => {
    try {
      const wasmModule = await getWasmModule();
      
      // Call the WASM lock_session function
      wasmModule.lock_session();
      
      // Notify background service worker
      const chromeApi = getChromeApi();
      await chromeApi.runtime.sendMessage({ type: 'SESSION_LOCK' });
      
      set({
        isSessionUnlocked: false,
        currentKeyspace: null,
        currentKeypair: null,
        availableKeypairs: [],
        isWebSocketConnected: false,
        webSocketUrl: null
      });
      
      return true;
    } catch (error) {
      console.error('Failed to lock session:', error);
      return false;
    }
  },
  
  // Create a new keyspace
  createKeyspace: async (keyspace: string, password: string) => {
    try {
      const wasmModule = await getWasmModule();
      
      // Call the WASM create_keyspace function
      await wasmModule.create_keyspace(keyspace, password);
      
      // Initialize Rhai environment
      wasmModule.init_rhai_env();
      
      // Notify background service worker
      const chromeApi = getChromeApi();
      await chromeApi.runtime.sendMessage({ type: 'SESSION_UNLOCK' });
      
      set({ 
        isSessionUnlocked: true,
        currentKeyspace: keyspace,
        currentKeypair: null,
        availableKeypairs: []
      });
      
      return true;
    } catch (error) {
      console.error('Failed to create keyspace:', error);
      return false;
    }
  },
  
  // List all keypairs in the current keyspace
  listKeypairs: async () => {
    try {
      console.log('Listing keypairs from WASM module');
      const wasmModule = await getWasmModule();
      console.log('WASM module loaded, calling list_keypairs');
      
      // Call the WASM list_keypairs function
      let keypairsJson;
      try {
        keypairsJson = await wasmModule.list_keypairs();
        console.log('Raw keypairs JSON from WASM:', keypairsJson);
      } catch (listError) {
        console.error('Error calling list_keypairs:', listError);
        throw new Error(`Failed to list keypairs: ${listError.message || listError}`);
      }
      
      let keypairs;
      try {
        keypairs = JSON.parse(keypairsJson);
        console.log('Parsed keypairs object:', keypairs);
      } catch (parseError) {
        console.error('Error parsing keypairs JSON:', parseError);
        throw new Error(`Failed to parse keypairs JSON: ${parseError.message}`);
      }
      
      // Transform the keypairs to our expected format
      const formattedKeypairs: KeypairMetadata[] = keypairs.map((keypair: any, index: number) => {
        console.log(`Processing keypair at index ${index}:`, keypair);
        return {
          id: keypair.id, // Use the actual keypair ID from the WASM module
          type: keypair.key_type || 'Unknown',
          name: keypair.metadata?.name,
          description: keypair.metadata?.description,
          createdAt: keypair.metadata?.created_at || Date.now()
        };
      });
      
      console.log('Formatted keypairs for UI:', formattedKeypairs);
      set({ availableKeypairs: formattedKeypairs });
      return formattedKeypairs;
    } catch (error) {
      console.error('Failed to list keypairs:', error);
      return [];
    }
  },
  
  // Select a keypair for use
  selectKeypair: async (keypairId: string) => {
    try {
      console.log('Selecting keypair with ID:', keypairId);
      
      // First, let's log the available keypairs to see what we have
      const { availableKeypairs } = get();
      console.log('Available keypairs:', JSON.stringify(availableKeypairs));
      
      const wasmModule = await getWasmModule();
      console.log('WASM module loaded, attempting to select keypair');
      
      try {
        // Call the WASM select_keypair function
        await wasmModule.select_keypair(keypairId);
        console.log('Successfully selected keypair in WASM');
      } catch (selectError) {
        console.error('Error in WASM select_keypair:', selectError);
        throw new Error(`select_keypair error: ${selectError.message || selectError}`);
      }
      
      // Find the keypair in our availableKeypairs list
      const selectedKeypair = availableKeypairs.find((kp: KeypairMetadata) => kp.id === keypairId);
      
      if (selectedKeypair) {
        console.log('Found keypair in available list, setting as current');
        set({ currentKeypair: selectedKeypair });
      } else {
        console.log('Keypair not found in available list, creating new entry from available data');
        // If not found in our list (rare case), create a new entry with what we know
        // Since we can't get metadata from WASM, use what we have from the keypair list
        const matchingKeypair = availableKeypairs.find(k => k.id === keypairId);
        
        if (matchingKeypair) {
          set({ currentKeypair: matchingKeypair });
        } else {
          // Last resort: create a minimal keypair entry
          const newKeypair: KeypairMetadata = {
            id: keypairId,
            type: 'Unknown',
            name: `Keypair ${keypairId.substring(0, 8)}...`,
            createdAt: Date.now()
          };
          
          set({ currentKeypair: newKeypair });
        }
      }
      
      return true;
    } catch (error) {
      console.error('Failed to select keypair:', error);
      throw error; // Re-throw to show error in UI
    }
  },
  
  // Create a new keypair
  createKeypair: async (type: string, metadata?: Record<string, any>) => {
    try {
      const wasmModule = await getWasmModule();
      
      // Format metadata for WASM
      const metadataJson = metadata ? JSON.stringify({
        name: metadata.name,
        description: metadata.description,
        created_at: Date.now()
      }) : undefined;
      
      // Call the WASM add_keypair function
      const keypairId = await wasmModule.add_keypair(type, metadataJson);
      
      // Refresh the keypair list
      await get().listKeypairs();
      
      return keypairId;
    } catch (error) {
      console.error('Failed to create keypair:', error);
      throw error;
    }
  },
  
  // Connect to a WebSocket server
  connectWebSocket: async (url: string) => {
    try {
      const wasmModule = await getWasmModule();
      const { currentKeypair } = get();
      
      if (!currentKeypair) {
        throw new Error('No keypair selected');
      }
      
      // Get the public key from WASM
      const publicKeyArray = await wasmModule.current_keypair_public_key();
      const publicKeyHex = Array.from(publicKeyArray)
        .map(b => b.toString(16).padStart(2, '0'))
        .join('');
      
      // Connect to WebSocket via background service worker
      const chromeApi = getChromeApi();
      const response = await chromeApi.runtime.sendMessage({
        type: 'CONNECT_WEBSOCKET',
        serverUrl: url,
        publicKey: publicKeyHex
      });
      
      if (response && response.success) {
        set({
          isWebSocketConnected: true,
          webSocketUrl: url
        });
        return true;
      } else {
        throw new Error(response?.error || 'Failed to connect to WebSocket server');
      }
    } catch (error) {
      console.error('Failed to connect to WebSocket:', error);
      return false;
    }
  },
  
  // Disconnect from WebSocket server
  disconnectWebSocket: async () => {
    try {
      // Disconnect via background service worker
      const chromeApi = getChromeApi();
      const response = await chromeApi.runtime.sendMessage({
        type: 'DISCONNECT_WEBSOCKET'
      });
      
      if (response && response.success) {
        set({
          isWebSocketConnected: false,
          webSocketUrl: null
        });
        return true;
      } else {
        throw new Error(response?.error || 'Failed to disconnect from WebSocket server');
      }
    } catch (error) {
      console.error('Failed to disconnect from WebSocket:', error);
      return false;
    }
  },
  
  // Execute a Rhai script
  executeScript: async (script: string) => {
    try {
      const wasmModule = await getWasmModule();
      
      // Call the WASM run_rhai function
      const result = await wasmModule.run_rhai(script);
      return result;
    } catch (error) {
      console.error('Failed to execute script:', error);
      throw error;
    }
  },
  
  // Sign a message with the current keypair
  signMessage: async (message: string) => {
    try {
      const wasmModule = await getWasmModule();
      
      // Convert message to Uint8Array
      const messageBytes = stringToUint8Array(message);
      
      // Call the WASM sign function
      const signature = await wasmModule.sign(messageBytes);
      return signature;
    } catch (error) {
      console.error('Failed to sign message:', error);
      throw error;
    }
  }
}));
