new-api/web/default/src/features/channels/components/dialogs/ollama-models-dialog.tsx

/*
Copyright (C) 2023-2026 QuantumNous

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as
published by the Free Software Foundation, either version 3 of the
License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.

For commercial licensing, please contact support@quantumnous.com
*/
import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
import { useQueryClient } from '@tanstack/react-query'
import { Loader2, RefreshCw, Trash2, Download, Search } from 'lucide-react'
import { useTranslation } from 'react-i18next'
import { toast } from 'sonner'
import { getCommonHeaders } from '@/lib/api'
import {
  AlertDialog,
  AlertDialogAction,
  AlertDialogCancel,
  AlertDialogContent,
  AlertDialogDescription,
  AlertDialogFooter,
  AlertDialogHeader,
  AlertDialogTitle,
} from '@/components/ui/alert-dialog'
import { Button } from '@/components/ui/button'
import { Checkbox } from '@/components/ui/checkbox'
import {
  Dialog,
  DialogContent,
  DialogDescription,
  DialogFooter,
  DialogHeader,
  DialogTitle,
} from '@/components/ui/dialog'
import { Input } from '@/components/ui/input'
import { Label } from '@/components/ui/label'
import { Progress } from '@/components/ui/progress'
import { Separator } from '@/components/ui/separator'
import {
  deleteOllamaModel,
  fetchModels as fetchModelsFromEndpoint,
  fetchUpstreamModels,
  updateChannel,
} from '../../api'
import { channelsQueryKeys, parseModelsString } from '../../lib'
import {
  formatBytes,
  normalizeOllamaModels,
  resolveOllamaBaseUrl,
  type OllamaModel,
  type PullProgress,
} from '../../lib/ollama-utils'
import { useChannels } from '../channels-provider'

const CHANNEL_TYPE_OLLAMA = 4

