// components/ProtectedRoute.tsx
import { Navigate, useLocation } from 'react-router-dom'
import { useAuthStore } from '@/stores/authStore'

interface ProtectedRouteProps {
  children: React.ReactNode
}

export function ProtectedRoute({ children }: ProtectedRouteProps) {
  const user = useAuthStore(state => state.user)
  const loading = useAuthStore(state => state.loading)
  const location = useLocation()

  // Still checking auth (initial load from localStorage)
  if (loading) {
    return (
      <div className="flex items-center justify-center min-h-screen">
        <div className="animate-spin rounded-full h-12 w-12 border-b-2" />
      </div>
    )
  }

  // Not authenticated - redirect with return URL
  if (!user) {
    return <Navigate to="/login" state={{ from: location.pathname }} replace />
  }

  // Authenticated - render children
  return <>{children}</>
}

import { BrowserRouter, Routes, Route } from 'react-router-dom'
import { ProtectedRoute } from './components/ProtectedRoute'

function App() {
  return (
    <BrowserRouter>
      <Routes>
        {/* Public Routes */}
        <Route path="/" element={<HomePage />} />
        <Route path="/login" element={<LoginPage />} />
        <Route path="/register" element={<RegisterPage />} />

        {/* Protected Routes */}
        <Route
          path="/chat"
          element={
            <ProtectedRoute>
              <ChatPage />
            </ProtectedRoute>
          }
        />
        <Route
          path="/settings"
          element={
            <ProtectedRoute>
              <SettingsPage />
            </ProtectedRoute>
          }
        />
      </Routes>
    </BrowserRouter>
  )
}

import { useLocation, useNavigate } from 'react-router-dom'
import { useAuthStore } from '@/stores/authStore'

function LoginPage() {
  const location = useLocation()
  const navigate = useNavigate()
  const login = useAuthStore(state => state.login)

  const handleSubmit = async (e: React.FormEvent) => {
    e.preventDefault()
    await login(email, password)

    // Navigate to return URL or default to /chat
    const returnUrl = location.state?.from || '/chat'
    navigate(returnUrl, { replace: true })
  }

  return (
    <form onSubmit={handleSubmit}>
      <input type="email" required />
      <input type="password" required />
      <button type="submit">Login</button>
    </form>
  )
}

// pages/LoginPage.tsx
import { useState } from 'react'
import { useNavigate, useLocation } from 'react-router-dom'
import { useAuthStore } from '@/stores/authStore'

function LoginPage() {
  const [email, setEmail] = useState('')
  const [password, setPassword] = useState('')
  const navigate = useNavigate()
  const location = useLocation()

  const login = useAuthStore(state => state.login)
  const loading = useAuthStore(state => state.loading)
  const error = useAuthStore(state => state.error)

  const handleSubmit = async (e: React.FormEvent) => {
    e.preventDefault()
    await login(email, password)

    // Navigate to return URL or default to /chat
    const returnUrl = location.state?.from || '/chat'
    navigate(returnUrl, { replace: true })
  }

  return (
    <form onSubmit={handleSubmit}>
      <input
        type="email"
        value={email}
        onChange={(e) => setEmail(e.target.value)}
        placeholder="Email"
        required
      />
      <input
        type="password"
        value={password}
        onChange={(e) => setPassword(e.target.value)}
        placeholder="Password"
        required
      />
      {error && <p className="error">{error}</p>}
      <button type="submit" disabled={loading}>
        {loading ? 'Logging in...' : 'Login'}
      </button>
    </form>
  )
}

// pages/RegisterPage.tsx
import { useState } from 'react'
import { useNavigate } from 'react-router-dom'
import { useAuthStore } from '@/stores/authStore'

