frontend/src/App.tsx

import React, { useState, useCallback } from 'react';
import { useDropzone } from 'react-dropzone';
import axios from 'axios';
import { ChatBubbleLeftIcon, PhotoIcon, ArrowUpTrayIcon } from '@heroicons/react/24/outline';

interface Message {
  type: 'user' | 'assistant';
  content: string;
  imageUrl?: string;
}

function App() {
  const [messages, setMessages] = useState<Message[]>([]);
  const [prompt, setPrompt] = useState('');
  const [isLoading, setIsLoading] = useState(false);
  const [selectedImage, setSelectedImage] = useState<File | null>(null);
  const [previewUrl, setPreviewUrl] = useState<string | null>(null);

  const onDrop = useCallback((acceptedFiles: File[]) => {
    const file = acceptedFiles[0];
    if (file) {
      setSelectedImage(file);
      const url = URL.createObjectURL(file);
      setPreviewUrl(url);
    }
  }, []);

  const { getRootProps, getInputProps, isDragActive } = useDropzone({
    onDrop,
    accept: {
      'image/*': ['.png', '.jpg', '.jpeg', '.gif']
    },
    maxFiles: 1
  });

  const handleSubmit = async (e: React.FormEvent) => {
    e.preventDefault();
    if (!selectedImage || !prompt.trim()) return;

    setIsLoading(true);
    const formData = new FormData();
    formData.append('file', selectedImage);
    formData.append('prompt', prompt);

    // Add user message
    setMessages(prev => [...prev, {
      type: 'user',
      content: prompt,
      imageUrl: previewUrl || undefined
    }]);

    try {
      const response = await axios.post('http://localhost:8000/api/chat', formData, {
        headers: {
          'Content-Type': 'multipart/form-data',
        },
      });

      // Add assistant message
      setMessages(prev => [...prev, {
        type: 'assistant',
        content: response.data.response
      }]);

      // Clear input
      setPrompt('');
      setSelectedImage(null);
      setPreviewUrl(null);
    } catch (error) {
      console.error('Error:', error);
      // Add error message
      setMessages(prev => [...prev, {
        type: 'assistant',
        content: 'Sorry, there was an error processing your request.'
      }]);
    } finally {
      setIsLoading(false);
    }
  };

  return (
    <div className="min-h-screen bg-gray-100">
      <div className="max-w-4xl mx-auto p-4">
        <header className="text-center py-8">
          <h1 className="text-4xl font-bold text-primary-600">LLaVA Chat</h1>
          <p className="text-gray-600 mt-2">Upload an image and chat with LLaVA about it</p>
        </header>

        <div className="bg-white rounded-lg shadow-lg p-4 mb-4">
          <div className="space-y-4">
            {messages.map((message, index) => (
              <div
                key={index}
                className={`flex ${message.type === 'user' ? 'justify-end' : 'justify-start'}`}
              >
                <div
                  className={`max-w-[80%] rounded-lg p-4 ${
                    message.type === 'user'
                      ? 'bg-primary-600 text-white'
                      : 'bg-gray-100 text-gray-800'
                  }`}
                >
                  {message.imageUrl && (
                    <img
                      src={message.imageUrl}
                      alt="Uploaded"
                      className="w-48 h-48 object-cover rounded-lg mb-2"
                    />
                  )}
                  <p className="whitespace-pre-wrap">{message.content}</p>
                </div>
              </div>
            ))}
          </div>
        </div>

        <form onSubmit={handleSubmit} className="bg-white rounded-lg shadow-lg p-4">
          {!selectedImage ? (
            <div
              {...getRootProps()}
              className={`border-2 border-dashed rounded-lg p-8 text-center cursor-pointer transition-colors
                ${isDragActive ? 'border-primary-500 bg-primary-50' : 'border-gray-300 hover:border-primary-500'}`}
            >
              <input {...getInputProps()} />
              <PhotoIcon className="mx-auto h-12 w-12 text-gray-400" />
              <p className="mt-2 text-sm text-gray-600">
                Drag and drop an image here, or click to select
              </p>
            </div>
          ) : (
            <div className="relative">
              <img
                src={previewUrl || ''}
                alt="Preview"
                className="w-full h-48 object-cover rounded-lg"
              />
              <button
                type="button"
                onClick={() => {
                  setSelectedImage(null);
                  setPreviewUrl(null);
                }}
                className="absolute top-2 right-2 bg-red-500 text-white p-1 rounded-full hover:bg-red-600"
              >
                ×
              </button>
            </div>
          )}

          <div className="mt-4 flex space-x-4">
            <input
              type="text"
              value={prompt}
              onChange={(e) => setPrompt(e.target.value)}
              placeholder="Ask about the image..."
              className="input-primary flex-1"
              disabled={!selectedImage || isLoading}
            />
            <button
              type="submit"
              disabled={!selectedImage || !prompt.trim() || isLoading}
              className="btn-primary disabled:opacity-50 disabled:cursor-not-allowed"
            >
              {isLoading ? (
                <div className="w-6 h-6 border-2 border-white border-t-transparent rounded-full animate-spin" />
              ) : (
                <ArrowUpTrayIcon className="h-6 w-6" />
              )}
            </button>
          </div>
        </form>
      </div>
    </div>
  );
}

export default App;