import { useState, useEffect, useRef, useCallback } from 'react'
import { useLocation, useNavigate, useParams } from 'react-router-dom'
import { uploadsApi } from '@/api/uploads'
import type { PendingUpload } from '@/types/upload'
import type { ForkMetadata, ActionItem, QuestionItem, ChatMessageResponse, TriageUpdate } from '@/types/ai-session'
import { aiSessionsApi } from '@/api/aiSessions'
import { useBranching } from '@/hooks/useBranching'
import { analytics } from '@/lib/analytics'
import { toast } from '@/lib/toast'
import type { ChatListItem, ConclusionOutcome } from '@/types/assistant-chat'
import type { SuggestedFlow } from '@/types/copilot'

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

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

export function useAssistantSession() {
  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 === urlSessionId || d.chatId === sessionStorage.getItem('rf-active-chat-id')) 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 === urlSessionId || d.chatId === sessionStorage.getItem('rf-active-chat-id')) 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 === urlSessionId || d.chatId === sessionStorage.getItem('rf-active-chat-id')) }
    } catch { /* ignore */ }
    return false
  })
  const [sidebarCollapsed, setSidebarCollapsed] = useState(() =>
    localStorage.getItem('rf-chat-sidebar-collapsed') === 'true'
  )

  const messagesEndRef = useRef<HTMLDivElement>(null)
  const inputRef = useRef<HTMLTextAreaElement>(null)
  const fileInputRef = useRef<HTMLInputElement>(null)
  const dragCounterRef = useRef(0)
  const prefillHandledRef = useRef(false)
  const currentChatRef = useRef<string | null>(activeChatId)
  const loadingRef = useRef(false)
  const initialLoadDoneRef = useRef(false)

  const toggleSidebarCollapse = () => {
    const next = !sidebarCollapsed
    setSidebarCollapsed(next)
    localStorage.setItem('rf-chat-sidebar-collapsed', String(next))
  }

  // 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 on mount
  useEffect(() => { loadChats() }, [])

  // Load session data on mount or when URL session changes.
  // On initial mount, always load even if activeChatId matches urlSessionId
  // (state is empty after view toggle between /assistant and /cockpit).
  useEffect(() => {
    if (urlSessionId) {
      if (!initialLoadDoneRef.current || urlSessionId !== activeChatId) {
        selectChat(urlSessionId)
      }
      initialLoadDoneRef.current = true
    } else if (!initialLoadDoneRef.current && activeChatId) {
      // Restore session from sessionStorage on mount (when URL has no session ID)
      selectChat(activeChatId)
      initialLoadDoneRef.current = true
    }
  }, [urlSessionId]) // eslint-disable-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])

  // 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])

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

  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
    }
  }

  // Callback for pages to handle triage data from selectChat.
  // Pages that care about triage (CockpitPage) set this; FlowPilotPage ignores it.
  const onSessionLoadedRef = useRef<((detail: {
    client_name: string | null
    asset_name: string | null
    issue_category: string | null
    triage_hypothesis: string | null
    evidence_items: Array<{ text: string; status: string }> | null
    psa_ticket_id: string | null
  }) => void) | null>(null)

  const selectChat = useCallback(async (chatId: string) => {
    currentChatRef.current = chatId
    setActiveChatId(chatId)
    // Clear task lane immediately to prevent the persist effect from
    // tagging the old session's data with the new chatId (race condition).
    // Server-side pending_task_lane will restore it below if it exists.
    setActiveQuestions([])
    setActiveActions([])
    setShowTaskLane(false)
    try {
      const detail = await aiSessionsApi.getSession(chatId)
      if (currentChatRef.current !== chatId) return
      setMessages(
        (detail.conversation_messages || []).map(m => ({
          role: m.role as 'user' | 'assistant',
          content: m.content,
        }))
      )
      // Notify page of triage data
      onSessionLoadedRef.current?.({
        client_name: detail.client_name ?? null,
        asset_name: detail.asset_name ?? null,
        issue_category: detail.issue_category ?? null,
        triage_hypothesis: detail.triage_hypothesis ?? null,
        evidence_items: detail.evidence_items ?? null,
        psa_ticket_id: detail.psa_ticket_id ?? null,
      })
      // Restore task lane from server 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) {
          const responses = (detail.pending_task_lane as Record<string, unknown>).responses as unknown[] | undefined
          if (responses && responses.length > 0) {
            try {
              sessionStorage.setItem(`rf-tasklane-state:${chatId}`, JSON.stringify(responses))
            } catch { /* ignore */ }
          }
          setActiveQuestions(q)
          setActiveActions(a)
          setShowTaskLane(true)
          return
        }
      }
      // Fallback: restore from sessionStorage (covers view toggles before backend fix)
      try {
        const saved = sessionStorage.getItem('rf-tasklane-meta')
        if (saved) {
          const d = JSON.parse(saved)
          if (d.chatId === chatId) {
            const q = d.questions || []
            const a = d.actions || []
            if (q.length > 0 || a.length > 0) {
              setActiveQuestions(q)
              setActiveActions(a)
              setShowTaskLane(d.show === true)
              return
            }
          }
        }
      } catch { /* ignore */ }
      // No task lane data from either source — clear
      setShowTaskLane(false)
      setActiveQuestions([])
      setActiveActions([])
    } catch {
      setMessages([])
      setShowTaskLane(false)
      setActiveQuestions([])
      setActiveActions([])
    }
  }, [])

  const handleNewChat = async () => {
    if (loadingRef.current) return
    loadingRef.current = true
    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)
      setMessages([])
      setShowTaskLane(false)
      setActiveQuestions([])
      setActiveActions([])
    } catch {
      toast.error('Failed to create chat')
    } finally {
      loadingRef.current = false
    }
  }

  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 () => {
    toast.info('Pin feature coming soon')
  }

  // Process an AI chat response — updates messages, task lane, branching.
  // Returns the response for page-specific handling (e.g. triage_update).
  const processResponse = useCallback((response: ChatMessageResponse, chatId: string) => {
    if (currentChatRef.current !== chatId) return null
    setMessages(prev => [
      ...prev,
      {
        role: 'assistant',
        content: response.content,
        suggestedFlows: response.suggested_flows,
        fork: response.fork,
        actions: response.actions,
        questions: response.questions,
      },
    ])
    if (response.fork && chatId) {
      branching.loadBranches(chatId)
    }
    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)
    }
    return response
  }, [branching])

  // Callback for pages to handle triage updates from chat responses.
  const onTriageUpdateRef = useRef<((update: TriageUpdate) => void) | null>(null)

  const sendMessage = useCallback(async (
    rawMessage: string,
    options?: {
      uploadIds?: string[]
      clearComposer?: boolean
    }
  ) => {
    const message = rawMessage.trim()
    if (!message || !activeChatId || loadingRef.current) return
    loadingRef.current = true

    const sendChatId = activeChatId
    const uploadIds = options?.uploadIds
    const completedUploadIds = uploadIds ?? pendingUploads
      .filter((u) => u.status === 'done' && u.result?.id)
      .map((u) => u.result!.id)
    if (options?.clearComposer !== false) {
      setInput('')
      if (!uploadIds) {
        setPendingUploads([])
      }
    }
    setMessages(prev => [...prev, { role: 'user', content: message }])
    setLoading(true)

    try {
      const response = await aiSessionsApi.sendChatMessage(sendChatId, {
        message,
        upload_ids: completedUploadIds.length > 0 ? completedUploadIds : undefined,
      })
      if (currentChatRef.current !== sendChatId) return
      analytics.aiFeatureUsed({ feature: 'assistant_chat' })
      processResponse(response, sendChatId)
      setChats(prev =>
        prev.map(c =>
          c.id === sendChatId
            ? { ...c, message_count: c.message_count + 2, title: c.message_count === 0 ? message.slice(0, 100) : c.title, updated_at: new Date().toISOString() }
            : c
        )
      )
      if (response.triage_update) onTriageUpdateRef.current?.(response.triage_update)
    } catch {
      if (currentChatRef.current !== sendChatId) return
      setMessages(prev => [
        ...prev,
        { role: 'assistant', content: 'Sorry, something went wrong. Please try again.' },
      ])
    } finally {
      loadingRef.current = false
      if (currentChatRef.current === sendChatId) {
        setLoading(false)
        requestAnimationFrame(() => inputRef.current?.focus())
      }
    }
  }, [activeChatId, pendingUploads, processResponse])

  const handleSend = async () => {
    await sendMessage(input)
  }

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

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

    const sendPrefill = async () => {
      if (loadingRef.current) return
      loadingRef.current = true
      setShowTaskLane(false)
      setActiveQuestions([])
      setActiveActions([])
      setLoading(true)
      if (logs) {
        setShowLogs(true)
        setLogContent(logs)
      }

      try {
        const initialMessage = logs
          ? `${prefill}\n\nAttached logs/output:\n\`\`\`\n${logs}\n\`\`\``
          : prefill
        const session = await aiSessionsApi.createChatSession({
          intake_type: 'free_text',
          intake_content: { text: initialMessage },
        })
        const prefillChatId = session.session_id
        currentChatRef.current = prefillChatId
        const chatItem: ChatListItem = {
          id: prefillChatId,
          title: session.title,
          message_count: 0,
          pinned: false,
          created_at: new Date().toISOString(),
          updated_at: new Date().toISOString(),
        }
        setChats(prev => [chatItem, ...prev])
        setActiveChatId(prefillChatId)
        setMessages([{ role: 'user', content: initialMessage }])

        const response = await aiSessionsApi.sendChatMessage(prefillChatId, {
          message: initialMessage,
          upload_ids: uploadIds?.length ? uploadIds : undefined,
        })
        if (currentChatRef.current !== prefillChatId) return
        processResponse(response, prefillChatId)
        setChats(prev =>
          prev.map(c =>
            c.id === prefillChatId
              ? { ...c, message_count: 2, title: prefill.slice(0, 100), updated_at: new Date().toISOString() }
              : c
          )
        )
        if (response.triage_update) onTriageUpdateRef.current?.(response.triage_update)
      } catch {
        toast.error('Failed to start AI conversation')
      } finally {
        loadingRef.current = false
        setLoading(false)
      }
    }

    sendPrefill()
  }, [location.state, navigate, processResponse])

  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',
      })
      return activeChatId
    } else if (outcome === 'escalated') {
      await aiSessionsApi.escalateSession(activeChatId, {
        escalation_reason: _notes || 'Escalated from assistant chat',
      })
      return activeChatId
    } else {
      await aiSessionsApi.pauseSession(activeChatId)
      return activeChatId
    }
  }

  const handleResumeNew = async (summary: string) => {
    if (loadingRef.current) return
    loadingRef.current = true
    setLoading(true)
    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 }])

      const resumeChatId = session.session_id
      const response = await aiSessionsApi.sendChatMessage(resumeChatId, { message: resumePrompt })
      if (currentChatRef.current !== resumeChatId) return
      processResponse(response, resumeChatId)
      setChats(prev =>
        prev.map(c =>
          c.id === resumeChatId
            ? { ...c, message_count: 2, title: resumePrompt.slice(0, 100), updated_at: new Date().toISOString() }
            : c
        )
      )
      if (response.triage_update) onTriageUpdateRef.current?.(response.triage_update)
    } catch {
      toast.error('Failed to create resume chat')
    } finally {
      loadingRef.current = false
      setLoading(false)
    }
  }

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

  // ── File handling ──────────────────────────────

  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])

  return {
    // State
    chats, activeChatId, messages, input, loading,
    showConclude, showStatusUpdate, branching,
    mobileSidebarOpen, showLogs, logContent,
    pendingUploads, isDragOver,
    activeQuestions, activeActions, showTaskLane,
    sidebarCollapsed,
    // Setters
    setInput, setShowConclude, setShowStatusUpdate,
    setMobileSidebarOpen, setShowLogs, setLogContent,
    setShowTaskLane, setActiveQuestions, setActiveActions,
    // Handlers
    loadChats, selectChat, handleNewChat, handleDeleteChat, handleTogglePin,
    handleSend, sendMessage, handleConclude, handleResumeNew,
    handleKeyDown, handlePaste,
    handleDragOver, handleDragEnter, handleDragLeave, handleDrop,
    handleFileSelect, handleRemoveUpload, retryUpload,
    toggleSidebarCollapse, handlePrefill, processResponse,
    // Refs
    messagesEndRef, inputRef, fileInputRef, currentChatRef, loadingRef,
    // Page-specific callbacks
    onSessionLoadedRef, onTriageUpdateRef,
    // Constants
    ACCEPTED_FILE_TYPES,
  }
}