function RegisterPage() {
  const [email, setEmail] = useState('')
  const [password, setPassword] = useState('')
  const [name, setName] = useState('')
  const navigate = useNavigate()

  const register = useAuthStore(state => state.register)
  const loading = useAuthStore(state => state.loading)
  const error = useAuthStore(state => state.error)

  const handleSubmit = async (e: React.FormEvent) => {
    e.preventDefault()
    await register(email, password, name)

    // Navigate to chat after successful registration
    navigate('/chat', { replace: true })
  }

  return (
    <form onSubmit={handleSubmit}>
      <input
        type="text"
        value={name}
        onChange={(e) => setName(e.target.value)}
        placeholder="Name"
        required
      />
      <input
        type="email"
        value={email}
        onChange={(e) => setEmail(e.target.value)}
        placeholder="Email"
        required
      />
      <input
        type="password"
        value={password}
        onChange={(e) => setPassword(e.target.value)}
        placeholder="Password"
        required
      />
      {error && <p className="error">{error}</p>}
      <button type="submit" disabled={loading}>
        {loading ? 'Creating account...' : 'Register'}
      </button>
    </form>
  )
}

// api/axiosInstance.ts
import axios from 'axios'
import { useAuthStore } from '@/stores/authStore'

const api = axios.create({
  baseURL: '/api/v1'
})

// Request interceptor: Add access token to every request
api.interceptors.request.use(
  (config) => {
    const token = useAuthStore.getState().access_token
    if (token) {
      config.headers.Authorization = `Bearer ${token}`
    }
    return config
  }
)

// Response interceptor: Handle 401 and refresh token
api.interceptors.response.use(
  (response) => response,  // Success - pass through
  async (error) => {
    const originalRequest = error.config

    // 401 error and haven't retried yet
    if (error.response?.status === 401 && !originalRequest._retry) {
      originalRequest._retry = true

      try {
        const { refresh_token } = useAuthStore.getState()

        // Call refresh endpoint
        const response = await axios.post('/api/v1/auth/refresh', {
          refresh_token
        })

        const { access_token, refresh_token: new_refresh } = response.data

        // Update tokens in store (rotation)
        useAuthStore.setState({
          access_token,
          refresh_token: new_refresh
        })

        // Retry original request with new token
        originalRequest.headers.Authorization = `Bearer ${access_token}`
        return api(originalRequest)
      } catch (refreshError) {
        // Refresh failed - logout and redirect
        useAuthStore.getState().logout()
        window.location.href = '/login'
        return Promise.reject(refreshError)
      }
    }

    return Promise.reject(error)
  }
)

export default api

// Usage in components
import api from '@/api/axiosInstance'

async function fetchChats() {
  const response = await api.get('/chats')  // Auto-adds token, auto-refreshes if expired
  return response.data
}

// stores/chatStore.ts
import { create } from 'zustand'
import api from '@/api/axiosInstance'

interface Chat {
  session_id: string
  title: string
  last_activity: string
  message_count: number
}

interface Message {
  message_id: string
  role: 'user' | 'assistant'
  content: string
  created_at: string
}

interface ChatState {
  chats: Chat[]
  messages: Record<string, Message[]>  // session_id -> messages
  loading: boolean

  fetchChats: () => Promise<void>
  loadMessages: (sessionId: string) => Promise<void>
}

export const useChatStore = create<ChatState>((set, get) => ({
  chats: [],
  messages: {},
  loading: false,

  fetchChats: async () => {
    set({ loading: true })
    try {
      // Auto-includes Authorization header via interceptor
      const response = await api.get('/chats')
      set({ chats: response.data.chats, loading: false })
    } catch (error) {
      console.error('Failed to fetch chats:', error)
      set({ loading: false })
    }
  },

  loadMessages: async (sessionId) => {
    try {
      const response = await api.get(`/chat/${sessionId}/messages`)
      set(state => ({
        messages: {
          ...state.messages,
          [sessionId]: response.data.messages
        }
      }))
    } catch (error) {
      console.error('Failed to load messages:', error)
    }
  }
}))

// components/ChatLayout.tsx
import { useEffect } from 'react'
import { useAuthStore } from '@/stores/authStore'
import { useChatStore } from '@/stores/chatStore'

