jan/web/screens/Settings/Engines/LocalEngineSettings.tsx

/* eslint-disable @typescript-eslint/no-explicit-any */
import React, { useCallback, useEffect, useMemo, useState } from 'react'

import {
  DownloadEvent,
  EngineEvent,
  events,
  InferenceEngine,
} from '@janhq/core'
import { Button, ScrollArea, Badge, Select, Progress } from '@janhq/joi'

import { twMerge } from 'tailwind-merge'

import {
  useGetDefaultEngineVariant,
  useGetInstalledEngines,
  useGetLatestReleasedEngine,
  setDefaultEngineVariant,
  installEngine,
  updateEngine,
  useGetReleasedEnginesByVersion,
} from '@/hooks/useEngineManagement'

import { formatDownloadPercentage } from '@/utils/converter'

import ExtensionSetting from '../ExtensionSetting'

import DeleteEngineVariant from './DeleteEngineVariant'
const os = () => {
  switch (PLATFORM) {
    case 'win32':
      return 'windows'
    case 'linux':
      return 'linux'

    default:
      return 'mac'
  }
}

const LocalEngineSettings = ({ engine }: { engine: InferenceEngine }) => {
  const { installedEngines, mutate: mutateInstalledEngines } =
    useGetInstalledEngines(engine)
  const { defaultEngineVariant, mutate: mutateDefaultEngineVariant } =
    useGetDefaultEngineVariant(engine)
  const { latestReleasedEngine } = useGetLatestReleasedEngine(engine, os())
  const { releasedEnginesByVersion } = useGetReleasedEnginesByVersion(
    engine,
    defaultEngineVariant?.version as string,
    os()
  )
  const [installingEngines, setInstallingEngines] = useState<
    Map<string, number>
  >(new Map())

  const isEngineUpdated =
    latestReleasedEngine &&
    latestReleasedEngine.every((item) =>
      item.name.includes(
        defaultEngineVariant?.version.replace(/^v/, '') as string
      )
    )

  const availableVariants = useMemo(
    () =>
      latestReleasedEngine?.map((e) =>
        e.name.replace(
          `${defaultEngineVariant?.version.replace(/^v/, '') as string}-`,
          ''
        )
      ),
    [latestReleasedEngine, defaultEngineVariant]
  )
  const options =
    installedEngines &&
    installedEngines
      .filter((x: any) => x.version === defaultEngineVariant?.version)
      .map((x: any) => ({
        name: x.name,
        value: x.name,
      }))

  const installedEngineByVersion = installedEngines?.filter(
    (x: any) => x.version === defaultEngineVariant?.version
  )

  const [selectedVariants, setSelectedVariants] = useState(
    defaultEngineVariant?.variant
  )

  const selectedVariant = useMemo(
    () =>
      options?.map((e) => e.value).includes(selectedVariants)
        ? selectedVariants
        : undefined,
    [selectedVariants, options]
  )

  useEffect(() => {
    if (defaultEngineVariant?.variant) {
      setSelectedVariants(defaultEngineVariant.variant || '')
    }
  }, [defaultEngineVariant])

  const handleEngineUpdate = useCallback(
    (event: { id: string; type: DownloadEvent; percent: number }) => {
      mutateInstalledEngines()
      mutateDefaultEngineVariant()
      // Backward compatible support - cortex.cpp returns full variant file name
      const variant: string | undefined = event.id.includes('.tar.gz')
        ? availableVariants?.find((e) => event.id.includes(`${e}.tar.gz`))
        : availableVariants?.find((e) => event.id.includes(e))

      if (!variant) {
        console.error(
          'Variant not found for event.id:',
          event.id,
          availableVariants
        )
        return
      }

      // Clone the existing Map to ensure immutability
      setInstallingEngines((prev) => {
        const updated = new Map(prev)
        if (
          event.type === DownloadEvent.onFileDownloadError ||
          event.type === DownloadEvent.onFileDownloadStopped ||
          event.type === DownloadEvent.onFileDownloadSuccess
        ) {
          // Remove the variant from the Map if download stops/errors/succeeds
          updated.delete(variant)
        } else {
          // Update the variant with the new percentage
          updated.set(variant, event.percent)
        }
        return updated
      })
    },
    [
      mutateDefaultEngineVariant,
      mutateInstalledEngines,
      setInstallingEngines,
      availableVariants,
    ]
  )

  useEffect(() => {
    events.on(EngineEvent.OnEngineUpdate, handleEngineUpdate)
    return () => {
      events.off(EngineEvent.OnEngineUpdate, handleEngineUpdate)
    }
  }, [handleEngineUpdate])

  const handleChangeVariant = (e: string) => {
    setSelectedVariants(e)
    setDefaultEngineVariant(engine, {
      variant: e,
      version: String(defaultEngineVariant?.version),
    })
  }

  return (
    <ScrollArea className="h-full w-full">
      <div className="block w-full px-4">
        <div className="mb-3 mt-4 border-b border-[hsla(var(--app-border))] pb-4">
          <div className="flex w-full flex-col items-start justify-between sm:flex-row">
            <div className="w-full flex-shrink-0 space-y-1.5">
              <div className="flex items-center justify-between gap-x-2">
                <div>
                  <h6 className="line-clamp-1 font-semibold">Engine Version</h6>
                </div>
                <div className="flex items-center gap-x-3">
                  <Badge variant="outline" theme="secondary">
                    {defaultEngineVariant?.version}
                  </Badge>
                </div>
              </div>
            </div>
          </div>
        </div>
      </div>

      <div className="block w-full px-4">
        <div className="mb-3 mt-4 border-b border-[hsla(var(--app-border))] pb-4">
          <div className="flex w-full flex-col items-start justify-between sm:flex-row">
            <div className="w-full flex-shrink-0 space-y-1.5">
              <div className="flex items-center justify-between gap-x-2">
                <div>
                  <h6 className="line-clamp-1 font-semibold">Check Updates</h6>
                </div>
                <div className="flex items-center gap-x-3">
                  <Button
                    disabled={isEngineUpdated}
                    onClick={() => {
                      updateEngine(engine)
                    }}
                  >
                    {!isEngineUpdated ? 'Update now' : 'Updated'}
                  </Button>
                </div>
              </div>
            </div>
          </div>
        </div>
      </div>

      <div className="block w-full px-4">
        <div className="mb-3 mt-4 border-b border-[hsla(var(--app-border))] pb-4">
          <div className="flex w-full flex-col items-start justify-between sm:flex-row">
            <div className="w-full flex-shrink-0 space-y-1.5">
              <div className="flex items-center justify-between gap-x-4">
                <div>
                  <h6 className="line-clamp-1 font-semibold">
                    {engine} Backend
                  </h6>
                  <div className="mt-2 w-full font-medium leading-relaxed text-[hsla(var(--text-secondary))]">
                    <p>
                      Choose the default variant that best suited for your
                      hardware. See&nbsp;
                      <a
                        href="https://jan.ai/docs/local-engines/llama-cpp"
                        className="cursor-pointer text-blue-600 dark:text-blue-400"
                        target="_blank"
                      >
                        our guides.
                      </a>
                    </p>
                  </div>
                </div>
                <div className="flex flex-shrink-0 items-center gap-x-3">
                  <div className="flex w-full min-w-[180px]">
                    <Select
                      value={selectedVariant}
                      placeholder="Select variant"
                      onValueChange={(e) => handleChangeVariant(e)}
                      options={options}
                      block
                    />
                  </div>
                </div>
              </div>
            </div>
          </div>
        </div>
      </div>

      <div className="block w-full px-4">
        <div className="mb-3 mt-4 pb-4">
          <div className="flex w-full flex-col items-start justify-between sm:flex-row">
            <div className="w-full flex-shrink-0 ">
              <div className="flex items-center justify-between gap-x-2">
                <div>
                  <h6 className="mb-2 line-clamp-1 font-semibold">Backends</h6>
                </div>
              </div>

              <div>
                {releasedEnginesByVersion &&
                  releasedEnginesByVersion?.map((item, i) => {
                    return (
                      <div
                        key={i}
                        className={twMerge(
                          'border border-b-0 border-[hsla(var(--app-border))] bg-[hsla(var(--tertiary-bg))] p-4 first:rounded-t-lg last:rounded-b-lg last:border-b',
                          releasedEnginesByVersion?.length === 1 && 'rounded-lg'
                        )}
                      >
                        <div className="flex flex-col items-start justify-start gap-4 sm:flex-row sm:items-center sm:justify-between">
                          <div className="flex w-full gap-x-8">
                            <div className="flex h-full w-full items-center justify-between gap-2">
                              <h6
                                className={twMerge(
                                  'font-medium lg:line-clamp-1 lg:min-w-[280px] lg:max-w-[280px]',
                                  'max-w-none text-[hsla(var(--text-secondary))]'
                                )}
                              >
                                {item.name}
                              </h6>

                              {installedEngineByVersion?.some(
                                (x) => x.name === item.name
                              ) ? (
                                <DeleteEngineVariant
                                  variant={item}
                                  engine={engine}
                                />
                              ) : (
                                <>
                                  {installingEngines.has(item.name) ? (
                                    <Button variant="soft">
                                      <div className="flex items-center space-x-2">
                                        <Progress
                                          className="inline-block h-2 w-[80px]"
                                          value={
                                            formatDownloadPercentage(
                                              installingEngines.get(
                                                item.name
                                              ) ?? 0,
                                              {
                                                hidePercentage: true,
                                              }
                                            ) as number
                                          }
                                        />
                                        <span className="tabular-nums">
                                          {formatDownloadPercentage(
                                            installingEngines.get(item.name) ??
                                              0
                                          )}
                                        </span>
                                      </div>
                                    </Button>
                                  ) : (
                                    <Button
                                      variant="soft"
                                      onClick={() => {
                                        setInstallingEngines((prev) => {
                                          const updated = new Map(prev)
                                          updated.set(item.name, 0)
                                          return updated
                                        })
                                        installEngine(engine, {
                                          variant: item.name,
                                          version: String(
                                            defaultEngineVariant?.version
                                          ),
                                        }).then(() => {
                                          if (selectedVariants === '') {
                                            setSelectedVariants(item.name)
                                          }
                                        })
                                      }}
                                    >
                                      Download
                                    </Button>
                                  )}
                                </>
                              )}
                            </div>
                          </div>
                        </div>
                      </div>
                    )
                  })}
              </div>
            </div>
          </div>
        </div>
        <div className="border-b border-[hsla(var(--app-border))]" />
        <div className="flex w-full border-b border-[hsla(var(--app-border))]">
          {/* TODO: Pull settings from engine when it's supported */}
          <ExtensionSetting extensionName="@janhq/inference-cortex-extension" />
        </div>
      </div>
    </ScrollArea>
  )
}

export default LocalEngineSettings