import React, { useState, useEffect } from 'react';
import KeyspaceManager from './KeyspaceManager';
import KeypairManager from './KeypairManager';
import SignMessage from './SignMessage';
import * as wasmHelper from './WasmHelper';

function App() {
  const [wasmState, setWasmState] = useState({
    loading: false,
    initialized: false,
    error: null
  });
  const [locked, setLocked] = useState(true);
  const [keyspaces, setKeyspaces] = useState([]);
  const [currentKeyspace, setCurrentKeyspace] = useState('');
  const [keypairs, setKeypairs] = useState([]); // [{id, label, publicKey}]
  const [selectedKeypair, setSelectedKeypair] = useState('');
  const [signature, setSignature] = useState('');
  const [loading, setLoading] = useState(false);
  const [status, setStatus] = useState('');

  // Load WebAssembly on component mount
  useEffect(() => {
    async function initWasm() {
      try {
        setStatus('Loading WebAssembly module...');
        await wasmHelper.loadWasmModule();
        setWasmState(wasmHelper.getWasmState());
        setStatus('WebAssembly module loaded');
        // Load session state
        await refreshStatus();
      } catch (error) {
        console.error('Failed to load WebAssembly:', error);
        setStatus('Error loading WebAssembly: ' + (error.message || 'Unknown error'));
      }
    }
    
    initWasm();
  }, []);

  // Fetch status from background on mount
  async function refreshStatus() {
    const state = await wasmHelper.getSessionState();
    setCurrentKeyspace(state.currentKeyspace || '');
    setKeypairs(state.keypairs || []);
    setSelectedKeypair(state.selectedKeypair || '');
    setLocked(!state.currentKeyspace);
    
    // For demo: collect all keyspaces from storage
    if (state.keypairs && state.keypairs.length > 0) {
      setKeyspaces([state.currentKeyspace]);
    } else {
      setKeyspaces([state.currentKeyspace].filter(Boolean));
    }
  }

  // Session unlock/create
  const handleUnlock = async (keyspace, password) => {
    if (!wasmState.initialized) {
      setStatus('WebAssembly module not loaded');
      return;
    }
    
    setLoading(true);
    setStatus('Unlocking...');
    try {
      await wasmHelper.initSession(keyspace, password);
      setCurrentKeyspace(keyspace);
      setLocked(false);
      setStatus('Session unlocked!');
      await refreshStatus();
    } catch (e) {
      setStatus('Unlock failed: ' + e);
    }
    setLoading(false);
  };
  
  const handleCreateKeyspace = async (keyspace, password) => {
    if (!wasmState.initialized) {
      setStatus('WebAssembly module not loaded');
      return;
    }
    
    setLoading(true);
    setStatus('Creating keyspace...');
    try {
      await wasmHelper.initSession(keyspace, password);
      setCurrentKeyspace(keyspace);
      setLocked(false);
      setStatus('Keyspace created and unlocked!');
      await refreshStatus();
    } catch (e) {
      setStatus('Create failed: ' + e);
    }
    setLoading(false);
  };
  
  const handleLock = async () => {
    if (!wasmState.initialized) {
      setStatus('WebAssembly module not loaded');
      return;
    }
    
    setLoading(true);
    setStatus('Locking...');
    try {
      await wasmHelper.lockSession();
      setLocked(true);
      setCurrentKeyspace('');
      setKeypairs([]);
      setSelectedKeypair('');
      setStatus('Session locked.');
      await refreshStatus();
    } catch (e) {
      setStatus('Lock failed: ' + e);
    }
    setLoading(false);
  };
  
  const handleSelectKeypair = async (id) => {
    if (!wasmState.initialized) {
      setStatus('WebAssembly module not loaded');
      return;
    }
    
    setLoading(true);
    setStatus('Selecting keypair...');
    try {
      await wasmHelper.selectKeypair(id);
      setSelectedKeypair(id);
      setStatus('Keypair selected.');
      await refreshStatus();
    } catch (e) {
      setStatus('Select failed: ' + e);
    }
    setLoading(false);
  };
  
  const handleCreateKeypair = async () => {
    if (!wasmState.initialized) {
      setStatus('WebAssembly module not loaded');
      return;
    }
    
    setLoading(true);
    setStatus('Creating keypair...');
    try {
      const keyId = await wasmHelper.addKeypair();
      setStatus('Keypair created. ID: ' + keyId);
      await refreshStatus();
    } catch (e) {
      setStatus('Create failed: ' + e);
    }
    setLoading(false);
  };
  
  const handleSign = async (message) => {
    if (!wasmState.initialized) {
      setStatus('WebAssembly module not loaded');
      return;
    }
    
    setLoading(true);
    setStatus('Signing message...');
    try {
      if (!selectedKeypair) {
        throw new Error('No keypair selected');
      }
      const sig = await wasmHelper.sign(message);
      setSignature(sig);
      setStatus('Message signed!');
    } catch (e) {
      setStatus('Signing failed: ' + e);
      setSignature('');
    }
    setLoading(false);
  };

  return (
    <div className="App">
      <h1>Modular Vault Extension</h1>
      {wasmState.error && (
        <div className="error">
          WebAssembly Error: {wasmState.error}
        </div>
      )}
      <KeyspaceManager
        keyspaces={keyspaces}
        onUnlock={handleUnlock}
        onCreate={handleCreateKeyspace}
        locked={locked}
        onLock={handleLock}
        currentKeyspace={currentKeyspace}
      />
      {!locked && (
        <>
          <KeypairManager
            keypairs={keypairs}
            onSelect={handleSelectKeypair}
            onCreate={handleCreateKeypair}
            selectedKeypair={selectedKeypair}
          />
          {selectedKeypair && (
            <SignMessage
              onSign={handleSign}
              signature={signature}
              loading={loading}
            />
          )}
        </>
      )}
      <div className="status" style={{marginTop: '1rem', minHeight: 24}}>
        {status}
      </div>
    </div>
  );
}

export default App;