import { useState, useEffect, useRef, useCallback } from 'react'
import { useLocation, useNavigate, useParams } from 'react-router-dom'
import { Sparkles, Send, Loader2, MessageSquare, Paperclip, Terminal, X, RotateCcw, ImagePlus, ListChecks, FileText, CheckCircle2, ArrowUpRight, MoreHorizontal, Pause } from 'lucide-react'
import { cn } from '@/lib/utils'
import { uploadsApi } from '@/api/uploads'
import type { PendingUpload } from '@/types/upload'
import type { ForkMetadata, ActionItem, QuestionItem } from '@/types/ai-session'
import { PageMeta } from '@/components/common/PageMeta'
import { aiSessionsApi } from '@/api/aiSessions'
import { useBranching } from '@/hooks/useBranching'
import { analytics } from '@/lib/analytics'
import { toast } from '@/lib/toast'
import { ChatSidebar, ChatSidebarCollapsedBar } from '@/components/assistant/ChatSidebar'
import { ChatMessage } from '@/components/assistant/ChatMessage'
import { TaskLane, clearTaskState } from '@/components/assistant/TaskLane'
import { ConcludeSessionModal } from '@/components/assistant/ConcludeSessionModal'
import { StatusUpdateModal } from '@/components/flowpilot/StatusUpdateModal'
import type { ChatListItem, ConclusionOutcome } from '@/types/assistant-chat'
import type { SuggestedFlow } from '@/types/copilot'

interface MessageWithMeta {
  role: 'user' | 'assistant'
  content: string
  suggestedFlows?: SuggestedFlow[]
  fork?: ForkMetadata | null
  actions?: ActionItem[] | null
  questions?: QuestionItem[] | null
  imageUrls?: string[]
}