export function OllamaModelsDialog({
  open,
  onOpenChange,
}: {
  open: boolean
  onOpenChange: (open: boolean) => void
}) {
  const { t } = useTranslation()
  const queryClient = useQueryClient()
  const { currentRow } = useChannels()

  const isOllamaChannel = currentRow?.type === CHANNEL_TYPE_OLLAMA
  const channelId = currentRow?.id

  const [isFetching, setIsFetching] = useState(false)
  const [models, setModels] = useState<OllamaModel[]>([])
  const [selected, setSelected] = useState<string[]>([])
  const [search, setSearch] = useState('')

  const [pullName, setPullName] = useState('')
  const [isPulling, setIsPulling] = useState(false)
  const [pullProgress, setPullProgress] = useState<PullProgress | null>(null)
  const pullAbortRef = useRef<AbortController | null>(null)

  const [deleteOpen, setDeleteOpen] = useState(false)
  const [deleteTarget, setDeleteTarget] = useState<string | null>(null)
  const [isDeleting, setIsDeleting] = useState(false)

  const filteredModels = useMemo(() => {
    if (!search.trim()) return models
    const keyword = search.trim().toLowerCase()
    return models.filter((m) => m.id.toLowerCase().includes(keyword))
  }, [models, search])

  const existingModels = useMemo(
    () => parseModelsString(currentRow?.models ?? ''),
    [currentRow?.models]
  )

  useEffect(() => {
    if (!open) {
      setModels([])
      setSelected([])
      setSearch('')
      setPullName('')
      setIsPulling(false)
      setPullProgress(null)
      pullAbortRef.current?.abort()
      pullAbortRef.current = null
      return
    }

    if (open && isOllamaChannel && channelId) {
      void fetchOllamaModels()
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [open, isOllamaChannel, channelId])

  const fetchOllamaModels = useCallback(async () => {
    if (!channelId) return
    setIsFetching(true)
    try {
      let normalized: OllamaModel[] = []
      let lastErr = ''

      // 1) Prefer live fetch for Ollama if base_url is set (more accurate / supports unsaved changes)
      const baseUrl = resolveOllamaBaseUrl(currentRow ?? null)
      if (isOllamaChannel && baseUrl) {
        try {
          const payloadLive = await fetchModelsFromEndpoint({
            base_url: baseUrl,
            type: CHANNEL_TYPE_OLLAMA,
            key: typeof currentRow?.key === 'string' ? currentRow.key : '',
          })
          if (payloadLive?.success) {
            normalized = normalizeOllamaModels(payloadLive.data)
          } else if (payloadLive?.message) {
            lastErr = String(payloadLive.message)
          }
        } catch (err: unknown) {
          lastErr = err instanceof Error ? err.message : ''
        }
      }

      // 2) Fallback to server-side fetch by channelId
      if (!normalized.length) {
        const payload = await fetchUpstreamModels(Number(channelId))
        if (payload?.success) {
          normalized = normalizeOllamaModels(payload.data)
          lastErr = ''
        } else {
          lastErr = String(payload?.message || '')
        }
      }

      if (!normalized.length && lastErr) {
        toast.error(lastErr || t('Failed to fetch models'))
      }

      setModels(normalized)
      setSelected((prev) => {
        if (!prev.length) return normalized.map((m) => m.id)
        const stillAvailable = prev.filter((id) =>
          normalized.some((m) => m.id === id)
        )
        return stillAvailable.length
          ? stillAvailable
          : normalized.map((m) => m.id)
      })
    } catch (err: unknown) {
      const msg = err instanceof Error ? err.message : undefined
      toast.error(msg || t('Failed to fetch models'))
      setModels([])
    } finally {
      setIsFetching(false)
    }
  }, [channelId, currentRow, isOllamaChannel, t])

  const toggleSelected = (modelId: string, checked: boolean) => {
    setSelected((prev) => {
      if (checked) return prev.includes(modelId) ? prev : [...prev, modelId]
      return prev.filter((id) => id !== modelId)
    })
  }

  const selectAllFiltered = () => {
    setSelected((prev) => {
      const next = new Set(prev)
      filteredModels.forEach((m) => next.add(m.id))
      return Array.from(next)
    })
  }

  const clearSelection = () => setSelected([])

  const applySelection = async (mode: 'append' | 'replace') => {
    if (!currentRow) return
    if (!selected.length) {
      toast.info(t('No models selected'))
      return
    }

    const next =
      mode === 'replace'
        ? Array.from(new Set(selected))
        : Array.from(new Set([...existingModels, ...selected]))

    try {
      const res = await updateChannel(currentRow.id, { models: next.join(',') })
      if (res.success) {
        toast.success(
          mode === 'replace'
            ? t('Models updated successfully')
            : t('Models appended successfully')
        )
        queryClient.invalidateQueries({ queryKey: channelsQueryKeys.lists() })
      } else {
        toast.error(res.message || t('Failed to update models'))
      }
    } catch (err: unknown) {
      toast.error(
        err instanceof Error ? err.message : t('Failed to update models')
      )
    }
  }

  const pullModel = async () => {
    if (!channelId) return
    if (!pullName.trim()) {
      toast.error(t('Please enter model name'))
      return
    }

    if (!resolveOllamaBaseUrl(currentRow)) {
      toast.error(t('Please set Ollama API Base URL first'))
      return
    }

    pullAbortRef.current?.abort()
    const controller = new AbortController()
    pullAbortRef.current = controller

    setIsPulling(true)
    setPullProgress({ status: 'starting', completed: 0, total: 0 })

    try {
      const response = await fetch('/api/channel/ollama/pull/stream', {
        method: 'POST',
        credentials: 'include',
        headers: {
          ...getCommonHeaders(),
          Accept: 'text/event-stream',
        },
        body: JSON.stringify({
          channel_id: channelId,
          model_name: pullName.trim(),
        }),
        signal: controller.signal,
      })

      if (!response.ok || !response.body) {
        throw new Error(`HTTP ${response.status}: ${response.statusText}`)
      }

      const reader = response.body.getReader()
      const decoder = new TextDecoder()
      let buffer = ''

      while (true) {
        const { done, value } = await reader.read()
        if (done) break

        buffer += decoder.decode(value, { stream: true })
        const lines = buffer.split('\n')
        buffer = lines.pop() || ''

        for (const line of lines) {
          if (!line.startsWith('data: ')) continue
          const eventData = line.slice(6)
          if (!eventData) continue

          if (eventData === '[DONE]') {
            setIsPulling(false)
            setPullProgress(null)
            pullAbortRef.current = null
            return
          }

          try {
            const data = JSON.parse(eventData)
            if (data?.status) {
              setPullProgress(data)
            } else if (data?.error) {
              toast.error(String(data.error))
              setIsPulling(false)
              setPullProgress(null)
              pullAbortRef.current = null
              return
            } else if (data?.message) {
              toast.success(String(data.message))
              setPullName('')
              setIsPulling(false)
              setPullProgress(null)
              pullAbortRef.current = null
              await fetchOllamaModels()
              queryClient.invalidateQueries({
                queryKey: channelsQueryKeys.lists(),
              })
              return
            }
          } catch {
            // ignore malformed events
          }
        }
      }

      setIsPulling(false)
      setPullProgress(null)
      pullAbortRef.current = null
      await fetchOllamaModels()
      queryClient.invalidateQueries({ queryKey: channelsQueryKeys.lists() })
    } catch (err: unknown) {
      const isAbort =
        typeof err === 'object' &&
        err !== null &&
        'name' in err &&
        (err as { name?: unknown }).name === 'AbortError'
      if (!isAbort) {
        const msg = err instanceof Error ? err.message : ''
        toast.error(t('Model pull failed: {{msg}}', { msg }))
      }
      setIsPulling(false)
      setPullProgress(null)
      pullAbortRef.current = null
    }
  }

  const deleteModel = async (modelName: string) => {
    if (!channelId) return
    try {
      setIsDeleting(true)
      const payload = await deleteOllamaModel({
        channel_id: Number(channelId),
        model_name: modelName,
      })
      if (payload?.success) {
        toast.success(t('Model deleted'))
        await fetchOllamaModels()
        queryClient.invalidateQueries({ queryKey: channelsQueryKeys.lists() })
        setDeleteOpen(false)
        setDeleteTarget(null)
      } else {
        toast.error(payload?.message || t('Failed to delete model'))
      }
    } catch (err: unknown) {
      const msg = err instanceof Error ? err.message : undefined
      toast.error(msg || t('Failed to delete model'))
    } finally {
      setIsDeleting(false)
    }
  }

  const close = () => {
    pullAbortRef.current?.abort()
    pullAbortRef.current = null
    onOpenChange(false)
  }

  if (!open) return null

  return (
    <Dialog open={open} onOpenChange={close}>
      <DialogContent className='max-h-[90vh] overflow-hidden sm:max-w-3xl'>
        <DialogHeader>
          <DialogTitle>{t('Ollama Models')}</DialogTitle>
          <DialogDescription>
            {t('Manage local models for:')} <strong>{currentRow?.name}</strong>
          </DialogDescription>
        </DialogHeader>

        {!isOllamaChannel ? (
          <div className='text-muted-foreground py-8 text-center'>
            {t('This channel is not an Ollama channel.')}
          </div>
        ) : (
          <div className='max-h-[78vh] space-y-4 overflow-y-auto py-2 pr-1'>
            <div className='flex flex-col gap-3 sm:flex-row sm:items-end sm:justify-between'>
              <div className='flex-1 space-y-2'>
                <Label htmlFor='ollama-pull'>{t('Pull model')}</Label>
                <div className='flex gap-2'>
                  <Input
                    id='ollama-pull'
                    placeholder={t('e.g. llama3.1:8b')}
                    value={pullName}
                    onChange={(e) => setPullName(e.target.value)}
                    disabled={!channelId || isPulling}
                  />
                  <Button
                    onClick={() => void pullModel()}
                    disabled={!channelId || isPulling}
                  >
                    {isPulling ? (
                      <>
                        <Loader2 className='mr-2 h-4 w-4 animate-spin' />
                        {t('Pulling...')}
                      </>
                    ) : (
                      <>
                        <Download className='mr-2 h-4 w-4' />
                        {t('Pull')}
                      </>
                    )}
                  </Button>
                </div>
                {pullProgress && (
                  <div className='space-y-2'>
                    <div className='text-muted-foreground text-xs'>
                      {t('Status:')} {String(pullProgress.status || '-')}
                    </div>
                    <Progress
                      value={
                        typeof pullProgress.completed === 'number' &&
                        typeof pullProgress.total === 'number' &&
                        pullProgress.total > 0
                          ? Math.min(
                              100,
                              Math.round(
                                (pullProgress.completed / pullProgress.total) *
                                  100
                              )
                            )
                          : 0
                      }
                    />
                  </div>
                )}
              </div>

              <div className='flex gap-2'>
                <Button
                  variant='outline'
                  onClick={() => void fetchOllamaModels()}
                  disabled={!channelId || isFetching}
                >
                  {isFetching ? (
                    <Loader2 className='mr-2 h-4 w-4 animate-spin' />
                  ) : (
                    <RefreshCw className='mr-2 h-4 w-4' />
                  )}
                  {t('Refresh')}
                </Button>
              </div>
            </div>

            <Separator />

            <div className='space-y-3'>
              <div className='flex flex-col gap-2 sm:flex-row sm:items-center sm:justify-between'>
                <div>
                  <p className='text-sm font-medium'>{t('Local models')}</p>
                  <p className='text-muted-foreground text-xs'>
                    {t('Select models and apply to channel models list.')}
                  </p>
                </div>
                <div className='relative sm:w-72'>
                  <Search className='text-muted-foreground absolute top-1/2 left-3 h-4 w-4 -translate-y-1/2' />
                  <Input
                    placeholder={t('Search models...')}
                    value={search}
                    onChange={(e) => setSearch(e.target.value)}
                    className='pl-9'
                  />
                </div>
              </div>

              <div className='flex flex-wrap gap-2'>
                <Button variant='outline' size='sm' onClick={selectAllFiltered}>
                  {t('Select all (filtered)')}
                </Button>
                <Button variant='outline' size='sm' onClick={clearSelection}>
                  {t('Clear selection')}
                </Button>
                <Button
                  size='sm'
                  onClick={() => void applySelection('append')}
                  disabled={!selected.length}
                >
                  {t('Append to channel')}
                </Button>
                <Button
                  variant='secondary'
                  size='sm'
                  onClick={() => void applySelection('replace')}
                  disabled={!selected.length}
                >
                  {t('Replace channel models')}
                </Button>
              </div>

              <div className='overflow-hidden rounded-md border'>
                <div className='max-h-[420px] overflow-y-auto'>
                  {filteredModels.length === 0 ? (
                    <div className='text-muted-foreground p-6 text-center text-sm'>
                      {t('No models found.')}
                    </div>
                  ) : (
                    <div className='divide-y'>
                      {filteredModels.map((m) => {
                        const checked = selected.includes(m.id)
                        return (
                          <div
                            key={m.id}
                            className='flex items-center justify-between gap-3 p-3'
                          >
                            <div className='flex min-w-0 items-start gap-3'>
                              <Checkbox
                                checked={checked}
                                onCheckedChange={(v) =>
                                  toggleSelected(m.id, !!v)
                                }
                                aria-label={`Select model ${m.id}`}
                              />
                              <div className='min-w-0'>
                                <div className='truncate font-mono text-sm'>
                                  {m.id}
                                </div>
                                <div className='text-muted-foreground flex flex-wrap gap-x-3 gap-y-1 text-xs'>
                                  <span>
                                    {t('Size:')} {formatBytes(m.size)}
                                  </span>
                                  {m.digest && (
                                    <span className='truncate'>
                                      {t('Digest:')} {String(m.digest)}
                                    </span>
                                  )}
                                </div>
                              </div>
                            </div>

                            <Button
                              variant='ghost'
                              size='sm'
                              className='text-destructive hover:text-destructive'
                              onClick={() => {
                                setDeleteTarget(m.id)
                                setDeleteOpen(true)
                              }}
                              disabled={!channelId}
                            >
                              <Trash2 className='h-4 w-4' />
                            </Button>
                          </div>
                        )
                      })}
                    </div>
                  )}
                </div>
              </div>
            </div>
          </div>
        )}

        <DialogFooter>
          <Button variant='outline' onClick={close}>
            {t('Close')}
          </Button>
        </DialogFooter>
      </DialogContent>

      <AlertDialog
        open={deleteOpen}
        onOpenChange={(v) => {
          setDeleteOpen(v)
          if (!v) setDeleteTarget(null)
        }}
      >
        <AlertDialogContent>
          <AlertDialogHeader>
            <AlertDialogTitle>{t('Confirm delete')}</AlertDialogTitle>
            <AlertDialogDescription>
              {t('Delete model "{{name}}"? This cannot be undone.', {
                name: deleteTarget || '',
              })}
            </AlertDialogDescription>
          </AlertDialogHeader>
          <AlertDialogFooter>
            <AlertDialogCancel disabled={isDeleting}>
              {t('Cancel')}
            </AlertDialogCancel>
            <AlertDialogAction
              className='bg-destructive text-destructive-foreground hover:bg-destructive/90'
              disabled={isDeleting || !deleteTarget}
              onClick={() => {
                if (!deleteTarget) return
                void deleteModel(deleteTarget)
              }}
            >
              {isDeleting ? (
                <Loader2 className='mr-2 h-4 w-4 animate-spin' />
              ) : null}
              {t('Delete')}
            </AlertDialogAction>
          </AlertDialogFooter>
        </AlertDialogContent>
      </AlertDialog>
    </Dialog>
  )
}