function ChatLayout() {
  const user = useAuthStore(state => state.user)
  const fetchChats = useChatStore(state => state.fetchChats)
  const chats = useChatStore(state => state.chats)

  // Load chats when user logs in or component mounts
  useEffect(() => {
    if (user) {
      fetchChats()  // Fetch user's chat list
    }
  }, [user, fetchChats])

  return (
    <div>
      <h2>Your Chats</h2>
      {chats.map(chat => (
        <ChatItem key={chat.session_id} chat={chat} />
      ))}
    </div>
  )
}

// When user selects a chat
function ChatPage({ sessionId }: { sessionId: string }) {
  const loadMessages = useChatStore(state => state.loadMessages)
  const messages = useChatStore(state => state.messages[sessionId])

  // Load messages for this chat when sessionId changes
  useEffect(() => {
    loadMessages(sessionId)
  }, [sessionId, loadMessages])

  return (
    <div>
      {messages?.map(msg => (
        <Message key={msg.message_id} {...msg} />
      ))}
    </div>
  )
}

# app/api/chat.py
from fastapi import APIRouter, Depends, HTTPException
from app.core.dependencies import get_current_user
from app.models.user import User
from app.models.session import UserSession
from app.models.message import Message

@router.get("/chat/{session_id}/messages")
async def get_session_messages(
    session_id: str,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_async_session)
):
    """Get all messages for a session (with ownership verification)"""

    # Verify session exists and user owns it
    result = await db.execute(
        select(UserSession).where(
            UserSession.session_id == session_id,
            UserSession.user_id == str(current_user.id)
        )
    )
    session = result.scalar_one_or_none()

    if not session:
        raise HTTPException(404, "Session not found or access denied")

    # Fetch all messages for this session, ordered by sequence
    result = await db.execute(
        select(Message)
        .where(Message.session_id == session_id)
        .order_by(Message.sequence)
    )
    messages = result.scalars().all()

    return {
        "session_id": session_id,
        "messages": [msg.to_dict() for msg in messages]
    }

// api/chat.ts
import api from './axiosInstance'

export interface Message {
  message_id: string
  session_id: string
  user_id: string
  sequence: number
  role: 'user' | 'assistant'
  content: string | ContentBlock[]
  tokens: number
  created_at: string
}

export async function getSessionMessages(
  sessionId: string
): Promise<Message[]> {
  try {
    const response = await api.get(`/chat/${sessionId}/messages`)
    return response.data.messages
  } catch (error) {
    // 404 is normal for new sessions (no messages yet)
    if (error.response?.status === 404) {
      return []
    }
    throw error
  }
}

// pages/ChatPage.tsx
import { useEffect, useState } from 'react'
import { useParams } from 'react-router-dom'
import { getSessionMessages } from '@/api/chat'

function ChatPage() {
  const { sessionId } = useParams<{ sessionId: string }>()
  const [messages, setMessages] = useState<Message[]>([])
  const [loading, setLoading] = useState(true)
  const [error, setError] = useState<string | null>(null)

  // Load conversation history when sessionId changes or on mount
  useEffect(() => {
    if (!sessionId) return

    const loadHistory = async () => {
      setLoading(true)
      setError(null)

      try {
        const history = await getSessionMessages(sessionId)
        setMessages(history)
      } catch (err) {
        console.error('Failed to load messages:', err)
        setError('Failed to load conversation history')
      } finally {
        setLoading(false)
      }
    }

    loadHistory()
  }, [sessionId])

  if (loading) {
    return <div className="loading">Loading conversation...</div>
  }

  if (error) {
    return <div className="error">{error}</div>
  }

  return (
    <div>
      {messages.length === 0 ? (
        <EmptyState />
      ) : (
        messages.map(msg => (
          <Message key={msg.message_id} {...msg} />
        ))
      )}
      <ChatInput onSubmit={handleSendMessage} />
    </div>
  )
}