import React, { useState, useEffect } from 'react';
import { ContentEditor } from '@/components/editors/ContentEditor';
import { FloatingButton } from '@/components/FloatingButton';
import { Header, Tabs } from '@/components/Header';
import { MemoryEditor } from '@/components/editors/MemoryEditor';
import { Sidebar } from '@/components/Sidebar';
import { StandardEditor } from '@/components/editors/StandardEditor';
import { useWebSocket, WebSocketState } from '@/hooks/useWebSocket';

export const LlmState = {
  IDLE: 'IDLE',
  INFERENCE: 'INFERENCE'
};

const App = () => {
  // Editor content state
  const [generatedContext, setGeneratedContext] = useState('');
  const [modifiedContext, setModifiedContext] = useState('');
  const [contextDirty, setContextDirty] = useState(false);
  const [generatedResponse, setGeneratedResponse] = useState('');
  const [modifiedResponse, setModifiedResponse] = useState('');
  const [input, setInput] = useState('');
  const [output, setOutput] = useState('');
  
  // LLM state
  const [llms, setLlms] = useState({});
  const [activeLlm, setActiveLlm] = useState(null);

  // Auto approver state
  const [autoApproverConfig, setAutoApproverConfig] = useState({
    context_enabled: false,
    response_enabled: false,
    context_timeout: 5.0,
    response_timeout: 10.0,
    llm_name: null
  });
  
  // UI state
  const [activeTab, setActiveTab] = useState(Tabs.CONTEXT);
  const [showDiff, setShowDiff] = useState(false);
  const [showSidebar, setShowSidebar] = useState(false);
  const [autoScroll, setAutoScroll] = useState(false);

  // WebSocket connections
  const wsRoot = `${window.location.protocol === 'https:' ? 'wss:' : 'ws:'}//${window.location.host}/ws`
  const llmWs = useWebSocket(`${wsRoot}/llms`);
  const contextWs = useWebSocket(`${wsRoot}/context`);
  const responseWs = useWebSocket(`${wsRoot}/response`);
  const stdoutWs = useWebSocket(`${wsRoot}/stdout`);
  const autoApproverWs = useWebSocket(`${wsRoot}/auto_approver`);
  const memoryWs = useWebSocket(`${wsRoot}/memory`);
  const [entries, setEntries] = useState([]);
  const [wsState, setWsState] = useState(WebSocketState.DISCONNECTED);

  // Handle llm state changes
  useEffect(() => {
    llmWs.addMessageHandler((data) => {
      setLlms(prevLlms => ({
        ...prevLlms,
        [data.llm]: data.state
      }));
    });
  }, []);

  // Handle context changes
  useEffect(() => {
    contextWs.addMessageHandler((data) => {
      setContextDirty(false);
      setModifiedContext(data.context);
      if (data.generated) {
        setGeneratedContext(data.context);
      }
    });
  }, []);

  // Handle response changes
  useEffect(() => {
    responseWs.addMessageHandler((data) => {
      setModifiedResponse(data.response);
      setGeneratedResponse(data.response);
    });
  }, []);

  // Handle stdout changes
  useEffect(() => {
    stdoutWs.addMessageHandler((data) => {
      setOutput(data.output);
    });
  }, []);

  // Handle auto approver config updates
  useEffect(() => {
    autoApproverWs.addMessageHandler((data) => {
      setAutoApproverConfig(data.config);
    });
  }, []);

  // Handle memory state changes
  useEffect(() => {
    memoryWs.addMessageHandler((data) => {
      if (data.type === 'memory_state') {
        setEntries(data.entries);
      }
    });
  }, []);

  // Initialize active llm
  useEffect(() => {
    if (activeLlm === null && Object.keys(llms).length > 0) {
      setActiveLlm(Object.keys(llms)[0]);
    }
  }, [llms, activeLlm]);

  const handleNextLlm = () => {
    const llmNames = Object.keys(llms);
    const currentIndex = llmNames.indexOf(activeLlm);
    const nextIndex = (currentIndex + 1) % llmNames.length;
    setActiveLlm(llmNames[nextIndex]);
  };

  const handleApprove = () => {
    fetch('/api/response/approve', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({ response: modifiedResponse })
    });
  };

  const handleInference = () => {
    fetch('/api/context', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({ context: modifiedContext })
    });

    fetch('/api/response', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({ response: modifiedResponse })
    });
    
    fetch(`/api/inference/${activeLlm}`, {
      method: 'POST',
    });
  };

  const handleStop = () => {
    // We still specify which LLM to stop inference for
    fetch(`/api/inference/${activeLlm}/stop`, {
      method: 'POST'
    });
  };

  const handleSendInput = () => {
    if (input.trim()) {
      fetch('/api/input', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ input: input })
      });
      setInput('');
    }
  };

  const handleClearOutput = () => {
    fetch('/api/clear', { method: 'POST' });
  };

  const handleContextEdit = (context) => {
    setContextDirty(true);
    // When context is edited, we need to reset all LLM states to IDLE
    // since responses are no longer valid
    const resetLlms = {};
    for (const llm of Object.keys(llms)) {
      resetLlms[llm] = LlmState.IDLE;
    }    
    setLlms(resetLlms);
    setModifiedContext(context);
    
    // Reset response buffers
    setGeneratedResponse('');
    setModifiedResponse('');
  };

  const handleResponseEdit = (response) => {
    setModifiedResponse(response);
  };

  const handleAutoApproverConfigChange = async (config) => {
    try {
      const response = await fetch('/api/auto_approver/config', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify(config)
      });
      if (!response.ok) {
        throw new Error('Failed to update auto approver config');
      }
    } catch (error) {
      console.error('Error updating auto approver config:', error);
    }
  };

  const handleCreateEntry = async (entry) => {
    const response = await fetch('/api/memory/entry', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify(entry)
    });
    if (!response.ok) throw new Error('Failed to create entry');
  };

  const handleSaveEntry = async (id, entry) => {
    const response = await fetch(`/api/memory/entry/${id}`, {
      method: 'PUT',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify(entry)
    });
    if (!response.ok) throw new Error('Failed to update entry');
  };

  const handleDeleteEntry = async (id) => {
    const response = await fetch(`/api/memory/entry/${id}`, {
      method: 'DELETE'
    });
    if (!response.ok) throw new Error('Failed to delete entry');
  };

  const handleResetEntry = async (id) => {
    const response = await fetch(`/api/memory/entry/${id}/reset`, {
      method: 'POST'
    });
    if (!response.ok) throw new Error('Failed to reset entry');
  };

  const handleUpdateEntry = async (id) => {
    const response = await fetch(`/api/memory/entry/${id}/update`, {
      method: 'POST'
    });
    if (!response.ok) throw new Error('Failed to update entry');
  };

  // Determine active LLM's current state for UI feedback
  const activeLlmState = activeLlm ? llms[activeLlm] : null;

  // Track overall WebSocket connection state
  useEffect(() => {
    if (llmWs.wsState === WebSocketState.CONNECTED && 
        contextWs.wsState === WebSocketState.CONNECTED && 
        responseWs.wsState === WebSocketState.CONNECTED && 
        stdoutWs.wsState === WebSocketState.CONNECTED &&
        autoApproverWs.wsState === WebSocketState.CONNECTED &&
        memoryWs.wsState === WebSocketState.CONNECTED) {
      setWsState(WebSocketState.CONNECTED);
    } else if (llmWs.wsState === WebSocketState.CONNECTING || 
        contextWs.wsState === WebSocketState.CONNECTING || 
        responseWs.wsState === WebSocketState.CONNECTING || 
        stdoutWs.wsState === WebSocketState.CONNECTING ||
        autoApproverWs.wsState === WebSocketState.CONNECTING ||
        memoryWs.wsState === WebSocketState.CONNECTING) {
      setWsState(WebSocketState.CONNECTING);
    } else {
      setWsState(WebSocketState.DISCONNECTED);
    }
  }, [llmWs.wsState, contextWs.wsState, responseWs.wsState, stdoutWs.wsState, autoApproverWs.wsState, memoryWs.wsState]);

  return (
    <div className="h-screen bg-gray-100 flex flex-col overflow-hidden">
      <Header
        wsState={wsState}
        agentState={activeLlmState}
        activeTab={activeTab}
        setActiveTab={setActiveTab}
      />

      <div className="flex flex-1 overflow-hidden">
        <main className="flex-1 p-4 overflow-auto">
          {activeTab === Tabs.CONTEXT && (
            <ContentEditor
              showDiff={showDiff}
              originalContent={generatedContext}
              modifiedContent={modifiedContext}
              onChange={handleContextEdit}
              autoScroll={autoScroll}
            />
          )}
          {activeTab === Tabs.RESPONSE && (
            <ContentEditor
              showDiff={showDiff}
              originalContent={generatedResponse}
              modifiedContent={modifiedResponse}
              onChange={handleResponseEdit}
            />
          )}
          {activeTab === Tabs.INPUT && (
            <StandardEditor
              content={input}
              onChange={value => setInput(value)}
              language="plaintext"
            />
          )}
          {activeTab === Tabs.OUTPUT && (
            <StandardEditor
              content={output}
              readOnly={true}
              language="plaintext"
            />
          )}
          {activeTab === Tabs.MEMORY && (
            <MemoryEditor
              entries={entries}
              onCreateEntry={handleCreateEntry}
              onSaveEntry={handleSaveEntry} 
              onDeleteEntry={handleDeleteEntry}
              onResetEntry={handleResetEntry}
              onUpdateEntry={handleUpdateEntry}
            />
          )}
        </main>

        <Sidebar
          showSidebar={showSidebar}
          llms={llms}
          activeLlm={activeLlm}
          llmState={activeLlmState}
          showDiff={showDiff}
          autoScroll={autoScroll}
          onAutoScrollChange={setAutoScroll}
          input={input}
          output={output}
          autoApproverConfig={autoApproverConfig}
          onApprove={handleApprove}
          onInference={handleInference}
          onStop={handleStop}
          onToggleDiff={() => setShowDiff(!showDiff)}
          onSendInput={handleSendInput}
          onClearOutput={handleClearOutput}
          onLlmChange={setActiveLlm}
          onNextLlm={handleNextLlm}
          onAutoApproverConfigChange={handleAutoApproverConfigChange}
        />

        <FloatingButton onClick={() => setShowSidebar(!showSidebar)} />
      </div>
    </div>
  );
};

export default App;