/*
 *  store.ts
 *  Zyquo Cloud Web
 *
 *  Author: Simon-Pierre Boucher
 *  Mail: contact@spboucher.ai
 *
 *  The app store (Zustand): settings, keys status, conversations, streaming
 *  state, and UI state. All persistence goes through storage/; all networking
 *  through providers/. Conversation behaviors ported from the native
 *  ConversationStore (send/stream/stop/regenerate/edit-resend/auto-title),
 *  extended with web-only variants + branching.
 */

import { create } from 'zustand'
import { cheapestModel, defaultModel, findModel, pricedCost } from '../features/catalogHelpers'
import { autoTrimForContext } from '../features/conversationTools'
import { clientFor } from '../providers/registry'
import { ProviderError } from '../providers/types'
import {
  deleteConversation as dbDeleteConversation,
  loadConversations,
  saveConversation,
} from '../storage/conversations'
import { configuredProviders, getKey } from '../storage/keys'
import { loadSettings, saveSettings } from '../storage/settings'
import { isLocked } from '../storage/vault'
import type {
  AIModel,
  Attachment,
  ChatParameters,
  Conversation,
  Message,
  MessageVariant,
  Provider,
  Settings,
} from '../types'

export function uid(): string {
  return crypto.randomUUID()
}

export interface Toast {
  id: string
  text: string
  kind: 'info' | 'error' | 'success'
  /** Optional undo action (soft deletes). */
  undo?: () => void
}

interface StreamingState {
  /** Message id currently being streamed into. */
  messageID: string
  abort: AbortController
  startedAt: number
  outputChars: number
}

export interface ZyquoState {
  // Settings
  settings: Settings
  updateSettings: (patch: Partial<Settings>) => void

  // Keys (presence only — values read from storage on demand)
  keyedProviders: Provider[]
  refreshKeyedProviders: () => void
  vaultLocked: boolean
  setVaultLocked: (locked: boolean) => void

  // Conversations
  conversations: Conversation[]
  selectedID: string | null
  searchText: string
  loaded: boolean
  loadAll: () => Promise<void>
  select: (id: string | null) => void
  setSearchText: (text: string) => void
  newConversation: (opts?: {
    model?: AIModel
    systemPrompt?: string
    personaID?: string
    parameters?: ChatParameters
  }) => Conversation
  updateConversation: (id: string, patch: Partial<Conversation>) => void
  removeConversation: (id: string) => void
  restoreConversation: (conversation: Conversation) => void

  // Chat actions
  send: (conversationID: string, text: string, attachments: Attachment[], overrideModel?: AIModel) => Promise<void>
  regenerate: (conversationID: string, messageID: string, overrideModel?: AIModel) => Promise<void>
  editAndResend: (conversationID: string, messageID: string, newText: string) => Promise<void>
  deleteMessage: (conversationID: string, messageID: string) => void
  stop: (conversationID: string) => void
  continueGeneration: (conversationID: string) => Promise<void>
  branchFrom: (conversationID: string, messageID: string) => Conversation | null
  setActiveVariant: (conversationID: string, messageID: string, index: number) => void

  // Streaming
  streaming: Record<string, StreamingState>

  // UI
  /** Prompt text to preload into the next-opened conversation's input. */
  draftSeed: string | null
  setDraftSeed: (seed: string | null) => void
  sidebarOpen: boolean
  setSidebarOpen: (open: boolean) => void
  modelMenuOpen: boolean
  setModelMenuOpen: (open: boolean) => void
  settingsOpen: boolean
  settingsTab: string
  openSettings: (tab?: string) => void
  closeSettings: () => void
  paletteOpen: boolean
  setPaletteOpen: (open: boolean) => void
  compareOpen: boolean
  setCompareOpen: (open: boolean) => void
  toasts: Toast[]
  toast: (text: string, kind?: Toast['kind'], undo?: () => void) => void
  dismissToast: (id: string) => void
}

