import { useState, useRef, useEffect } from 'react'
import { Client } from '@langchain/langgraph-sdk'
import ReactMarkdown from 'react-markdown'
import remarkGfm from 'remark-gfm'
import { useAuth } from '../../context/useAuth.js'
import Sidebar from '../Sidebar/Sidebar.jsx'
import TimeTravelPanel from '../TimeTravel/TimeTravelPanel.jsx'
import styles from './Chat.module.css'

const API_BASE = import.meta.env.VITE_API_BASE || 'https://cripple-lee-personal-assistant-server.hf.space'

// Module-scope client whose Authorization header tracks the current session.
// The SDK reads the protected `defaultHeaders` field on every request, so
// mutating it on login/logout is enough — no need to recreate the client.
const client = new Client({ apiUrl: API_BASE })

export function setClientAccessToken(token) {
  // `new Client(config)` builds separate sub-client instances (threads, runs,
  // …) and each BaseClient constructor copies `config.defaultHeaders` into its
  // OWN object. So `client.threads.defaultHeaders` and `client.defaultHeaders`
  // are different objects — mutating only the parent never reaches the
  // sub-clients that actually issue requests, and the Authorization header is
  // silently dropped (server responds 401 "Missing Authorization header").
  // Set/clear the header on the parent AND every sub-client.
  const targets = [client, client.threads, client.runs, client.assistants, client.crons, client.store]
  for (const c of targets) {
    if (!c) continue
    if (token) {
      c.defaultHeaders = {
        ...(c.defaultHeaders || {}),
        Authorization: `Bearer ${token}`,
      }
    } else if (c.defaultHeaders) {
      delete c.defaultHeaders.Authorization
    }
  }
}

// Fetch the current user's monthly token usage from the custom /usage route.
async function fetchUsage() {
  const headers = { ...(client.defaultHeaders || {}) }
  const res = await fetch(`${API_BASE}/usage`, { headers })
  if (!res.ok) throw new Error(`Failed to load usage (${res.status})`)
  return res.json()
}

function ChatMessage({ role, content, isStreaming }) {
  return (
    <div className={`${styles.message} ${role === 'human' ? styles.user : styles.agent}`}>
      <div className={styles.bubble}>
        <span className={styles.role}>{role === 'human' ? 'You' : 'AI'}</span>
        {role === 'human' ? (
          <p className={styles.content}>
            {content}
            {role === 'agent' && isStreaming && <span className={styles.cursor} aria-hidden="true" />}
          </p>
        ) : (
          <div className={styles.content}>
            <ReactMarkdown remarkPlugins={[remarkGfm]}>{content}</ReactMarkdown>
            {isStreaming && <span className={styles.cursor} aria-hidden="true" />}
          </div>
        )}
      </div>
    </div>
  )
}

const GREETING = { role: 'assistant', content: 'Hello! How can I help you today?', id: 0 }

