import { useState, useEffect, createContext, useContext, useMemo, useCallback } 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"
import { ThemeProvider, useTheme } from "./components/theme-provider"

// --- Auth Context ---
interface AuthContextType {
  isAuthenticated: boolean
  setIsAuthenticated: (val: boolean) => void
  user: UserResponse | null
  setUser: (user: UserResponse | null) => void
}

const AuthContext = createContext<AuthContextType | undefined>(undefined)

function AuthProvider({ children }: { children: React.ReactNode }) {
  const [isAuthenticated, setIsAuthenticated] = useState(false)
  const [user, setUser] = useState<UserResponse | null>(null)

  const value = useMemo(() => ({
    isAuthenticated,
    setIsAuthenticated,
    user,
    setUser
  }), [isAuthenticated, user])

  return (
    <AuthContext.Provider value={value}>
      {children}
    </AuthContext.Provider>
  )
}

function useAuth() {
  const context = useContext(AuthContext)
  if (context === undefined) {
    throw new Error("useAuth must be used within an AuthProvider")
  }
  return context
}

// --- App Content ---
function AppContent() {
  const { setThemeLocal } = useTheme()
  const { isAuthenticated, setIsAuthenticated, user, setUser } = useAuth()
  
  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(() => {
    registerUnauthorizedCallback(() => {
      setIsAuthenticated(false)
      setUser(null)
      setConversations([])
      setSelectedThreadId(null)
      setThreadMessages({})
      setThemeLocal("light")
    })

    const initAuth = async () => {
      try {
        const userData = await AuthService.getMe()
        setUser(userData)
        setIsAuthenticated(true)
        if (userData.theme_preference) {
          setThemeLocal(userData.theme_preference)
        }
        loadHistory()
      } catch (err: unknown) {
        console.log("No active session found", err)
        setIsAuthenticated(false)
      } finally {
        setIsLoading(false)
      }
    }

    initAuth()
  }, [setIsAuthenticated, setThemeLocal, setUser])

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

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

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

  const handleSelectConversation = useCallback(async (id: string) => {
    setSelectedThreadId(id)
    try {
      const msgs = await ChatService.getMessages(id)
      setThreadMessages(prev => ({ ...prev, [id]: msgs }))
    } catch (err) {
      console.error("Failed to fetch messages:", err)
    }
  }, [])

  const handleCreateConversation = useCallback(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 = useCallback(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 = useCallback(async (id: string) => {
    try {
      await ChatService.deleteConversation(id)
      setConversations(prev => prev.filter(c => c.id !== id))
      setThreadMessages(prev => {
        const next = { ...prev }
        delete next[id]
        return next
      })
      if (selectedThreadId === id) {
        setSelectedThreadId(null)
      }
    } catch (err) {
      console.error("Failed to delete conversation:", err)
    }
  }, [selectedThreadId])

  const handleMessagesFinal = useCallback((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} 
                      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>
  )
}

function ThemeWrapper({ children }: { children: React.ReactNode }) {
  const { isAuthenticated } = useAuth()
  return (
    <ThemeProvider isAuthenticated={isAuthenticated}>
      {children}
    </ThemeProvider>
  )
}

function App() {
  return (
    <AuthProvider>
      <ThemeWrapper>
        <AppContent />
      </ThemeWrapper>
    </AuthProvider>
  )
}

export default App