function applyThemeToDocument(settings: Settings): void {
  const root = document.documentElement
  const dark =
    settings.theme === 'dark' ||
    (settings.theme === 'system' && matchMedia('(prefers-color-scheme: dark)').matches)
  if (dark) root.setAttribute('data-theme', 'dark')
  else root.removeAttribute('data-theme')
  if (settings.accent !== 'indigo') root.setAttribute('data-accent', settings.accent)
  else root.removeAttribute('data-accent')
  root.style.setProperty('--z-chat-font-size', `${settings.chatFontSize}px`)
}

export const useStore = create<ZyquoState>((set, get) => {
  const initialSettings = loadSettings()
  applyThemeToDocument(initialSettings)

  const persist = (conversation: Conversation) => {
    void saveConversation(conversation).catch(() => {
      get().toast('Failed to save conversation locally', 'error')
    })
  }

  const mutateConversation = (
    id: string,
    mutate: (c: Conversation) => Conversation,
    save = true
  ) => {
    set((state) => {
      const conversations = state.conversations.map((c) => {
        if (c.id !== id) return c
        const next = mutate(c)
        if (save) persist(next)
        return next
      })
      return { conversations }
    })
  }

  const modelFor = (conversation: Conversation): AIModel | undefined =>
    findModel(conversation.provider, conversation.modelID)

  /** Streams a reply into a fresh assistant message appended to the conversation. */
  const generateReply = async (conversationID: string, overrideModel?: AIModel) => {
    const state = get()
    const conversation = state.conversations.find((c) => c.id === conversationID)
    if (!conversation) return
    if (state.streaming[conversationID]) state.stop(conversationID)

    const model = overrideModel ?? modelFor(conversation)
    if (!model) {
      state.toast(`Model ${conversation.modelID} not found in catalog`, 'error')
      return
    }
    const apiKey = getKey(model.provider)
    if (!apiKey) {
      state.toast(ProviderError.missingAPIKey(model.provider).message, 'error')
      get().openSettings('providers')
      return
    }
    if (overrideModel) {
      mutateConversation(conversationID, (c) => ({
        ...c,
        modelID: overrideModel.id,
        provider: overrideModel.provider,
      }))
    }

    const placeholder: Message = {
      id: uid(),
      role: 'assistant',
      text: '',
      modelID: model.id,
      provider: model.provider,
      createdAt: Date.now(),
    }
    mutateConversation(
      conversationID,
      (c) => ({ ...c, messages: [...c.messages, placeholder], updatedAt: Date.now() }),
      false
    )

    const abort = new AbortController()
    set((s) => ({
      streaming: {
        ...s.streaming,
        [conversationID]: {
          messageID: placeholder.id,
          abort,
          startedAt: Date.now(),
          outputChars: 0,
        },
      },
    }))

    const source = get().conversations.find((c) => c.id === conversationID)
    if (!source) return
    // Request context: all messages except the placeholder and errored turns,
    // auto-trimmed (oldest first, system prompt kept) when near the window.
    const { messages: trimmedMessages, trimmed } = autoTrimForContext({
      ...source,
      messages: source.messages.filter((m) => m.id !== placeholder.id),
    })
    if (trimmed > 0) {
      get().toast(`Context nearly full — trimmed ${trimmed} oldest turn${trimmed > 1 ? 's' : ''}`)
    }
    const requestMessages = trimmedMessages
    const proxy = get().settings.proxyBaseURLs[model.provider]

    const patchMessage = (patch: (m: Message) => Message, save = false) => {
      mutateConversation(
        conversationID,
        (c) => ({
          ...c,
          messages: c.messages.map((m) => (m.id === placeholder.id ? patch(m) : m)),
        }),
        save
      )
    }

    try {
      const client = clientFor(model)
      const events = client.streamChat(
        {
          model,
          ...(source.systemPrompt ? { systemPrompt: source.systemPrompt } : {}),
          messages: requestMessages,
          parameters: source.parameters,
          stream: true,
          ...(proxy ? { baseURLOverride: proxy } : {}),
          ...(model.customBaseURL ? { baseURLOverride: model.customBaseURL } : {}),
        },
        apiKey,
        abort.signal
      )
      for await (const event of events) {
        switch (event.type) {
          case 'textDelta':
            set((s) => {
              const st = s.streaming[conversationID]
              return st
                ? {
                    streaming: {
                      ...s.streaming,
                      [conversationID]: { ...st, outputChars: st.outputChars + event.text.length },
                    },
                  }
                : {}
            })
            patchMessage((m) => ({ ...m, text: m.text + event.text }))
            break
          case 'reasoningDelta':
            patchMessage((m) => ({ ...m, reasoning: (m.reasoning ?? '') + event.text }))
            break
          case 'citations':
            patchMessage((m) => ({ ...m, citations: event.citations }))
            break
          case 'usage':
            patchMessage((m) => ({
              ...m,
              usage: event.usage,
              ...(model.pricing
                ? { estimatedCost: pricedCost(model, event.usage) }
                : {}),
            }))
            break
          case 'finished':
            break
        }
      }
      patchMessage((m) => m, true) // final persist
      void autoTitleIfNeeded(conversationID)
    } catch (err) {
      const cancelled = err instanceof ProviderError && err.kind === 'cancelled'
      if (!cancelled) {
        const text = err instanceof Error ? err.message : String(err)
        patchMessage((m) => ({ ...m, errorText: text }), true)
        get().toast(text, 'error')
      } else {
        patchMessage((m) => m, true)
      }
    } finally {
      set((s) => {
        const streaming = { ...s.streaming }
        delete streaming[conversationID]
        return { streaming }
      })
    }
  }

  /** Native auto-title behavior: after the first completed exchange. */
  const autoTitleIfNeeded = async (conversationID: string) => {
    const conversation = get().conversations.find((c) => c.id === conversationID)
    if (!conversation || !conversation.hasAutoTitle) return
    const assistantTurns = conversation.messages.filter(
      (m) => m.role === 'assistant' && m.text.trim() !== ''
    )
    if (assistantTurns.length !== 1) return
    const model = cheapestModel(conversation.provider)
    const apiKey = model ? getKey(model.provider) : undefined
    if (!model || !apiKey) return
    const firstUser = conversation.messages.find((m) => m.role === 'user')?.text ?? ''
    const lastAssistant = assistantTurns[assistantTurns.length - 1]?.text ?? ''
    try {
      const result = await clientFor(model).complete(
        {
          model,
          messages: [
            {
              id: uid(),
              role: 'user',
              text:
                'Write a title of at most 5 words for this conversation. Reply with the title only, no quotes.\n' +
                `User: ${firstUser.slice(0, 500)}\nAssistant: ${lastAssistant.slice(0, 500)}`,
              createdAt: Date.now(),
            },
          ],
          parameters: { maxTokens: 24 },
          stream: false,
        },
        apiKey
      )
      const title = result.text
        .trim()
        .replace(/^["'“”]+|["'“”]+$/g, '')
        .slice(0, 60)
      const current = get().conversations.find((c) => c.id === conversationID)
      if (title !== '' && current?.hasAutoTitle) {
        mutateConversation(conversationID, (c) => ({ ...c, title }))
      }
    } catch {
      // Silent failure — title stays "New Chat".
    }
  }

  return {
    settings: initialSettings,
    updateSettings: (patch) => {
      const settings = { ...get().settings, ...patch }
      saveSettings(settings)
      applyThemeToDocument(settings)
      set({ settings })
    },

    keyedProviders: [],
    refreshKeyedProviders: () => set({ keyedProviders: configuredProviders() }),
    vaultLocked: isLocked(),
    setVaultLocked: (locked) => set({ vaultLocked: locked }),

    conversations: [],
    selectedID: null,
    searchText: '',
    loaded: false,
    loadAll: async () => {
      const conversations = await loadConversations()
      set({
        conversations,
        loaded: true,
        selectedID: conversations[0]?.id ?? null,
        keyedProviders: configuredProviders(),
      })
    },
    select: (id) => set({ selectedID: id, sidebarOpen: false }),
    setSearchText: (text) => set({ searchText: text }),

    newConversation: (opts = {}) => {
      const settings = get().settings
      const model =
        opts.model ??
        findModel(settings.defaultProvider, settings.defaultModelID) ??
        defaultModel()
      const systemPrompt = opts.systemPrompt ?? settings.defaultSystemPrompt
      const conversation: Conversation = {
        id: uid(),
        title: 'New Chat',
        createdAt: Date.now(),
        updatedAt: Date.now(),
        modelID: model.id,
        provider: model.provider,
        ...(systemPrompt !== '' ? { systemPrompt } : {}),
        parameters: opts.parameters ?? { ...settings.defaultParameters },
        messages: [],
        pinned: false,
        hasAutoTitle: true,
        ...(opts.personaID ? { personaID: opts.personaID } : {}),
      }
      set((s) => ({
        conversations: [conversation, ...s.conversations],
        selectedID: conversation.id,
        sidebarOpen: false,
      }))
      persist(conversation)
      return conversation
    },

    updateConversation: (id, patch) => {
      mutateConversation(id, (c) => ({ ...c, ...patch, updatedAt: Date.now() }))
    },

    removeConversation: (id) => {
      const conversation = get().conversations.find((c) => c.id === id)
      get().stop(id)
      set((s) => {
        const conversations = s.conversations.filter((c) => c.id !== id)
        return {
          conversations,
          selectedID: s.selectedID === id ? (conversations[0]?.id ?? null) : s.selectedID,
        }
      })
      void dbDeleteConversation(id)
      if (conversation) {
        get().toast(`Deleted “${conversation.title}”`, 'info', () =>
          get().restoreConversation(conversation)
        )
      }
    },

    restoreConversation: (conversation) => {
      set((s) => ({
        conversations: [conversation, ...s.conversations].sort(
          (a, b) => b.updatedAt - a.updatedAt
        ),
        selectedID: conversation.id,
      }))
      persist(conversation)
    },

    send: async (conversationID, text, attachments, overrideModel) => {
      const conversation = get().conversations.find((c) => c.id === conversationID)
      if (!conversation) return
      const model = overrideModel ?? findModel(conversation.provider, conversation.modelID)
      const userMessage: Message = {
        id: uid(),
        role: 'user',
        text,
        ...(attachments.length > 0 ? { attachments } : {}),
        ...(model ? { modelID: model.id, provider: model.provider } : {}),
        createdAt: Date.now(),
      }
      mutateConversation(conversationID, (c) => ({
        ...c,
        messages: [...c.messages, userMessage],
        updatedAt: Date.now(),
      }))
      await generateReply(conversationID, overrideModel)
    },

    regenerate: async (conversationID, messageID, overrideModel) => {
      const conversation = get().conversations.find((c) => c.id === conversationID)
      if (!conversation) return
      const index = conversation.messages.findIndex((m) => m.id === messageID)
      const message = conversation.messages[index]
      if (!message || message.role !== 'assistant') return
      // Web improvement over native: keep the old answer as a variant instead
      // of dropping it, unless it was empty/errored.
      const keepVariant = message.text.trim() !== '' && !message.errorText
      if (index === conversation.messages.length - 1) {
        mutateConversation(conversationID, (c) => {
          const messages = c.messages.slice(0, -1)
          return { ...c, messages }
        })
        const old: MessageVariant | null = keepVariant
          ? {
              text: message.text,
              ...(message.reasoning ? { reasoning: message.reasoning } : {}),
              ...(message.citations ? { citations: message.citations } : {}),
              ...(message.modelID ? { modelID: message.modelID } : {}),
              ...(message.provider ? { provider: message.provider } : {}),
              ...(message.usage ? { usage: message.usage } : {}),
              ...(message.estimatedCost !== undefined
                ? { estimatedCost: message.estimatedCost }
                : {}),
              createdAt: message.createdAt,
            }
          : null
        await generateReply(conversationID, overrideModel)
        if (old) {
          // Attach the previous answer as a variant of the fresh message.
          const fresh = get().conversations.find((c) => c.id === conversationID)
          const last = fresh?.messages[fresh.messages.length - 1]
          if (last && last.role === 'assistant') {
            mutateConversation(conversationID, (c) => ({
              ...c,
              messages: c.messages.map((m) =>
                m.id === last.id
                  ? {
                      ...m,
                      variants: [...(message.variants ?? []), old],
                    }
                  : m
              ),
            }))
          }
        }
      }
    },

    editAndResend: async (conversationID, messageID, newText) => {
      const conversation = get().conversations.find((c) => c.id === conversationID)
      if (!conversation) return
      const index = conversation.messages.findIndex((m) => m.id === messageID)
      if (index === -1) return
      mutateConversation(conversationID, (c) => ({
        ...c,
        messages: c.messages
          .slice(0, index + 1)
          .map((m) => (m.id === messageID ? { ...m, text: newText } : m)),
        updatedAt: Date.now(),
      }))
      await generateReply(conversationID)
    },

    deleteMessage: (conversationID, messageID) => {
      mutateConversation(conversationID, (c) => ({
        ...c,
        messages: c.messages.filter((m) => m.id !== messageID),
        updatedAt: Date.now(),
      }))
    },

    stop: (conversationID) => {
      const st = get().streaming[conversationID]
      if (st) st.abort.abort()
    },

    continueGeneration: async (conversationID) => {
      const conversation = get().conversations.find((c) => c.id === conversationID)
      if (!conversation) return
      const continueMessage: Message = {
        id: uid(),
        role: 'user',
        text: 'Continue exactly where you left off.',
        createdAt: Date.now(),
      }
      mutateConversation(conversationID, (c) => ({
        ...c,
        messages: [...c.messages, continueMessage],
        updatedAt: Date.now(),
      }))
      await generateReply(conversationID)
    },

    branchFrom: (conversationID, messageID) => {
      const conversation = get().conversations.find((c) => c.id === conversationID)
      if (!conversation) return null
      const index = conversation.messages.findIndex((m) => m.id === messageID)
      if (index === -1) return null
      const branch: Conversation = {
        ...conversation,
        id: uid(),
        title: `${conversation.title} (branch)`,
        createdAt: Date.now(),
        updatedAt: Date.now(),
        messages: conversation.messages.slice(0, index + 1).map((m) => ({ ...m, id: uid() })),
        pinned: false,
        hasAutoTitle: false,
        branchedFrom: { conversationID, messageID },
      }
      set((s) => ({
        conversations: [branch, ...s.conversations],
        selectedID: branch.id,
      }))
      persist(branch)
      return branch
    },

    setActiveVariant: (conversationID, messageID, index) => {
      mutateConversation(conversationID, (c) => ({
        ...c,
        messages: c.messages.map((m) => {
          if (m.id !== messageID || !m.variants) return m
          const variants = m.variants
          const target = index === -1 ? null : variants[index]
          if (index === -1 || !target) return { ...m, activeVariant: undefined }
          // Swap: current content goes into the variant slot, variant becomes active.
          return { ...m, activeVariant: index }
        }),
      }))
    },

    streaming: {},

    draftSeed: null,
    setDraftSeed: (seed) => set({ draftSeed: seed }),
    sidebarOpen: false,
    setSidebarOpen: (open) => set({ sidebarOpen: open }),
    modelMenuOpen: false,
    setModelMenuOpen: (open) => set({ modelMenuOpen: open }),
    settingsOpen: false,
    settingsTab: 'providers',
    openSettings: (tab = 'providers') => set({ settingsOpen: true, settingsTab: tab }),
    closeSettings: () => set({ settingsOpen: false }),
    paletteOpen: false,
    setPaletteOpen: (open) => set({ paletteOpen: open }),
    compareOpen: false,
    setCompareOpen: (open) => set({ compareOpen: open }),

    toasts: [],
    toast: (text, kind = 'info', undo) => {
      const id = uid()
      set((s) => ({ toasts: [...s.toasts, { id, text, kind, ...(undo ? { undo } : {}) }] }))
      setTimeout(() => get().dismissToast(id), 6000)
    },
    dismissToast: (id) => set((s) => ({ toasts: s.toasts.filter((t) => t.id !== id) })),
  }
})