ea-chatbot-lg/frontend/src/App.tsx

import { useState, useEffect } from "react"
import { Routes, Route } from "react-router-dom"
import { MainLayout } from "./components/layout/MainLayout"
import { LoginForm } from "./components/auth/LoginForm"
import { RegisterForm } from "./components/auth/RegisterForm"
import { ChatInterface } from "./components/chat/ChatInterface"
import { AuthService, type UserResponse } from "./services/auth"
import { ChatService, type MessageResponse } from "./services/chat"
import { type Conversation } from "./components/layout/HistorySidebar"
import { registerUnauthorizedCallback } from "./services/api"
import { Button } from "./components/ui/button"

function App() {
  const [isAuthenticated, setIsAuthenticated] = useState(false)
  const [user, setUser] = useState<UserResponse | null>(null)
  const [authMode, setAuthMode] = useState<"login" | "register">("login")
  const [isLoading, setIsLoading] = useState(true)
  const [selectedThreadId, setSelectedThreadId] = useState<string | null>(null)
  const [conversations, setConversations] = useState<Conversation[]>([])
  const [threadMessages, setThreadMessages] = useState<Record<string, MessageResponse[]>>({})

  useEffect(() => {
    // Register callback to handle session expiration from anywhere in the app
    registerUnauthorizedCallback(() => {
      setIsAuthenticated(false)
      setUser(null)
      setConversations([])
      setSelectedThreadId(null)
      setThreadMessages({})
    })

    const initAuth = async () => {
      try {
        const userData = await AuthService.getMe()
        setUser(userData)
        setIsAuthenticated(true)
        // Load history after successful auth
        loadHistory()
      } catch (err: unknown) {
        console.log("No active session found", err)
        setIsAuthenticated(false)
      } finally {
        setIsLoading(false)
      }
    }

    initAuth()
  }, [])

  const loadHistory = async () => {
    try {
      const history = await ChatService.listConversations()
      setConversations(history)
    } catch (err) {
      console.error("Failed to load conversation history:", err)
    }
  }

  const handleAuthSuccess = async () => {
    try {
      const userData = await AuthService.getMe()
      setUser(userData)
      setIsAuthenticated(true)
      loadHistory()
    } catch (err: unknown) {
      console.error("Failed to fetch user profile after login:", err)
    }
  }

  const handleLogout = async () => {
    try {
      await AuthService.logout()
    } catch (err: unknown) {
      console.error("Logout failed:", err)
    } finally {
      setIsAuthenticated(false)
      setUser(null)
      setSelectedThreadId(null)
      setConversations([])
      setThreadMessages({})
    }
  }

  const handleSelectConversation = async (id: string) => {
    setSelectedThreadId(id)
    // Always fetch messages to avoid stale cache issues when switching back
    // or if the session was updated from elsewhere
    try {
      const msgs = await ChatService.getMessages(id)
      setThreadMessages(prev => ({ ...prev, [id]: msgs }))
    } catch (err) {
      console.error("Failed to fetch messages:", err)
    }
  }

  const handleCreateConversation = async () => {
    try {
      const newConv = await ChatService.createConversation()
      setConversations(prev => [newConv, ...prev])
      setSelectedThreadId(newConv.id)
      setThreadMessages(prev => ({ ...prev, [newConv.id]: [] }))
    } catch (err) {
      console.error("Failed to create conversation:", err)
    }
  }

  const handleRenameConversation = async (id: string, name: string) => {
    try {
      const updated = await ChatService.renameConversation(id, name)
      setConversations(prev => prev.map(c => c.id === id ? updated : c))
    } catch (err) {
      console.error("Failed to rename conversation:", err)
    }
  }

  const handleDeleteConversation = async (id: string) => {
    try {
      await ChatService.deleteConversation(id)
      setConversations(prev => prev.filter(c => c.id !== id))
      // Also clear from cache
      setThreadMessages(prev => {
        const next = { ...prev }
        delete next[id]
        return next
      })
      if (selectedThreadId === id) {
        setSelectedThreadId(null)
      }
    } catch (err) {
      console.error("Failed to delete conversation:", err)
    }
  }

  const handleMessagesFinal = (id: string, messages: MessageResponse[]) => {
    setThreadMessages(prev => ({ ...prev, [id]: messages }))
  }

  const queryParams = new URLSearchParams(window.location.search)
  const externalError = queryParams.get("error")

  if (isLoading) {
    return (
      <div className="min-h-screen flex items-center justify-center bg-background">
        <div className="animate-spin rounded-full h-8 w-8 border-b-2 border-primary"></div>
      </div>
    )
  }

  return (
    <Routes>
      <Route
        path="*"
        element={
          !isAuthenticated ? (
            <div className="min-h-screen flex items-center justify-center bg-background p-4">
              {authMode === "login" ? (
                <LoginForm
                  onSuccess={handleAuthSuccess}
                  onToggleMode={() => setAuthMode("register")}
                  externalError={externalError === "oidc_failed" ? "SSO authentication failed. Please try again." : null}
                />
              ) : (
                <RegisterForm
                  onSuccess={handleAuthSuccess}
                  onToggleMode={() => setAuthMode("login")}
                />
              )}
            </div>
          ) : (
            <MainLayout
              conversations={conversations}
              selectedId={selectedThreadId}
              onSelect={handleSelectConversation}
              onCreate={handleCreateConversation}
              onRename={handleRenameConversation}
              onDelete={handleDeleteConversation}
            >
              <div className="flex flex-col h-full gap-4">
                <div className="flex justify-between items-center shrink-0">
                  <div>
                    <h1 className="text-xl font-bold">
                      Welcome, {user?.display_name || user?.email || "User"}!
                    </h1>
                  </div>
                  <Button
                    variant="link"
                    onClick={handleLogout}
                    className="text-sm text-muted-foreground hover:text-primary h-auto p-0 font-normal underline"
                  >
                    Logout
                  </Button>
                </div>
                
                <div className="flex-1 min-h-0">
                  {selectedThreadId ? (
                    <ChatInterface 
                      key={selectedThreadId} // Force remount on thread change
                      threadId={selectedThreadId} 
                      initialMessages={threadMessages[selectedThreadId] || []}
                      onMessagesFinal={(msgs) => handleMessagesFinal(selectedThreadId, msgs)}
                    />
                  ) : (
                    <div className="flex flex-col items-center justify-center h-full text-center space-y-4 bg-muted/30 rounded-xl border border-dashed p-12">
                      <div className="p-4 bg-background rounded-full shadow-sm">
                        <svg
                          className="w-12 h-12 text-primary"
                          fill="none"
                          stroke="currentColor"
                          viewBox="0 0 24 24"
                          xmlns="http://www.w3.org/2000/svg"
                        >
                          <path
                            strokeLinecap="round"
                            strokeLinejoin="round"
                            strokeWidth={2}
                            d="M8 10h.01M12 10h.01M16 10h.01M9 16H5a2 2 0 01-2-2V6a2 2 0 012-2h14a2 2 0 012 2v8a2 2 0 01-2 2h-5l-5 5v-5z"
                          />
                        </svg>
                      </div>
                      <div className="max-w-xs space-y-2">
                        <h2 className="text-lg font-semibold">Ready to analyze election data?</h2>
                        <p className="text-sm text-muted-foreground">
                          Create a new conversation in the sidebar to start asking questions.
                        </p>
                        <Button 
                          onClick={handleCreateConversation}
                          className="mt-4"
                        >
                          Start New Chat
                        </Button>
                      </div>
                    </div>
                  )}
                </div>
              </div>
            </MainLayout>
          )
        }
      />
    </Routes>
  )
}

export default App