export default function Chat() {
  const { accessToken, user, logout } = useAuth()
  const [messages, setMessages] = useState([GREETING])
  const [input, setInput] = useState('')
  const [loading, setLoading] = useState(false)
  const [error, setError] = useState(null)
  const [streamingId, setStreamingId] = useState(null)
  const [threadId, setThreadId] = useState(null)
  const [threads, setThreads] = useState([])
  const [showTimeTravel, setShowTimeTravel] = useState(false)
  const [history, setHistory] = useState([])
  const [historyLoading, setHistoryLoading] = useState(false)
  const [activeCheckpointId, setActiveCheckpointId] = useState(null)
  // Remembers the forked checkpoint per thread, so reopening a conversation
  // from the sidebar shows the messages at the fork point instead of the tip.
  const [forkedCheckpoints, setForkedCheckpoints] = useState({})
  // Monthly token usage: { tokens_used, token_limit, remaining, limit_reached }.
  const [usage, setUsage] = useState(null)
  const messagesEndRef = useRef(null)
  // Identifies the in-flight server-side run ({ threadId, runId }) so
  // Pause can cancel it on the server, not just close the local stream.
  const activeRunRef = useRef(null)

  // Point the SDK client at the current session's access token.
  useEffect(() => {
    setClientAccessToken(accessToken)
    return () => setClientAccessToken(null)
  }, [accessToken])

  // Load monthly token usage on mount / login so the banner and composer
  // reflect the current allowance. State updates are deferred to a microtask
  // so no setState runs synchronously in the effect body.
  useEffect(() => {
    let cancelled = false
    queueMicrotask(() => {
      if (!accessToken) {
        if (!cancelled) setUsage(null)
        return
      }
      fetchUsage()
        .then((u) => { if (!cancelled) setUsage(u) })
        .catch((err) => console.error('Failed to load token usage:', err))
    })
    return () => { cancelled = true }
  }, [accessToken])

  const scrollToBottom = () => {
    messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' })
  }

  useEffect(() => {
    scrollToBottom()
  }, [messages])

  // Load persisted threads on mount so past conversations can be reopened.
  useEffect(() => {
    const loadThreads = async () => {
      try {
        const result = await client.threads.search({ limit: 20 })
        // Fetch message previews for each thread to show summaries in sidebar
        const threadsWithPreviews = await Promise.all(
          result.map(async (thread) => {
            try {
              const state = await client.threads.getState(thread.thread_id)
              const messages = state.values?.messages || []
              const firstUserMessage = messages.find((m) => m.type === 'human')
              const preview = firstUserMessage?.content
                ? String(firstUserMessage.content).slice(0, 50) +
                  (String(firstUserMessage.content).length > 50 ? '…' : '')
                : 'New conversation'
              return { ...thread, preview }
            } catch {
              return { ...thread, preview: 'New conversation' }
            }
          })
        )
        setThreads(threadsWithPreviews)
      } catch (err) {
        console.error('Failed to load threads:', err)
      }
    }
    loadThreads()
  }, [])

  const loadThreadHistory = async (id) => {
    if (loading) return
    setLoading(true)
    setError(null)
    try {
      // If this conversation was forked at a checkpoint earlier, check out
      // the messages at that checkpoint instead of the thread tip.
      const forkCheckpointId = forkedCheckpoints[id]
      const state = forkCheckpointId
        ? await client.threads.getState(id, forkCheckpointId)
        : await client.threads.getState(id)
      const history = (state.values?.messages || [])
        .filter((m) => m?.content)
        .map((m, i) => ({
          role: m.type === 'human' ? 'human' : 'assistant',
          content: typeof m.content === 'string' ? m.content : '',
          id: m.id || i,
        }))
      setThreadId(id)
      // Drop the previous thread's checkpoints; the sync effect reloads them
      // for the newly selected thread if the panel is open.
      setHistory([])
      setMessages(history.length > 0 ? [GREETING, ...history] : [GREETING])
      setStreamingId(null)
      setInput('')
    } catch (err) {
      setError(err.message)
    } finally {
      setLoading(false)
    }
  }

  // Fetch the checkpoint history for the given thread (newest first).
  const refreshHistory = async (id) => {
    if (!id) {
      setHistory([])
      return
    }
    setHistoryLoading(true)
    setError(null)
    try {
      const states = await client.threads.getHistory(id, { limit: 50 })
      setHistory(states)
    } catch (err) {
      setError(err.message)
    } finally {
      setHistoryLoading(false)
    }
  }

  // Keep the time travel panel in sync with the selected conversation:
  // whenever the active thread changes while the panel is open, reload that
  // thread's checkpoints. The fetch is deferred to a microtask so no setState
  // runs synchronously in the effect body.
  useEffect(() => {
    if (!showTimeTravel) return
    let cancelled = false
    queueMicrotask(async () => {
      setHistoryLoading(true)
      try {
        const states = await client.threads.getHistory(threadId, { limit: 50 })
        if (!cancelled) setHistory(states)
      } catch (err) {
        if (!cancelled) setError(err.message)
      } finally {
        if (!cancelled) setHistoryLoading(false)
      }
    })
    return () => {
      cancelled = true
    }
  }, [showTimeTravel, threadId])

  const handleToggleTimeTravel = () => {
    setShowTimeTravel((open) => !open)
  }

  const handleForkFromCheckpoint = async (state) => {
    console.log('Forking from checkpoint state:', state)
    if (loading || !threadId) return
    const checkpointId = state.checkpoint?.checkpoint_id
    if (!checkpointId) return
    setLoading(true)
    setError(null)
    try {
      // Load the state at the chosen checkpoint so the chat view shows the
      // conversation exactly as it was at that point in time.
      const pastState = await client.threads.getState(threadId, checkpointId)
      const pastMessages = (pastState.values?.messages || [])
        .filter((m) => m?.content)
        .map((m, i) => ({
          role: m.type === 'human' ? 'human' : 'assistant',
          content: typeof m.content === 'string' ? m.content : '',
          id: m.id || i,
        }))
      // Restore the last human message as the input so the user can continue
      // the conversation from that point, and show the rest of the messages in the chat view.
      const lastMessage = pastMessages.pop()
      setInput(lastMessage?.role === 'human' ? lastMessage.content : '')
      setMessages(pastMessages.length > 0 ? [GREETING, ...pastMessages] : [GREETING])
      // Remember the fork point; the next run passes it as checkpointId so
      // LangGraph forks the thread from this checkpoint instead of the tip.
      setActiveCheckpointId(checkpointId)
      // Associate the fork with this thread so selecting the conversation in
      // the sidebar later checks out the messages at this checkpoint.
      setForkedCheckpoints((prev) => ({ ...prev, [threadId]: checkpointId }))
      // Trim the checkpoint list to the timeline at the fork point: history
      // is newest-first, so keep everything from the forked checkpoint's
      // index onward (i.e. that checkpoint and all earlier ones).
      setHistory((prev) => {
        const forkIndex = prev.findIndex(
          (s) => s.checkpoint?.checkpoint_id === checkpointId
        )
        return forkIndex >= 0 ? prev.slice(forkIndex) : prev
      })
      // Clear the remaining transient states so the chat starts fresh from
      // the forked checkpoint: close the panel and reset input/streaming.
      setStreamingId(null)
      setShowTimeTravel(false)
    } catch (err) {
      setError(err.message)
    } finally {
      setLoading(false)
    }
  }

  const handleNewConversation = async () => {
    if (loading) return
    setLoading(true)
    setError(null)
    try {
      const thread = await client.threads.create()
      setThreadId(thread.thread_id)
      setActiveCheckpointId(null)
      setHistory([])
      setThreads((prev) => [{ ...thread, preview: 'New conversation' }, ...prev])
      setMessages([GREETING])
      setStreamingId(null)
      setInput('')
    } catch (err) {
      setError(err.message)
    } finally {
      setLoading(false)
    }
  }

  const handleDeleteThread = async (id) => {
    if (loading) return
    setLoading(true)
    setError(null)
    try {
      await client.threads.delete(id)
      setThreads((prev) => prev.filter((t) => t.thread_id !== id))
      setForkedCheckpoints((prev) => {
        const next = { ...prev }
        delete next[id]
        return next
      })
      // If the deleted conversation is currently open, reset the chat view.
      if (id === threadId) {
        setThreadId(null)
        setActiveCheckpointId(null)
        setMessages([GREETING])
        setStreamingId(null)
        setInput('')
        setShowTimeTravel(false)
        setHistory([])
      }
    } catch (err) {
      setError(err.message)
    } finally {
      setLoading(false)
    }
  }

  const handleStopStream = async (e) => {
    e.preventDefault()
    // Cancel the run on the server so the agent stops generating; the
    // server closing the run also ends the SSE stream and the loop below.
    const activeRun = activeRunRef.current
    if (activeRun) {
      activeRunRef.current = null
      try {
        await client.runs.cancel(activeRun.threadId, activeRun.runId)
      } catch (err) {
        console.error('Failed to cancel run:', err)
      }
    }
  }


  const handleSubmit = async (e) => {
    e.preventDefault()
    const text = input.trim()
    if (!text || loading) return

    // Client-side guard: don't even create a run when the limit is known to
    // be reached (the server enforces it too, returning 403).
    if (usage?.limit_reached) {
      setError(
        `Monthly token limit reached (${usage.tokens_used}/${usage.token_limit} tokens used). ` +
        'Your allowance refreshes on the first day of next month.'
      )
      return
    }

    const humanMessage = { role: 'human', content: text, id: messages.length }

    const updatedMessages = [...messages, humanMessage]
    setMessages(updatedMessages)
    setInput('')
    setLoading(true)
    setError(null)

    updatedMessages.forEach(
      message => console.log('Message: <' + message.role + '> ' + message.content)
    )

    try {
      const assistantId = 'agent'
      let currentThreadId = threadId
      if (!currentThreadId) {
        const thread = await client.threads.create()
        currentThreadId = thread.thread_id
        setThreadId(currentThreadId)
        setThreads((prev) => [{ ...thread, preview: text.slice(0, 50) + (text.length > 50 ? '…' : '') }, ...prev])
      } else {
        // Update preview with latest user message for existing thread
        setThreads((prev) =>
          prev.map((t) =>
            t.thread_id === currentThreadId
              ? { ...t, preview: text.slice(0, 50) + (text.length > 50 ? '…' : '') }
              : t
          )
        )
      }

      let stream = null
      if (activeCheckpointId) {
        console.log('Sending message with checkpointId:', activeCheckpointId)

        const newConfig = await client.threads.updateState(
          currentThreadId,
          {
            values: { messages: [humanMessage] },
            checkpointId: activeCheckpointId,
          }
        )

        stream = client.runs.stream(
          currentThreadId,
          assistantId,
          {
            input: null,
            streamMode: ['events'],
            checkpointId: newConfig.configurable?.checkpoint_id,
            onRunCreated: ({ run_id }) => {
              activeRunRef.current = { threadId: currentThreadId, runId: run_id }
            },
            onDisconnect: "cancel",
          }
        )
      } else {
        console.log('Sending message without checkpointId')
        stream = client.runs.stream(
          currentThreadId,
          assistantId,
          {
            input: { messages: [humanMessage] },
            streamMode: ['events'],
            onRunCreated: ({ run_id }) => {
              activeRunRef.current = { threadId: currentThreadId, runId: run_id }
            },
            onDisconnect: "cancel",
          },
        )
      }

      let hasStartedStreaming = false
      // Depth of nested tool execution. While > 0 we are inside a subagent
      // (navigate/search) tool call, and any chat-model chunks belong to the
      // subagent — NOT to the supervisor's own answer. Tracking this prevents
      // subagent output from being interleaved/duplicated into the final text.
      let toolDepth = 0

      for await (const event of stream) {
        console.log('Event:', event) // Debug log to see what events we're getting

        // Some stream events (e.g. metadata, updates, end) carry a null or
        // shapeless data payload. Skip them before reading the inner event
        // name, otherwise `event.data.event` throws and aborts the stream.
        if (!event.data || typeof event.data !== 'object') continue

        const eventName = event.data.event

        if (eventName === 'on_tool_start') {
          toolDepth += 1
          continue
        }
        if (eventName === 'on_tool_end') {
          toolDepth = Math.max(0, toolDepth - 1)
          continue
        }

        if (eventName !== 'on_chat_model_stream') {
          continue
        }

        // Ignore chunks produced by subagents running inside a tool call.
        if (toolDepth > 0) {
          continue
        }

        const chunk = event.data.data?.chunk
        if (!chunk) continue

        // Skip chunks that are part of a tool call. Qwen emits raw tool-call
        // markup (e.g. `<function=...>` / `</function>`) as content tokens;
        // these carry tool-call metadata and must not be treated as
        // user-facing text.
        if (chunk?.additional_kwargs?.tool_calls) {
          continue
        }

        const value = chunk?.content
        if (value === undefined || value === null || value === '') {
          continue
        }

        const text = String(value)

        // Defensive filter: drop any residual tool-markup fragments that
        // leaked into plain content (the source of the `se}</function>`
        // artifacts).
        if (
          text.includes('<function') ||
          text.includes('</function') ||
          text.includes('<tool') ||
          text.includes('</tool')
        ) {
          continue
        }

        hasStartedStreaming = true

        setMessages((messages) => {
          // Find the last message if it's from the agent and currently streaming
          const lastMessageIndex = messages.length - 1
          if (lastMessageIndex >= 0 && messages[lastMessageIndex].role === 'assistant') {
            // Update the last message with accumulated content
            const updatedMessages = [...messages]
            updatedMessages[lastMessageIndex] = {
              ...updatedMessages[lastMessageIndex],
              content: updatedMessages[lastMessageIndex].content + text
            }
            return updatedMessages
          } else {
            // Create the first assistant message
            setStreamingId(messages.length)
            return [...messages, { role: 'assistant', content: text, id: messages.length }]
          }
        })
      }

      // The stream ended (completed or stopped); forget the run handle.
      activeRunRef.current = null
      setStreamingId(null)
      // The run created a new checkpoint at the thread tip; clear the fork
      // point so following messages continue from the latest state, and drop
      // the thread's fork record since the fork is now part of its history.
      setActiveCheckpointId(null)
      setForkedCheckpoints((prev) => {
        const next = { ...prev }
        delete next[currentThreadId]
        return next
      })
      // Keep the time travel panel up to date with the new checkpoint.
      if (showTimeTravel) {
        refreshHistory(currentThreadId)
      }

      // If we didn't get any streaming content, show a message
      if (!hasStartedStreaming) {
        setMessages((messages) => [...messages, { role: 'assistant', content: 'There are some problems with the agent.' }])
      }

      // Refresh the usage banner now that the run consumed tokens.
      fetchUsage().then(setUsage).catch(() => {})
    } catch (err) {
      // A 403 from runs.create means the monthly token limit was reached;
      // surface the server's message and refresh the usage banner.
      if (err?.status === 403 || /token limit/i.test(err?.message || '')) {
        fetchUsage().then(setUsage).catch(() => {})
        setError(
          err.message ||
          'Monthly token limit reached. Your allowance refreshes on the first day of next month.'
        )
      } else {
        setError(err.message)
      }
    } finally {
      activeRunRef.current = null
      setLoading(false)
    }
}

  return (
    <div className={styles.layout}>
      <Sidebar
        threads={threads}
        activeThreadId={threadId}
        loading={loading}
        onSelectThread={loadThreadHistory}
        onNewChat={handleNewConversation}
        onDeleteThread={handleDeleteThread}
      />
      <div className={styles.chat}>
        <header className={styles.header}>
          <h1>ASSISTANT</h1>
          <div className={styles.headerActions}>
            <button
              type="button"
              className={styles.timeTravelToggle}
              onClick={handleToggleTimeTravel}
              disabled={loading || !threadId}
              title={threadId ? 'Browse checkpoints and fork the conversation' : 'Start a conversation first'}
              aria-pressed={showTimeTravel}
            >
              ⏱ Time travel
            </button>
            <span className={styles.userName} title={user?.username}>
              {user?.display_name || user?.username}
            </span>
            <button
              type="button"
              className={styles.logoutButton}
              onClick={logout}
              title="Sign out"
            >
              Sign out
            </button>
          </div>
        </header>

        <div className={styles.messages}>
          {usage?.limit_reached && (
            <div className={styles.usageBanner} role="alert">
              Monthly token limit reached ({usage.tokens_used}/{usage.token_limit} tokens used).
              Your allowance refreshes on the first day of next month.
            </div>
          )}
          {messages.map((msg, index) => (
            <ChatMessage key={index} role={msg.role} content={msg.content} isStreaming={msg.id === streamingId} />
          ))}
          {loading && (
            <div className={styles.loading}>
              <span className={styles.spinner} aria-hidden="true" />
              Agent is thinking…
            </div>
          )}
          {error && <div className={styles.error}>Error: {error}</div>}
          <div ref={messagesEndRef} />
        </div>

        <form className={styles.inputArea} onSubmit={handleSubmit}>
          <input
            type="text"
            value={input}
            onChange={(e) => setInput(e.target.value)}
            placeholder={
              usage?.limit_reached
                ? 'Monthly token limit reached — refreshes on the 1st of next month'
                : 'Type your message…'
            }
            disabled={loading || usage?.limit_reached}
            aria-label="Message"
          />
          {loading ? (
            <button
              type="button"
              className={styles.pauseButton}
              onClick={handleStopStream}
              title="Stop generating"
            >
              Pause
            </button>
          ) : (
            <button type="submit" disabled={!input.trim() || usage?.limit_reached}>
              Send
            </button>
          )}
        </form>
      </div>
      {showTimeTravel && (
        <TimeTravelPanel
          history={history}
          loading={historyLoading}
          activeCheckpointId={activeCheckpointId}
          onFork={handleForkFromCheckpoint}
          onClose={() => setShowTimeTravel(false)}
        />
      )}
    </div>
  )
}