import React, { useState } from 'react';
import './App.css';

interface DocumentResponse {
  text: string;
  type: string;
  score?: number;
}

function App() {
  const [file, setFile] = useState<File | null>(null);
  const [query, setQuery] = useState('');
  const [results, setResults] = useState<DocumentResponse[]>([]);
  const [loading, setLoading] = useState(false);
  const [message, setMessage] = useState('');
  const [sessionId, setSessionId] = useState<string | null>(null);
  const [currentAnswer, setCurrentAnswer] = useState<string>('');

  const handleFileChange = (e: React.ChangeEvent<HTMLInputElement>) => {
    if (e.target.files && e.target.files[0]) {
      setFile(e.target.files[0]);
    }
  };

  const handleUpload = async () => {
    if (!file) {
      setMessage('Please select a file first');
      return;
    }

    setLoading(true);
    const formData = new FormData();
    formData.append('file', file);

    try {
      const response = await fetch('http://localhost:9000/upload', {
        method: 'POST',
        body: formData,
      });

      if (!response.ok) {
        throw new Error('Upload failed');
      }

      const data = await response.json();
      setMessage(data.message);
      setSessionId(data.session_id);
      setResults([]);
      setCurrentAnswer('');
    } catch (error) {
      setMessage('Error uploading file');
    } finally {
      setLoading(false);
    }
  };

  const handleQuery = async () => {
    if (!query.trim()) {
      setMessage('Please enter a query');
      return;
    }

    if (!sessionId) {
      setMessage('Please upload a document first');
      return;
    }

    setLoading(true);
    setCurrentAnswer('');
    try {
      const response = await fetch(`http://localhost:9000/query/${sessionId}`, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
        },
        body: JSON.stringify({ text: query, k: 4 }),
      });

      if (!response.ok) {
        throw new Error('Query failed');
      }

      // Handle streaming response
      const reader = response.body?.getReader();
      const decoder = new TextDecoder();
      let buffer = '';

      if (!reader) {
        throw new Error('No reader available');
      }

      while (true) {
        const { done, value } = await reader.read();
        if (done) break;

        buffer += decoder.decode(value, { stream: true });
        const lines = buffer.split('\n');
        buffer = lines.pop() || '';

        for (const line of lines) {
          if (line.startsWith('data: ')) {
            const data = JSON.parse(line.slice(6));
            if (data.type === 'token') {
              setCurrentAnswer(prev => prev + data.text);
            } else if (data.type === 'context') {
              setResults(data.context);
            }
          }
        }
      }
    } catch (error) {
      setMessage('Error querying documents');
    } finally {
      setLoading(false);
    }
  };

  return (
    <div className="App">
      <header className="App-header">
        <h1>RAG System</h1>
      </header>
      
      <main className="App-main">
        <section className="upload-section">
          <h2>Upload PDF Document</h2>
          <input type="file" accept=".pdf" onChange={handleFileChange} />
          <button onClick={handleUpload} disabled={loading || !file}>
            {loading ? 'Uploading...' : 'Upload'}
          </button>
        </section>

        <section className="query-section">
          <h2>Query Documents</h2>
          <input
            type="text"
            value={query}
            onChange={(e) => setQuery(e.target.value)}
            placeholder="Enter your query..."
            disabled={!sessionId}
          />
          <button onClick={handleQuery} disabled={loading || !query.trim() || !sessionId}>
            {loading ? 'Searching...' : 'Search'}
          </button>
        </section>

        {message && <div className="message">{message}</div>}

        <section className="results-section">
          <h2>Results</h2>
          {currentAnswer && (
            <div className="llm-response">
              <h3>AI Response</h3>
              <p className="result-text">{currentAnswer}</p>
            </div>
          )}
          
          {results.length > 0 && (
            <div className="relevant-chunks">
              <h3>Relevant Document Chunks</h3>
              {results.map((result, index) => (
                <div key={index} className="result-item">
                  <p className="result-text">{result.text}</p>
                  {result.score !== undefined && (
                    <p className="result-score">Relevance Score: {result.score.toFixed(2)}</p>
                  )}
                </div>
              ))}
            </div>
          )}
        </section>
      </main>
    </div>
  );
}

export default App;