src/App.js

import React, { useState, useEffect, useRef } from "react";
import {
  Send,
  Moon,
  Sun,
  User,
  Bot,
  AlertCircle,
  CheckCircle,
  Clock,
  Heart,
  Phone,
  Calendar,
  MapPin,
  FileText,
} from "lucide-react";
import { apiService } from "./services/apiService";

const MedicalChatUI = () => {
  // UI State
  const [isDarkMode, setIsDarkMode] = useState(false);
  const [currentMessage, setCurrentMessage] = useState("");
  const [messages, setMessages] = useState([]);
  const [isTyping, setIsTyping] = useState(false);
  const [isConnected, setIsConnected] = useState(false);
  const [connectionError, setConnectionError] = useState("");

  // Patient Data State
  const [patientRecord, setPatientRecord] = useState({
    personal_info: {
      full_name: "",
      dob: "",
      gender: "",
      phone: "",
      address: "",
    },
    medical_info: {
      current_symptoms: [],
      symptom_start_date: "",
      previous_medications: [],
      allergies: [],
      severity_level: "",
    },
  });

  // Session State
  const [sessionStatus, setSessionStatus] = useState("idle"); // idle | collecting | completed
  const [confidenceScores, setConfidenceScores] = useState({
    extraction: 0,
    validation: 0,
  });

  const messagesEndRef = useRef(null);

  // Auto-scroll to bottom
  const scrollToBottom = () => {
    messagesEndRef.current?.scrollIntoView({ behavior: "smooth" });
  };

  useEffect(() => {
    scrollToBottom();
  }, [messages]);

  // Initialize connection and start chat
  useEffect(() => {
    initializeChat();
  }, []);

  const initializeChat = async () => {
    try {
      // Health check
      await apiService.healthCheck();
      setIsConnected(true);
      setConnectionError("");

      // Start chat session
      const response = await apiService.startChat();

      // Add initial AI message
      setMessages([
        {
          id: Date.now(),
          text: response.message,
          sender: "ai",
          timestamp: new Date().toISOString(),
        },
      ]);

      setSessionStatus("collecting");
    } catch (error) {
      setIsConnected(false);
      setConnectionError(error.message);
      console.error("Failed to initialize chat:", error);
    }
  };

  const handleSendMessage = async () => {
    if (!currentMessage.trim() || isTyping) return;

    const userMessage = {
      id: Date.now(),
      text: currentMessage.trim(),
      sender: "user",
      timestamp: new Date().toISOString(),
    };

    // Add user message to chat
    setMessages((prev) => [...prev, userMessage]);
    setCurrentMessage("");
    setIsTyping(true);

    try {
      // Send to backend
      const response = await apiService.sendMessage(userMessage.text);

      // Update patient record
      setPatientRecord(response.patient_record);
      setConfidenceScores(response.confidence_scores);
      setSessionStatus(response.session_status);

      // Add AI response
      const aiMessage = {
        id: Date.now() + 1,
        text: response.ai_message,
        sender: "ai",
        timestamp: new Date().toISOString(),
        metadata: response.metadata,
      };

      setMessages((prev) => [...prev, aiMessage]);

      // If conversation completed, show completion message
      if (response.should_end) {
        setTimeout(() => {
          const completionMessage = {
            id: Date.now() + 2,
            text: "✅ Information collection completed! Click 'Generate Summary' to create a doctor report.",
            sender: "system",
            timestamp: new Date().toISOString(),
          };
          setMessages((prev) => [...prev, completionMessage]);
        }, 1000);
      }
    } catch (error) {
      console.error("Failed to send message:", error);

      // Add error message
      const errorMessage = {
        id: Date.now() + 1,
        text: `❌ Error: ${error.message}. Please try again.`,
        sender: "system",
        timestamp: new Date().toISOString(),
      };
      setMessages((prev) => [...prev, errorMessage]);
    } finally {
      setIsTyping(false);
    }
  };

  const handleGenerateSummary = async () => {
    try {
      setIsTyping(true);
      const summary = await apiService.generateSummary();

      const summaryMessage = {
        id: Date.now(),
        text: `📋 **Doctor Summary Generated**\n\n**Summary:** ${
          summary.doctor_summary
        }\n\n**Urgency Level:** ${
          summary.urgency_level
        }\n\n**Key Findings:**\n${summary.key_findings
          .map((f) => `• ${f}`)
          .join(
            "\n"
          )}\n\n**Recommended Questions:**\n${summary.recommended_questions
          .map((q) => `• ${q}`)
          .join("\n")}`,
        sender: "system",
        timestamp: new Date().toISOString(),
      };

      setMessages((prev) => [...prev, summaryMessage]);
    } catch (error) {
      console.error("Failed to generate summary:", error);
    } finally {
      setIsTyping(false);
    }
  };

  const handleResetChat = async () => {
    try {
      await apiService.resetSession();
      setMessages([]);
      setPatientRecord({
        personal_info: {
          full_name: "",
          dob: "",
          gender: "",
          phone: "",
          address: "",
        },
        medical_info: {
          current_symptoms: [],
          symptom_start_date: "",
          previous_medications: [],
          allergies: [],
          severity_level: "",
        },
      });
      setSessionStatus("idle");
      setConfidenceScores({ extraction: 0, validation: 0 });

      // Restart chat
      await initializeChat();
    } catch (error) {
      console.error("Failed to reset chat:", error);
    }
  };

  // Components
  const MessageBubble = ({ message }) => {
    const isUser = message.sender === "user";
    const isSystem = message.sender === "system";

    return (
      <div className={`flex ${isUser ? "justify-end" : "justify-start"} mb-4`}>
        <div className="flex items-start space-x-3 max-w-3xl">
          {!isUser && (
            <div
              className={`flex-shrink-0 w-8 h-8 rounded-full flex items-center justify-center ${
                isSystem
                  ? isDarkMode
                    ? "bg-yellow-900/50 text-yellow-300"
                    : "bg-yellow-100 text-yellow-700"
                  : isDarkMode
                  ? "bg-blue-600"
                  : "bg-blue-500"
              }`}
            >
              {isSystem ? (
                <AlertCircle className="w-4 h-4" />
              ) : (
                <Bot className="w-4 h-4 text-white" />
              )}
            </div>
          )}

          <div
            className={`px-4 py-3 rounded-2xl ${
              isUser
                ? isDarkMode
                  ? "bg-blue-600 text-white"
                  : "bg-blue-500 text-white"
                : isSystem
                ? isDarkMode
                  ? "bg-yellow-900/20 text-yellow-200 border border-yellow-700"
                  : "bg-yellow-50 text-yellow-800 border border-yellow-200"
                : isDarkMode
                ? "bg-gray-700 text-gray-200"
                : "bg-gray-100 text-gray-900"
            }`}
          >
            <div className="whitespace-pre-wrap">{message.text}</div>

            {message.metadata && (
              <div className="text-xs mt-2 opacity-70">
                Exchange #{message.metadata.exchange_count} | Confidence:{" "}
                {(message.metadata.confidence || 0).toFixed(2)}
              </div>
            )}
          </div>

          {isUser && (
            <div className="flex-shrink-0 w-8 h-8 rounded-full bg-gray-400 flex items-center justify-center">
              <User className="w-4 h-4 text-white" />
            </div>
          )}
        </div>
      </div>
    );
  };

  const TypingIndicator = () => (
    <div className="flex justify-start mb-4">
      <div className="flex items-center space-x-3 max-w-xs">
        <div
          className={`flex-shrink-0 w-8 h-8 rounded-full flex items-center justify-center ${
            isDarkMode ? "bg-blue-600" : "bg-blue-500"
          }`}
        >
          <Bot className="w-4 h-4 text-white" />
        </div>
        <div
          className={`px-4 py-3 rounded-2xl ${
            isDarkMode ? "bg-gray-700" : "bg-gray-100"
          }`}
        >
          <div className="flex space-x-1">
            <div className="w-2 h-2 bg-blue-500 rounded-full animate-bounce"></div>
            <div className="w-2 h-2 bg-blue-500 rounded-full animate-bounce [animation-delay:0.1s]"></div>
            <div className="w-2 h-2 bg-blue-500 rounded-full animate-bounce [animation-delay:0.2s]"></div>
          </div>
        </div>
      </div>
    </div>
  );

  const PatientPanel = () => (
    <div
      className={`w-96 ${
        isDarkMode ? "bg-gray-800 border-gray-700" : "bg-white border-gray-200"
      } border-l flex flex-col`}
    >
      {/* Header */}
      <div
        className={`p-4 border-b ${
          isDarkMode ? "border-gray-700" : "border-gray-200"
        }`}
      >
        <h2
          className={`text-lg font-semibold ${
            isDarkMode ? "text-gray-200" : "text-gray-900"
          } flex items-center`}
        >
          <FileText className="w-5 h-5 mr-2" />
          Patient Record
        </h2>
        <div className="flex items-center mt-2 space-x-4">
          <div
            className={`text-sm ${
              isDarkMode ? "text-gray-400" : "text-gray-600"
            }`}
          >
            Status: {sessionStatus}
          </div>
          <div className="text-xs bg-blue-100 text-blue-800 px-2 py-1 rounded-full">
            Confidence:{" "}
            {Math.round(
              ((confidenceScores.extraction + confidenceScores.validation) /
                2) *
                100
            )}
            %
          </div>
        </div>
      </div>

      {/* Content */}
      <div className="flex-1 overflow-y-auto p-4 space-y-6">
        {/* Personal Information */}
        <div>
          <h3
            className={`text-sm font-medium ${
              isDarkMode ? "text-gray-300" : "text-gray-700"
            } mb-3 flex items-center`}
          >
            <User className="w-4 h-4 mr-2" />
            Personal Information
          </h3>

          <div className="space-y-3">
            <div>
              <span
                className={`text-sm ${
                  isDarkMode ? "text-gray-400" : "text-gray-600"
                }`}
              >
                Name:
              </span>
              <p
                className={`text-sm mt-1 ${
                  isDarkMode ? "text-gray-200" : "text-gray-900"
                }`}
              >
                {patientRecord.personal_info.full_name || "Not provided"}
              </p>
            </div>

            <div>
              <span
                className={`text-sm ${
                  isDarkMode ? "text-gray-400" : "text-gray-600"
                }`}
              >
                Date of Birth:
              </span>
              <p
                className={`text-sm mt-1 ${
                  isDarkMode ? "text-gray-200" : "text-gray-900"
                } flex items-center`}
              >
                <Calendar className="w-3 h-3 mr-1" />
                {patientRecord.personal_info.dob || "Not provided"}
              </p>
            </div>

            <div>
              <span
                className={`text-sm ${
                  isDarkMode ? "text-gray-400" : "text-gray-600"
                }`}
              >
                Phone:
              </span>
              <p
                className={`text-sm mt-1 ${
                  isDarkMode ? "text-gray-200" : "text-gray-900"
                } flex items-center`}
              >
                <Phone className="w-3 h-3 mr-1" />
                {patientRecord.personal_info.phone || "Not provided"}
              </p>
            </div>

            <div>
              <span
                className={`text-sm ${
                  isDarkMode ? "text-gray-400" : "text-gray-600"
                }`}
              >
                Address:
              </span>
              <p
                className={`text-sm mt-1 ${
                  isDarkMode ? "text-gray-200" : "text-gray-900"
                } flex items-center`}
              >
                <MapPin className="w-3 h-3 mr-1" />
                {patientRecord.personal_info.address || "Not provided"}
              </p>
            </div>
          </div>
        </div>

        {/* Medical Information */}
        <div>
          <h3
            className={`text-sm font-medium ${
              isDarkMode ? "text-gray-300" : "text-gray-700"
            } mb-3 flex items-center`}
          >
            <Heart className="w-4 h-4 mr-2" />
            Medical Information
          </h3>

          <div className="space-y-3">
            <div>
              <span
                className={`text-sm ${
                  isDarkMode ? "text-gray-400" : "text-gray-600"
                }`}
              >
                Symptoms:
              </span>
              <p
                className={`text-sm mt-1 ${
                  isDarkMode ? "text-gray-200" : "text-gray-900"
                }`}
              >
                {patientRecord.medical_info.current_symptoms?.length > 0
                  ? patientRecord.medical_info.current_symptoms.join(", ")
                  : "Not provided"}
              </p>
            </div>

            <div>
              <span
                className={`text-sm ${
                  isDarkMode ? "text-gray-400" : "text-gray-600"
                }`}
              >
                Duration:
              </span>
              <p
                className={`text-sm mt-1 ${
                  isDarkMode ? "text-gray-200" : "text-gray-900"
                } flex items-center`}
              >
                <Clock className="w-3 h-3 mr-1" />
                {patientRecord.medical_info.symptom_start_date ||
                  "Not provided"}
              </p>
            </div>

            <div>
              <span
                className={`text-sm ${
                  isDarkMode ? "text-gray-400" : "text-gray-600"
                }`}
              >
                Severity:
              </span>
              <p
                className={`text-sm mt-1 ${
                  isDarkMode ? "text-gray-200" : "text-gray-900"
                }`}
              >
                {patientRecord.medical_info.severity_level || "Not assessed"}
              </p>
            </div>

            <div>
              <span
                className={`text-sm ${
                  isDarkMode ? "text-gray-400" : "text-gray-600"
                }`}
              >
                Medications:
              </span>
              <p
                className={`text-sm mt-1 ${
                  isDarkMode ? "text-gray-200" : "text-gray-900"
                }`}
              >
                {patientRecord.medical_info.previous_medications?.length > 0
                  ? patientRecord.medical_info.previous_medications.join(", ")
                  : "None reported"}
              </p>
            </div>

            <div>
              <span
                className={`text-sm ${
                  isDarkMode ? "text-gray-400" : "text-gray-600"
                }`}
              >
                Allergies:
              </span>
              <p
                className={`text-sm mt-1 ${
                  isDarkMode ? "text-gray-200" : "text-gray-900"
                }`}
              >
                {patientRecord.medical_info.allergies?.length > 0
                  ? patientRecord.medical_info.allergies.join(", ")
                  : "None reported"}
              </p>
            </div>
          </div>
        </div>
      </div>

      {/* Actions */}
      <div
        className={`p-4 border-t ${
          isDarkMode ? "border-gray-700" : "border-gray-200"
        } space-y-2`}
      >
        <button
          onClick={handleGenerateSummary}
          disabled={sessionStatus === "idle"}
          className={`w-full px-4 py-2 rounded-lg transition-colors ${
            sessionStatus === "idle"
              ? `${
                  isDarkMode
                    ? "bg-gray-700 text-gray-500 cursor-not-allowed"
                    : "bg-gray-200 text-gray-400 cursor-not-allowed"
                }`
              : `${
                  isDarkMode
                    ? "bg-blue-600 hover:bg-blue-700"
                    : "bg-blue-500 hover:bg-blue-600"
                } text-white`
          }`}
        >
          Generate Summary
        </button>
        <button
          onClick={handleResetChat}
          className={`w-full px-4 py-2 rounded-lg transition-colors ${
            isDarkMode
              ? "border border-gray-600 hover:bg-gray-700 text-gray-300"
              : "border border-gray-300 hover:bg-gray-50 text-gray-700"
          }`}
        >
          New Session
        </button>
      </div>
    </div>
  );

  // Connection error screen
  if (!isConnected && connectionError) {
    return (
      <div
        className={`h-screen flex items-center justify-center ${
          isDarkMode ? "bg-gray-900" : "bg-gray-50"
        }`}
      >
        <div
          className={`text-center p-8 rounded-lg ${
            isDarkMode ? "bg-gray-800" : "bg-white"
          } shadow-lg`}
        >
          <AlertCircle
            className={`w-16 h-16 mx-auto mb-4 ${
              isDarkMode ? "text-red-400" : "text-red-500"
            }`}
          />
          <h2
            className={`text-xl font-semibold mb-2 ${
              isDarkMode ? "text-gray-200" : "text-gray-900"
            }`}
          >
            Connection Failed
          </h2>
          <p
            className={`mb-4 ${isDarkMode ? "text-gray-400" : "text-gray-600"}`}
          >
            {connectionError}
          </p>
          <button
            onClick={initializeChat}
            className={`px-4 py-2 rounded-lg ${
              isDarkMode
                ? "bg-blue-600 hover:bg-blue-700"
                : "bg-blue-500 hover:bg-blue-600"
            } text-white transition-colors`}
          >
            Retry Connection
          </button>
        </div>
      </div>
    );
  }

  return (
    <div
      className={`h-screen flex ${isDarkMode ? "bg-gray-900" : "bg-gray-50"}`}
    >
      {/* Main Content */}
      <div className="flex-1 flex flex-col">
        {/* Header */}
        <header
          className={`${
            isDarkMode
              ? "bg-gray-800 border-gray-700"
              : "bg-white border-gray-200"
          } border-b px-6 py-4 flex items-center justify-between`}
        >
          <div className="flex items-center space-x-4">
            <h1
              className={`text-xl font-semibold ${
                isDarkMode ? "text-gray-200" : "text-gray-900"
              }`}
            >
              Medical Multi-Agent System
            </h1>
          </div>

          <div className="flex items-center space-x-2">
            <div
              className={`px-3 py-1 rounded-full text-sm flex items-center space-x-2 ${
                isConnected
                  ? isDarkMode
                    ? "bg-green-900/50 text-green-300 border border-green-700"
                    : "bg-green-100 text-green-700 border border-green-200"
                  : isDarkMode
                  ? "bg-red-900/50 text-red-300 border border-red-700"
                  : "bg-red-100 text-red-700 border border-red-200"
              }`}
            >
              {isConnected ? (
                <CheckCircle className="w-3 h-3" />
              ) : (
                <AlertCircle className="w-3 h-3" />
              )}
              <span>
                {isConnected ? "Connected to Qwen 3-4B" : "Disconnected"}
              </span>
            </div>
            <button
              onClick={() => setIsDarkMode(!isDarkMode)}
              className={`p-2 rounded-lg ${
                isDarkMode
                  ? "hover:bg-gray-700 text-gray-300"
                  : "hover:bg-gray-100 text-gray-600"
              } transition-colors`}
            >
              {isDarkMode ? (
                <Sun className="w-5 h-5" />
              ) : (
                <Moon className="w-5 h-5" />
              )}
            </button>
          </div>
        </header>

        <div className="flex-1 flex min-h-0 overflow-hidden">
          {/* Chat Area */}
          <div className="flex-1 flex flex-col">
            {/* Messages */}
            <div className="flex-1 overflow-y-auto p-6">
              {messages.length === 0 && (
                <div className="flex items-center justify-center h-full">
                  <div className="text-center">
                    <Bot
                      className={`w-16 h-16 mx-auto mb-4 ${
                        isDarkMode ? "text-gray-600" : "text-gray-400"
                      }`}
                    />
                    <p
                      className={`text-lg ${
                        isDarkMode ? "text-gray-400" : "text-gray-600"
                      }`}
                    >
                      Initializing Medical Multi-Agent System...
                    </p>
                  </div>
                </div>
              )}

              {messages.map((message) => (
                <MessageBubble key={message.id} message={message} />
              ))}
              {isTyping && <TypingIndicator />}
              <div ref={messagesEndRef} />
            </div>

            {/* Input Area */}
            <div
              className={`p-6 border-t ${
                isDarkMode ? "border-gray-700" : "border-gray-200"
              }`}
            >
              <div className="flex space-x-4">
                <input
                  type="text"
                  value={currentMessage}
                  onChange={(e) => setCurrentMessage(e.target.value)}
                  onKeyPress={(e) => e.key === "Enter" && handleSendMessage()}
                  placeholder="Describe your symptoms or provide requested information..."
                  disabled={!isConnected}
                  className={`flex-1 px-4 py-3 rounded-lg border ${
                    isDarkMode
                      ? "bg-gray-700 border-gray-600 text-gray-200 placeholder-gray-400"
                      : "bg-white border-gray-300 text-gray-900 placeholder-gray-500"
                  } focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-transparent disabled:opacity-50`}
                />
                <button
                  onClick={handleSendMessage}
                  disabled={!currentMessage.trim() || isTyping || !isConnected}
                  className={`px-6 py-3 rounded-lg transition-colors flex items-center space-x-2 ${
                    !currentMessage.trim() || isTyping || !isConnected
                      ? `${
                          isDarkMode
                            ? "bg-gray-700 text-gray-500"
                            : "bg-gray-200 text-gray-400"
                        } cursor-not-allowed`
                      : `${
                          isDarkMode
                            ? "bg-blue-600 hover:bg-blue-700"
                            : "bg-blue-500 hover:bg-blue-600"
                        } text-white`
                  }`}
                >
                  <Send className="w-4 h-4" />
                  <span>{isTyping ? "Sending..." : "Send"}</span>
                </button>
              </div>
            </div>
          </div>

          {/* Patient Panel */}
          <PatientPanel />
        </div>
      </div>
    </div>
  );
};

export default MedicalChatUI;