export default function AssistantChatPage() {
  const location = useLocation()
  const navigate = useNavigate()
  const { sessionId: urlSessionId } = useParams<{ sessionId?: string }>()
  const [chats, setChats] = useState<ChatListItem[]>([])
  const [activeChatId, setActiveChatId] = useState<string | null>(() => {
    if (urlSessionId) return urlSessionId
    try { return sessionStorage.getItem('rf-active-chat-id') } catch { return null }
  })
  const [messages, setMessages] = useState<MessageWithMeta[]>([])
  const [input, setInput] = useState('')
  const [loading, setLoading] = useState(false)
  const [showConclude, setShowConclude] = useState(false)
  const [showStatusUpdate, setShowStatusUpdate] = useState(false)
  const branching = useBranching()
  const [mobileSidebarOpen, setMobileSidebarOpen] = useState(false)
  const [showLogs, setShowLogs] = useState(false)
  const [logContent, setLogContent] = useState('')
  const [pendingUploads, setPendingUploads] = useState<PendingUpload[]>([])
  const [isDragOver, setIsDragOver] = useState(false)
  const [activeQuestions, setActiveQuestions] = useState<QuestionItem[]>(() => {
    try {
      const saved = sessionStorage.getItem('rf-tasklane-meta')
      if (saved) { const d = JSON.parse(saved); if (d.chatId === activeChatId) return d.questions || [] }
    } catch { /* ignore */ }
    return []
  })
  const [activeActions, setActiveActions] = useState<ActionItem[]>(() => {
    try {
      const saved = sessionStorage.getItem('rf-tasklane-meta')
      if (saved) { const d = JSON.parse(saved); if (d.chatId === activeChatId) return d.actions || [] }
    } catch { /* ignore */ }
    return []
  })
  const [showTaskLane, setShowTaskLane] = useState(() => {
    try {
      const saved = sessionStorage.getItem('rf-tasklane-meta')
      if (saved) { const d = JSON.parse(saved); return d.show === true && d.chatId === activeChatId }
    } catch { /* ignore */ }
    return false
  })
  const [sidebarCollapsed, setSidebarCollapsed] = useState(() =>
    localStorage.getItem('rf-chat-sidebar-collapsed') === 'true'
  )
  const [activeSessionStatus, setActiveSessionStatus] = useState<string | null>(null)
  const [activePsaTicketId, setActivePsaTicketId] = useState<string | null>(null)
  const [showOverflow, setShowOverflow] = useState(false)
  const toggleSidebarCollapse = () => {
    const next = !sidebarCollapsed
    setSidebarCollapsed(next)
    localStorage.setItem('rf-chat-sidebar-collapsed', String(next))
  }
  const messagesEndRef = useRef<HTMLDivElement>(null)
  const inputRef = useRef<HTMLTextAreaElement>(null)
  const fileInputRef = useRef<HTMLInputElement>(null)
  const dragCounterRef = useRef(0)
  const prefillHandledRef = useRef(false)
  // Tracks the most recently requested active chat ID so in-flight selectChat
  // calls that complete after the user switches chats don't clobber new state.
  const currentChatRef = useRef<string | null>(activeChatId)

  // Persist active chat ID to sessionStorage
  useEffect(() => {
    try {
      if (activeChatId) sessionStorage.setItem('rf-active-chat-id', activeChatId)
      else sessionStorage.removeItem('rf-active-chat-id')
    } catch { /* ignore */ }
  }, [activeChatId])

  // Load chat list from ai_sessions
  useEffect(() => {
    loadChats()
  }, [])

  // If URL has a session ID, load it
  useEffect(() => {
    if (urlSessionId && urlSessionId !== activeChatId) {
      selectChat(urlSessionId)
    }
  }, [urlSessionId]) // eslint-disable-line react-hooks/exhaustive-deps

  // Restore session from sessionStorage on mount (when URL has no session ID)
  useEffect(() => {
    if (!urlSessionId && activeChatId) {
      selectChat(activeChatId)
    }
  }, []) // eslint-disable-line react-hooks/exhaustive-deps

  // Handle prefill from command palette / dashboard handoff
  useEffect(() => {
    const state = location.state as { prefill?: string; uploadIds?: string[] } | null
    const prefill = state?.prefill
    const uploadIds = state?.uploadIds
    if (!prefill || prefillHandledRef.current) return
    prefillHandledRef.current = true

    navigate(location.pathname, { replace: true, state: {} })

    const sendPrefill = async () => {
      // Clear stale task lane from previous session
      setShowTaskLane(false)
      setActiveQuestions([])
      setActiveActions([])
      setActiveSessionStatus('active')
      setActivePsaTicketId(null)

      try {
        const session = await aiSessionsApi.createChatSession({
          intake_type: 'free_text',
          intake_content: { text: prefill },
        })
        const chatItem: ChatListItem = {
          id: session.session_id,
          title: session.title,
          message_count: 0,
          pinned: false,
          created_at: new Date().toISOString(),
          updated_at: new Date().toISOString(),
        }
        setChats(prev => [chatItem, ...prev])
        setActiveChatId(session.session_id)
        setMessages([{ role: 'user', content: prefill }])
        setLoading(true)

        const response = await aiSessionsApi.sendChatMessage(session.session_id, {
          message: prefill,
          upload_ids: uploadIds?.length ? uploadIds : undefined,
        })
        setMessages(prev => [
          ...prev,
          { role: 'assistant', content: response.content, suggestedFlows: response.suggested_flows, fork: response.fork, actions: response.actions, questions: response.questions },
        ])
        setChats(prev =>
          prev.map(c =>
            c.id === session.session_id
              ? { ...c, message_count: 2, title: prefill.slice(0, 100), updated_at: new Date().toISOString() }
              : c
          )
        )
        // Show task lane if AI sent questions or actions
        if (response.fork && session.session_id) {
          branching.loadBranches(session.session_id)
        }
        const hasQuestions = response.questions && response.questions.length > 0
        const hasActions = response.actions && response.actions.length > 0
        if (hasQuestions || hasActions) {
          setActiveQuestions(response.questions || [])
          setActiveActions(response.actions || [])
          setShowTaskLane(true)
        }
      } catch {
        toast.error('Failed to start AI conversation')
      } finally {
        setLoading(false)
      }
    }

    sendPrefill()
  // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [])

  // Persist task lane metadata to sessionStorage
  useEffect(() => {
    try {
      sessionStorage.setItem('rf-tasklane-meta', JSON.stringify({
        show: showTaskLane,
        chatId: activeChatId,
        questions: activeQuestions,
        actions: activeActions,
      }))
    } catch { /* ignore */ }
  }, [showTaskLane, activeChatId, activeQuestions, activeActions])

  // Auto-scroll
  useEffect(() => {
    messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' })
  }, [messages])

  const loadChats = async () => {
    try {
      const sessions = await aiSessionsApi.listSessions({ session_type: 'chat', limit: 100 })
      setChats(sessions.map(s => ({
        id: s.id,
        title: s.title || s.problem_summary || 'New Chat',
        message_count: s.step_count,
        pinned: false,
        created_at: s.created_at,
        updated_at: s.created_at,
      })))
    } catch {
      // silently handle
    }
  }

  const selectChat = useCallback(async (chatId: string) => {
    currentChatRef.current = chatId
    setActiveChatId(chatId)
    // Clear TaskLane when switching chats — will restore from backend if available
    setShowTaskLane(false)
    setActiveQuestions([])
    setActiveActions([])
    setActiveSessionStatus(null)
    setActivePsaTicketId(null)
    try {
      const detail = await aiSessionsApi.getSession(chatId)
      // Guard: if the user switched to a different chat while this API call was
      // in flight (e.g. clicked "New Chat"), discard stale results so we don't
      // clobber the new session's task lane state.
      if (currentChatRef.current !== chatId) return
      setActiveSessionStatus(detail.status)
      setActivePsaTicketId(detail.psa_ticket_id)
      setMessages(
        (detail.conversation_messages || []).map(m => ({
          role: m.role as 'user' | 'assistant',
          content: m.content,
        }))
      )
      // Restore task lane from persisted state
      if (detail.pending_task_lane) {
        const q = detail.pending_task_lane.questions || []
        const a = detail.pending_task_lane.actions || []
        if (q.length > 0 || a.length > 0) {
          // Pre-load user's saved responses into sessionStorage BEFORE setting props
          // so TaskLane can restore them on mount/prop-change
          const responses = detail.pending_task_lane.responses
          if (responses && responses.length > 0) {
            try {
              sessionStorage.setItem(`rf-tasklane-state:${chatId}`, JSON.stringify(responses))
            } catch { /* ignore */ }
          }
          setActiveQuestions(q)
          setActiveActions(a)
          setShowTaskLane(true)
        }
      }
    } catch {
      setMessages([])
    }
  }, [])

  const handleNewChat = async () => {
    // Invalidate currentChatRef BEFORE the await so any in-flight handleSend/handleTaskSubmit
    // for the previous session sees a mismatch and bails — prevents stale task lane appearing
    // in the new empty session (same pattern as selectChat, which sets ref before its await).
    currentChatRef.current = null
    // Clear stale state immediately — don't wait for API to return
    setShowTaskLane(false)
    setActiveQuestions([])
    setActiveActions([])
    setMessages([])
    setActiveSessionStatus('active')
    setActivePsaTicketId(null)
    try {
      const session = await aiSessionsApi.createChatSession({
        intake_type: 'free_text',
        intake_content: { text: '' },
      })
      const chatItem: ChatListItem = {
        id: session.session_id,
        title: session.title,
        message_count: 0,
        pinned: false,
        created_at: new Date().toISOString(),
        updated_at: new Date().toISOString(),
      }
      currentChatRef.current = session.session_id
      setChats(prev => [chatItem, ...prev])
      setActiveChatId(session.session_id)
    } catch {
      toast.error('Failed to create chat')
    }
  }

  const handleDeleteChat = async (chatId: string) => {
    try {
      await aiSessionsApi.deleteSession(chatId)
      setChats(prev => prev.filter(c => c.id !== chatId))
      if (activeChatId === chatId) {
        setActiveChatId(null)
        setMessages([])
      }
    } catch {
      toast.error('Failed to delete chat')
    }
  }

  const handleTogglePin = async () => {
    // Pin/unpin not yet supported on unified sessions — no-op for now
    toast.info('Pin feature coming soon')
  }

  const handleSend = async () => {
    if (!input.trim() || !activeChatId || loading) return

    const userMessage = input.trim()
    const completedUploads = pendingUploads.filter((u) => u.status === 'done' && u.result?.id)
    const completedUploadIds = completedUploads.map((u) => u.result!.id)
    const imageUrls = completedUploads
      .filter((u) => u.preview)
      .map((u) => u.preview)
    setInput('')
    setPendingUploads([])
    setMessages(prev => [...prev, { role: 'user', content: userMessage, imageUrls: imageUrls.length > 0 ? imageUrls : undefined }])
    setLoading(true)

    const sentForChatId = activeChatId
    try {
      const response = await aiSessionsApi.sendChatMessage(activeChatId, {
        message: userMessage,
        upload_ids: completedUploadIds.length > 0 ? completedUploadIds : undefined,
      })
      // Guard: discard if user switched to a different chat while this was in flight
      if (currentChatRef.current !== sentForChatId) return
      analytics.aiFeatureUsed({ feature: 'assistant_chat' })
      setMessages(prev => [
        ...prev,
        { role: 'assistant', content: response.content, suggestedFlows: response.suggested_flows, fork: response.fork, actions: response.actions, questions: response.questions },
      ])
      setChats(prev =>
        prev.map(c =>
          c.id === sentForChatId
            ? { ...c, message_count: c.message_count + 2, title: c.message_count === 0 ? userMessage.slice(0, 100) : c.title, updated_at: new Date().toISOString() }
            : c
        )
      )
      // Load branches if fork was created
      if (response.fork && sentForChatId) {
        branching.loadBranches(sentForChatId)
      }
      // Show task lane if AI sent questions or actions
      const hasQuestions = response.questions && response.questions.length > 0
      const hasActions = response.actions && response.actions.length > 0
      if (hasQuestions || hasActions) {
        clearTaskState(sentForChatId)
        setActiveQuestions(response.questions || [])
        setActiveActions(response.actions || [])
        setShowTaskLane(true)
      }
    } catch (err: unknown) {
      console.error('[AssistantChat] sendChatMessage failed:', err)
      const status = (err as { response?: { status?: number } })?.response?.status
      let errorMsg: string
      if (status === 429) {
        errorMsg = "You're sending messages too quickly. Wait a moment and try again."
      } else if (status === 502 || status === 503) {
        errorMsg = "The AI is temporarily unavailable. Please try again in a few seconds."
      } else {
        errorMsg = "Something went wrong sending your message. Please try again."
      }
      // Remove the optimistic user message and restore it to the input so they can retry
      setMessages(prev => prev.slice(0, -1))
      setInput(userMessage)
      toast.error(errorMsg)
    } finally {
      setLoading(false)
      requestAnimationFrame(() => inputRef.current?.focus())
    }
  }

  const handleTaskSubmit = async (responses: Array<{ type: string; state: string; value: string; text?: string; label?: string }>) => {
    if (!activeChatId || loading) return

    // Format task responses into a structured message for the AI.
    // Pending tasks are included so the AI knows they weren't completed yet.
    const parts: string[] = []
    for (const r of responses) {
      const name = r.type === 'question' ? `Q: ${r.text}` : r.label || 'Check'
      if (r.state === 'done' && r.value.trim()) {
        parts.push(`**${name}:**\n\`\`\`\n${r.value.trim()}\n\`\`\``)
      } else if (r.state === 'skipped') {
        parts.push(`**${name}:** _(skipped)_`)
      } else {
        parts.push(`**${name}:** _(not yet completed)_`)
      }
    }
    const userMessage = parts.join('\n\n')

    setMessages(prev => [...prev, { role: 'user', content: userMessage }])
    setLoading(true)

    const sentForChatId = activeChatId
    try {
      const response = await aiSessionsApi.sendChatMessage(activeChatId, { message: userMessage })
      // Guard: discard if user switched to a different chat while this was in flight
      if (currentChatRef.current !== sentForChatId) return
      setMessages(prev => [
        ...prev,
        { role: 'assistant', content: response.content, suggestedFlows: response.suggested_flows, fork: response.fork, actions: response.actions, questions: response.questions },
      ])
      if (response.fork && sentForChatId) {
        branching.loadBranches(sentForChatId)
      }
      // Update task lane based on AI response
      const hasQuestions = response.questions && response.questions.length > 0
      const hasActions = response.actions && response.actions.length > 0
      clearTaskState(sentForChatId)
      if (hasQuestions || hasActions) {
        setActiveQuestions(response.questions || [])
        setActiveActions(response.actions || [])
        setShowTaskLane(true)
      } else {
        // AI sent no new tasks — clear the lane
        setShowTaskLane(false)
        setActiveQuestions([])
        setActiveActions([])
      }
    } catch (err: unknown) {
      console.error('[AssistantChat] handleTaskSubmit failed:', err)
      const status = (err as { response?: { status?: number } })?.response?.status
      const errorMsg = status === 429
        ? "You're sending messages too quickly. Wait a moment and try again."
        : "Something went wrong submitting your responses. Please try again."
      setMessages(prev => prev.slice(0, -1))
      toast.error(errorMsg)
    } finally {
      setLoading(false)
    }
  }

  const handleConclude = async (outcome: ConclusionOutcome, _notes: string): Promise<string> => {
    if (!activeChatId) throw new Error('No active chat')

    if (outcome === 'resolved') {
      await aiSessionsApi.resolveSession(activeChatId, {
        resolution_summary: _notes || 'Resolved via assistant chat',
      })
      setActiveSessionStatus('resolved')
      return activeChatId
    } else if (outcome === 'escalated') {
      await aiSessionsApi.escalateSession(activeChatId, {
        escalation_reason: _notes || 'Escalated from assistant chat',
      })
      setActiveSessionStatus('escalated')
      return activeChatId
    } else {
      await aiSessionsApi.pauseSession(activeChatId)
      setActiveSessionStatus('paused')
      return activeChatId
    }
  }

  const handleResumeNew = async (summary: string) => {
    // Invalidate currentChatRef BEFORE the await — same guard as handleNewChat
    currentChatRef.current = null
    // Clear stale state immediately — don't wait for API to return
    setShowTaskLane(false)
    setActiveQuestions([])
    setActiveActions([])
    setActiveSessionStatus('active')
    setActivePsaTicketId(null)
    try {
      const resumePrompt = `I'm continuing a previous troubleshooting session. Here's where we left off:\n\n${summary}\n\nPlease review this context and help me continue from where we stopped.`
      const session = await aiSessionsApi.createChatSession({
        intake_type: 'free_text',
        intake_content: { text: resumePrompt },
      })
      const chatItem: ChatListItem = {
        id: session.session_id,
        title: session.title,
        message_count: 0,
        pinned: false,
        created_at: new Date().toISOString(),
        updated_at: new Date().toISOString(),
      }
      currentChatRef.current = session.session_id
      setChats(prev => [chatItem, ...prev])
      setActiveChatId(session.session_id)
      setMessages([{ role: 'user', content: resumePrompt }])
      setLoading(true)

      const response = await aiSessionsApi.sendChatMessage(session.session_id, { message: resumePrompt })
      // Guard: discard if user switched to a different chat while this was in flight
      if (currentChatRef.current !== session.session_id) return
      setMessages(prev => [
        ...prev,
        { role: 'assistant', content: response.content, suggestedFlows: response.suggested_flows, fork: response.fork, actions: response.actions, questions: response.questions },
      ])
      setChats(prev =>
        prev.map(c =>
          c.id === session.session_id
            ? { ...c, message_count: 2, title: resumePrompt.slice(0, 100), updated_at: new Date().toISOString() }
            : c
        )
      )
      // Show task lane if AI sent questions or actions
      if (response.fork && session.session_id) {
        branching.loadBranches(session.session_id)
      }
      const hasQuestions = response.questions && response.questions.length > 0
      const hasActions = response.actions && response.actions.length > 0
      if (hasQuestions || hasActions) {
        setActiveQuestions(response.questions || [])
        setActiveActions(response.actions || [])
        setShowTaskLane(true)
      }
    } catch {
      toast.error('Failed to create resume chat')
    } finally {
      setLoading(false)
    }
  }

  const handleKeyDown = (e: React.KeyboardEvent) => {
    if (e.key === 'Enter' && !e.shiftKey) {
      e.preventDefault()
      handleSend()
    }
  }

  // Auto-grow textarea
  useEffect(() => {
    const el = inputRef.current
    if (!el) return
    el.style.height = 'auto'
    el.style.height = `${Math.min(el.scrollHeight, 150)}px`
  }, [input])

  // ── File handling ──────────────────────────────
  const ACCEPTED_FILE_TYPES = 'image/png,image/jpeg,image/gif,image/webp,.txt,.log,.csv,.pdf,.docx'

  const processFiles = useCallback((files: File[]) => {
    if (files.length === 0) return
    const newUploads: PendingUpload[] = files.map((file) => ({
      id: crypto.randomUUID(),
      file,
      preview: file.type.startsWith('image/') ? URL.createObjectURL(file) : '',
      status: 'uploading' as const,
    }))
    setPendingUploads((prev) => [...prev, ...newUploads])
    newUploads.forEach((upload) => {
      uploadsApi.upload(upload.file)
        .then((result) => {
          setPendingUploads((prev) => prev.map((u) => u.id === upload.id ? { ...u, status: 'done' as const, result } : u))
        })
        .catch((err) => {
          const is503 = err?.response?.status === 503
          if (is503) {
            toast.warning('Image attachments are not available yet — describe the issue in text instead')
          } else {
            toast.error(`Upload failed: ${err?.message || 'Unknown error'}`)
          }
          setPendingUploads((prev) => prev.filter((u) => u.id !== upload.id))
        })
    })
  }, [])

  const handlePaste = useCallback((e: React.ClipboardEvent<HTMLTextAreaElement>) => {
    const items = e.clipboardData?.items
    if (!items) return
    const imageFiles: File[] = []
    for (let i = 0; i < items.length; i++) {
      if (items[i].type.startsWith('image/')) {
        const file = items[i].getAsFile()
        if (file) imageFiles.push(file)
      }
    }
    if (imageFiles.length > 0) {
      e.preventDefault()
      processFiles(imageFiles)
    }
  }, [processFiles])

  const handleDragOver = useCallback((e: React.DragEvent) => { e.preventDefault(); e.dataTransfer.dropEffect = 'copy' }, [])
  const handleDragEnter = useCallback((e: React.DragEvent) => { e.preventDefault(); dragCounterRef.current++; if (dragCounterRef.current === 1) setIsDragOver(true) }, [])
  const handleDragLeave = useCallback((e: React.DragEvent) => { e.preventDefault(); dragCounterRef.current--; if (dragCounterRef.current === 0) setIsDragOver(false) }, [])
  const handleDrop = useCallback((e: React.DragEvent) => { e.preventDefault(); dragCounterRef.current = 0; setIsDragOver(false); processFiles(Array.from(e.dataTransfer.files)) }, [processFiles])
  const handleFileSelect = useCallback((e: React.ChangeEvent<HTMLInputElement>) => { if (e.target.files) { processFiles(Array.from(e.target.files)); e.target.value = '' } }, [processFiles])
  const handleRemoveUpload = useCallback((uploadId: string) => {
    setPendingUploads((prev) => { const toRemove = prev.find((u) => u.id === uploadId); if (toRemove?.preview) URL.revokeObjectURL(toRemove.preview); return prev.filter((u) => u.id !== uploadId) })
  }, [])
  const retryUpload = useCallback((uploadId: string) => {
    const upload = pendingUploads.find((u) => u.id === uploadId)
    if (!upload) return
    setPendingUploads((prev) => prev.map((u) => u.id === uploadId ? { ...u, status: 'uploading' as const, error: undefined } : u))
    uploadsApi.upload(upload.file)
      .then((result) => { setPendingUploads((prev) => prev.map((u) => u.id === uploadId ? { ...u, status: 'done' as const, result } : u)) })
      .catch((err) => { setPendingUploads((prev) => prev.map((u) => u.id === uploadId ? { ...u, status: 'error' as const, error: err?.message || 'Upload failed' } : u)) })
  }, [pendingUploads])

  // Cleanup blob URLs on unmount
  useEffect(() => { return () => { pendingUploads.forEach((u) => { if (u.preview) URL.revokeObjectURL(u.preview) }) } }, []) // eslint-disable-line react-hooks/exhaustive-deps

  return (
    <>
    <PageMeta title="AI Assistant" />
    <div className="flex h-[calc(100vh-3.5rem)]">
      {/* Sidebar — hidden on mobile, collapsed to top bar or full sidebar on desktop */}
      {!sidebarCollapsed && (
        <div className="hidden sm:block">
          <ChatSidebar
            chats={chats}
            activeChatId={activeChatId}
            onSelectChat={selectChat}
            onNewChat={handleNewChat}
            onDeleteChat={handleDeleteChat}
            onTogglePin={handleTogglePin}
            onToggleCollapse={toggleSidebarCollapse}
          />
        </div>
      )}
      <div className="sm:hidden">
        <ChatSidebar
          chats={chats}
          activeChatId={activeChatId}
          onSelectChat={selectChat}
          onNewChat={handleNewChat}
          onDeleteChat={handleDeleteChat}
          onTogglePin={handleTogglePin}
          mobileOpen={mobileSidebarOpen}
          onMobileClose={() => setMobileSidebarOpen(false)}
        />
      </div>

      {/* Main chat area + optional branch sidebar */}
      <div className="flex-1 flex flex-col min-w-0">

      {/* Collapsed sidebar top bar — desktop only */}
      {sidebarCollapsed && (
        <div className="hidden sm:block">
          <ChatSidebarCollapsedBar
            chats={chats}
            activeChatId={activeChatId}
            onNewChat={handleNewChat}
            onExpand={toggleSidebarCollapse}
          />
        </div>
      )}

      {/* Chat content row: chat column + TaskLane side by side */}
      <div className="flex-1 flex min-w-0 min-h-0">
      <div className="flex-1 flex flex-col min-w-0">
        {/* Mobile header with chat history toggle */}
        <div className="sm:hidden flex items-center gap-2 px-3 py-2 border-b border-border shrink-0">
          <button
            onClick={() => setMobileSidebarOpen(true)}
            className="flex items-center gap-2 rounded-lg px-3 py-2 text-sm text-muted-foreground hover:text-foreground hover:bg-[var(--color-bg-elevated)] transition-colors"
          >
            <MessageSquare size={16} />
            Chats
          </button>
          <div className="flex-1" />
          <button
            onClick={handleNewChat}
            className="rounded-lg px-3 py-2 text-sm font-medium text-primary hover:bg-primary/10 transition-colors"
          >
            + New
          </button>
        </div>

        {activeChatId ? (
          <>
            {/* Session header — title + lifecycle actions */}
            {(() => {
              const chatTitle = chats.find(c => c.id === activeChatId)?.title
              const isActive = activeSessionStatus === 'active' || activeSessionStatus === null
              const canAct = messages.length >= 2 && isActive && !loading
              const updateLabel = activePsaTicketId ? 'Update Ticket' : 'Share Update'
              return (
                <div className="flex items-center gap-3 border-b border-border px-3 sm:px-5 py-2.5 shrink-0">
                  <span className="flex h-7 w-7 items-center justify-center rounded-lg bg-accent-dim shrink-0">
                    <Sparkles size={14} className="text-primary" />
                  </span>
                  <div className="flex-1 min-w-0">
                    <h1 className="font-heading text-sm font-semibold text-foreground truncate">
                      {chatTitle || 'AI Assistant'}
                    </h1>
                    {activeSessionStatus && activeSessionStatus !== 'active' && (
                      <span className={cn(
                        'text-[0.625rem] font-medium uppercase tracking-wide',
                        activeSessionStatus === 'resolved' ? 'text-success' :
                        activeSessionStatus === 'escalated' || activeSessionStatus === 'requesting_escalation' ? 'text-warning' :
                        activeSessionStatus === 'paused' ? 'text-muted-foreground' : 'text-muted-foreground'
                      )}>
                        {activeSessionStatus === 'requesting_escalation' ? 'Escalated' : activeSessionStatus.charAt(0).toUpperCase() + activeSessionStatus.slice(1)}
                      </span>
                    )}
                  </div>

                  {/* Desktop actions — shown when session is active and has messages */}
                  <div className="hidden sm:flex items-center gap-1.5">
                    {isActive && (
                      <>
                        <button
                          onClick={() => setShowConclude(true)}
                          disabled={!canAct}
                          data-conclude-outcome="resolved"
                          className="flex items-center gap-1.5 rounded-lg bg-success-dim border border-success/20 px-3 py-1.5 text-xs font-medium text-success hover:bg-success/20 disabled:opacity-40 disabled:pointer-events-none transition-colors"
                        >
                          <CheckCircle2 size={13} />
                          Resolve
                        </button>
                        <button
                          onClick={() => setShowConclude(true)}
                          disabled={!canAct}
                          data-conclude-outcome="escalated"
                          className="flex items-center gap-1.5 rounded-lg bg-warning-dim border border-warning/20 px-3 py-1.5 text-xs font-medium text-warning hover:bg-warning/20 disabled:opacity-40 disabled:pointer-events-none transition-colors"
                        >
                          <ArrowUpRight size={13} />
                          Escalate
                        </button>
                      </>
                    )}
                    {messages.length >= 2 && (
                      <button
                        onClick={() => setShowStatusUpdate(true)}
                        disabled={loading}
                        className="flex items-center gap-1.5 rounded-lg bg-accent-dim border border-accent/20 px-3 py-1.5 text-xs font-medium text-accent hover:bg-accent/20 disabled:opacity-40 disabled:pointer-events-none transition-colors"
                      >
                        <FileText size={13} />
                        {updateLabel}
                      </button>
                    )}
                    {/* Overflow: Pause / — */}
                    {isActive && messages.length >= 2 && (
                      <div className="relative">
                        <button
                          onClick={() => setShowOverflow(!showOverflow)}
                          className="flex items-center justify-center rounded-lg px-2 py-1.5 text-muted-foreground hover:text-foreground hover:bg-[rgba(255,255,255,0.06)] transition-colors"
                        >
                          <MoreHorizontal size={16} />
                        </button>
                        {showOverflow && (
                          <>
                            <div className="fixed inset-0 z-40" onClick={() => setShowOverflow(false)} />
                            <div className="absolute right-0 top-full mt-1 z-50 w-36 rounded-lg border border-border bg-card py-1 shadow-lg">
                              <button
                                onClick={() => { setShowOverflow(false); aiSessionsApi.pauseSession(activeChatId).then(() => setActiveSessionStatus('paused')).catch(() => toast.error('Failed to pause')) }}
                                className="flex w-full items-center gap-2 px-3 py-2 text-xs text-muted-foreground hover:text-foreground hover:bg-[rgba(255,255,255,0.06)] transition-colors"
                              >
                                <Pause size={13} />
                                Pause
                              </button>
                            </div>
                          </>
                        )}
                      </div>
                    )}
                  </div>

                  {/* Mobile: single overflow menu */}
                  {messages.length >= 2 && (
                    <div className="sm:hidden relative">
                      <button
                        onClick={() => setShowOverflow(!showOverflow)}
                        className="flex items-center justify-center rounded-lg px-2 py-1.5 text-muted-foreground hover:text-foreground hover:bg-[rgba(255,255,255,0.06)] transition-colors"
                      >
                        <MoreHorizontal size={18} />
                      </button>
                      {showOverflow && (
                        <>
                          <div className="fixed inset-0 z-40" onClick={() => setShowOverflow(false)} />
                          <div className="absolute right-0 top-full mt-1 z-50 w-44 rounded-lg border border-border bg-card py-1 shadow-lg">
                            {isActive && (
                              <>
                                <button
                                  onClick={() => { setShowOverflow(false); setShowConclude(true) }}
                                  disabled={!canAct}
                                  className="flex w-full items-center gap-2 px-3 py-2.5 text-xs text-success hover:bg-success-dim transition-colors disabled:opacity-40"
                                >
                                  <CheckCircle2 size={14} />
                                  Resolve
                                </button>
                                <button
                                  onClick={() => { setShowOverflow(false); setShowConclude(true) }}
                                  disabled={!canAct}
                                  className="flex w-full items-center gap-2 px-3 py-2.5 text-xs text-warning hover:bg-warning-dim transition-colors disabled:opacity-40"
                                >
                                  <ArrowUpRight size={14} />
                                  Escalate
                                </button>
                              </>
                            )}
                            <button
                              onClick={() => { setShowOverflow(false); setShowStatusUpdate(true) }}
                              disabled={loading}
                              className="flex w-full items-center gap-2 px-3 py-2.5 text-xs text-accent hover:bg-accent-dim transition-colors disabled:opacity-40"
                            >
                              <FileText size={14} />
                              {updateLabel}
                            </button>
                            {isActive && (
                              <button
                                onClick={() => { setShowOverflow(false); aiSessionsApi.pauseSession(activeChatId).then(() => setActiveSessionStatus('paused')).catch(() => toast.error('Failed to pause')) }}
                                className="flex w-full items-center gap-2 px-3 py-2.5 text-xs text-muted-foreground hover:text-foreground hover:bg-[rgba(255,255,255,0.06)] transition-colors"
                              >
                                <Pause size={14} />
                                Pause
                              </button>
                            )}
                          </div>
                        </>
                      )}
                    </div>
                  )}
                </div>
              )
            })()}

            {/* Messages */}
            <div className="flex-1 overflow-y-auto px-4 sm:px-6 py-4 space-y-4">
              {messages.length === 0 && !loading && (
                <div className="flex flex-col items-center justify-center h-full text-center">
                  <div className="w-16 h-16 rounded-full bg-accent-dim flex items-center justify-center mb-4">
                    <Sparkles size={28} className="text-primary" />
                  </div>
                  <h2 className="text-lg font-heading font-semibold text-foreground mb-2">
                    AI Assistant
                  </h2>
                  <p className="text-sm text-muted-foreground max-w-md">
                    Ask me anything about IT infrastructure, networking, Active Directory,
                    cloud platforms, or troubleshooting. I'll also suggest relevant flows from your team's library.
                  </p>
                </div>
              )}
              {messages.map((msg, i) => (
                <ChatMessage
                  key={i}
                  role={msg.role}
                  content={msg.content}
                  suggestedFlows={msg.suggestedFlows}
                  imageUrls={msg.imageUrls}
                />
              ))}
              {loading && (
                <div className="flex gap-3">
                  <div className="w-8 h-8 rounded-full bg-primary/15 flex items-center justify-center">
                    <Sparkles size={14} className="text-primary" />
                  </div>
                  <div className="bg-input border border-border rounded-2xl px-4 py-3">
                    <Loader2 size={16} className="animate-spin text-primary" />
                  </div>
                </div>
              )}
              <div ref={messagesEndRef} />
            </div>

            {/* Rich Input */}
            <div className="px-3 sm:px-6 py-3 shrink-0">
              <div
                className="max-w-3xl mx-auto"
                onDragOver={handleDragOver}
                onDragEnter={handleDragEnter}
                onDragLeave={handleDragLeave}
                onDrop={handleDrop}
              >
                <div className={cn(
                  'relative rounded-xl border transition-all',
                  loading ? 'border-border/50 opacity-50' :
                  isDragOver ? 'border-primary/50 bg-primary/5' :
                  'border-border focus-within:border-[rgba(96,165,250,0.3)] focus-within:ring-1 focus-within:ring-primary/20'
                )} style={{ background: 'var(--color-bg-card)' }}>
                  {/* Drag overlay */}
                  {isDragOver && (
                    <div className="absolute inset-0 z-10 flex items-center justify-center rounded-xl border-2 border-dashed border-primary/50 bg-primary/5 pointer-events-none">
                      <div className="flex items-center gap-2 text-sm text-primary">
                        <ImagePlus size={18} />
                        Drop files to attach
                      </div>
                    </div>
                  )}

                  {/* Textarea */}
                  <textarea
                    ref={inputRef}
                    value={input}
                    onChange={e => setInput(e.target.value)}
                    onKeyDown={handleKeyDown}
                    onPaste={handlePaste}
                    placeholder={loading ? 'AI is thinking...' : 'Type a message, paste a screenshot, or drag a file...'}
                    disabled={loading}
                    rows={1}
                    className="w-full resize-none bg-transparent px-4 pt-3 pb-1 text-sm text-foreground placeholder:text-muted-foreground focus:outline-none disabled:cursor-not-allowed"
                    style={{ minHeight: '40px', maxHeight: '150px' }}
                  />

                  {/* Thumbnail strip */}
                  {pendingUploads.length > 0 && (
                    <div className="flex gap-2 flex-wrap px-4 pb-1">
                      {pendingUploads.map((upload) => (
                        <div key={upload.id} className="relative w-12 h-12 rounded-lg overflow-hidden border border-border bg-background">
                          {upload.preview ? (
                            <img src={upload.preview} alt="" className="w-full h-full object-cover" />
                          ) : (
                            <div className="w-full h-full flex items-center justify-center text-[0.5rem] text-muted-foreground px-1 text-center">
                              {upload.file.name.split('.').pop()?.toUpperCase()}
                            </div>
                          )}
                          {upload.status === 'uploading' && (
                            <div className="absolute inset-0 bg-background/60 flex items-center justify-center">
                              <Loader2 size={12} className="animate-spin text-primary" />
                            </div>
                          )}
                          {upload.status === 'done' && (
                            <button type="button" onClick={() => handleRemoveUpload(upload.id)} className="absolute -top-1 -right-1 w-4 h-4 rounded-full bg-background border border-border flex items-center justify-center hover:bg-danger/20 transition-colors">
                              <X size={8} className="text-muted-foreground" />
                            </button>
                          )}
                          {upload.status === 'error' && (
                            <div className="absolute inset-0 bg-danger/20 border-2 border-danger flex items-center justify-center cursor-pointer" onClick={() => retryUpload(upload.id)}>
                              <RotateCcw size={10} className="text-danger" />
                            </div>
                          )}
                        </div>
                      ))}
                    </div>
                  )}

                  {/* Logs textarea */}
                  {showLogs && (
                    <div className="px-4 pb-1">
                      <div className="flex items-center justify-between mb-1">
                        <span className="text-[0.625rem] uppercase tracking-wide text-muted-foreground font-sans">Paste logs or error output</span>
                        <button type="button" onClick={() => { setShowLogs(false); setLogContent('') }} className="text-muted-foreground hover:text-foreground"><X size={14} /></button>
                      </div>
                      <textarea
                        value={logContent}
                        onChange={(e) => setLogContent(e.target.value)}
                        placeholder="Paste event viewer logs, error messages, PowerShell output..."
                        rows={3}
                        className="w-full resize-none rounded-lg border border-border bg-background p-2 font-mono text-xs text-foreground placeholder:text-muted-foreground focus:border-[rgba(96,165,250,0.3)] focus:outline-none"
                      />
                    </div>
                  )}

                  {/* Bottom toolbar */}
                  <div className="flex items-center justify-between px-3 py-1.5 border-t border-border/50">
                    <div className="flex items-center gap-0.5">
                      <input ref={fileInputRef} type="file" multiple accept={ACCEPTED_FILE_TYPES} onChange={handleFileSelect} className="hidden" />
                      <button type="button" onClick={() => fileInputRef.current?.click()} disabled={loading} className="flex items-center gap-1.5 rounded-lg px-2 py-1.5 text-xs text-muted-foreground hover:text-foreground hover:bg-secondary transition-colors disabled:opacity-40" title="Attach files">
                        <Paperclip size={14} />
                        <span className="hidden sm:inline">Attach</span>
                      </button>
                      {!showLogs && (
                        <button type="button" onClick={() => setShowLogs(true)} disabled={loading} className="flex items-center gap-1.5 rounded-lg px-2 py-1.5 text-xs text-muted-foreground hover:text-foreground hover:bg-secondary transition-colors disabled:opacity-40" title="Paste logs">
                          <Terminal size={14} />
                          <span className="hidden sm:inline">Paste Logs</span>
                        </button>
                      )}
                      {!showTaskLane && (activeQuestions.length > 0 || activeActions.length > 0) && (
                        <button
                          type="button"
                          onClick={() => setShowTaskLane(true)}
                          className="flex items-center gap-1.5 rounded-lg px-2 py-1.5 text-xs text-accent-text hover:text-foreground hover:bg-accent-dim transition-colors"
                          title="Show task panel"
                        >
                          <ListChecks size={14} />
                          Tasks ({activeQuestions.length + activeActions.length})
                        </button>
                      )}
                    </div>
                    <button type="button" onClick={handleSend} disabled={!input.trim() || loading} className={cn(
                      'flex h-8 w-8 items-center justify-center rounded-lg transition-all',
                      input.trim() && !loading ? 'bg-primary text-white hover:brightness-110 active:scale-95' : 'bg-secondary text-muted-foreground cursor-not-allowed'
                    )} title="Send message">
                      <Send size={15} />
                    </button>
                  </div>
                </div>
              </div>
            </div>
          </>
        ) : (
          <div className="flex flex-col items-center justify-center h-full text-center">
            <div className="w-20 h-20 rounded-full bg-accent-dim flex items-center justify-center mb-4">
              <Sparkles size={32} className="text-primary" />
            </div>
            <h2 className="text-xl font-heading font-semibold text-foreground mb-2">
              AI Assistant
            </h2>
            <p className="text-sm text-muted-foreground max-w-md mb-6">
              Your Senior Systems & Network Engineer. Ask anything about IT infrastructure,
              or start a new chat to get personalized help with your team's flows.
            </p>
            <button
              onClick={handleNewChat}
              className="bg-primary text-white font-semibold text-sm rounded-lg px-6 py-2.5 hover:brightness-110 active:scale-[0.98] transition-all"
            >
              Start a Conversation
            </button>
          </div>
        )}
      </div>

      {/* Task lane — slides in when AI sends questions or actions */}
      {showTaskLane && (activeQuestions.length > 0 || activeActions.length > 0) && (
        <TaskLane
          questions={activeQuestions}
          actions={activeActions}
          sessionId={activeChatId}
          onSubmit={handleTaskSubmit}
          onClose={() => {
            setShowTaskLane(false)
          }}
          loading={loading}
        />
      )}

      {/* Branch map hidden — branching is now silent/background only.
          Branches are tracked in the DB but not shown to the user.
          The AI manages branch context internally. */}
      </div>{/* close chat content row */}
      </div>{/* close outer flex-col */}

      {/* Conclude Session Modal */}
      <ConcludeSessionModal
        isOpen={showConclude}
        onClose={() => setShowConclude(false)}
        onConclude={handleConclude}
        onResumeNew={handleResumeNew}
        chatTitle={chats.find(c => c.id === activeChatId)?.title ?? 'Chat'}
        sessionId={activeChatId}
      />

      {/* Status Update Modal */}
      {activeChatId && (
        <StatusUpdateModal
          open={showStatusUpdate}
          onClose={() => setShowStatusUpdate(false)}
          onGenerate={(audience, length, context) =>
            aiSessionsApi.generateStatusUpdate(activeChatId, { audience, length, context })
          }
          context="status"
        />
      )}
    </div>
    </>
  )
}
