LocalAI/core/gallery/models.go

package gallery

import (
	"context"
	"errors"
	"fmt"
	"net/url"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"time"

	"dario.cat/mergo"
	lconfig "github.com/mudler/LocalAI/core/config"
	"github.com/mudler/LocalAI/pkg/downloader"
	"github.com/mudler/LocalAI/pkg/model"
	"github.com/mudler/LocalAI/pkg/modelartifacts"
	"github.com/mudler/LocalAI/pkg/system"
	"github.com/mudler/LocalAI/pkg/utils"
	"github.com/mudler/LocalAI/pkg/vram"
	"github.com/mudler/LocalAI/pkg/xsysinfo"

	"github.com/mudler/xlog"
	"gopkg.in/yaml.v3"
)

/*

description: |
    foo
license: ""

urls:
-
-

name: "bar"

config_file: |
    # Note, name will be injected. or generated by the alias wanted by the user
    threads: 14

files:
    - filename: ""
      sha: ""
      uri: ""

prompt_templates:
    - name: ""
      content: ""

*/
// ModelConfig is the model configuration which contains all the model details
// This configuration is read from the gallery endpoint and is used to download and install the model
// It is the internal structure, separated from the request
type ModelConfig struct {
	Description     string           `yaml:"description"`
	Icon            string           `yaml:"icon"`
	License         string           `yaml:"license"`
	URLs            []string         `yaml:"urls"`
	Name            string           `yaml:"name"`
	ConfigFile      string           `yaml:"config_file"`
	Files           []File           `yaml:"files"`
	PromptTemplates []PromptTemplate `yaml:"prompt_templates"`

	// The fields below record how an entry carrying variants was resolved,
	// so a reinstall or upgrade can honor the same pin and so operators can
	// see which variant a stable model name is actually backed by.
	EntryName       string `yaml:"entry_name,omitempty"`
	ResolvedVariant string `yaml:"resolved_variant,omitempty"`
	PinnedVariant   string `yaml:"pinned_variant,omitempty"`
}

type File struct {
	Filename string `yaml:"filename" json:"filename"`
	SHA256   string `yaml:"sha256" json:"sha256"`
	URI      string `yaml:"uri" json:"uri"`
}

type PromptTemplate struct {
	Name    string `yaml:"name"`
	Content string `yaml:"content"`
}

// variantOptions turns an entry's declared variants into selectable options,
// resolving each referenced gallery entry so the selector gets the backend it
// filters on, and appends the entry itself as the base option.
//
// The base is appended rather than special-cased downstream so there is exactly
// one selection path, and so an operator can pin the entry's own name to
// decline an upgrade their hardware would otherwise take.
func variantOptions(models []*GalleryModel, entry *GalleryModel, env ResolveEnv) ([]VariantOption, error) {
	options := make([]VariantOption, 0, len(entry.Variants)+1)
	for _, v := range entry.Variants {
		// Every referenced entry is looked up here rather than lazily after
		// selection, because the backend that decides whether a variant is even
		// eligible lives on that entry, not on the variant.
		target := FindGalleryElement(models, v.Model)
		if target == nil {
			return nil, fmt.Errorf("model %q references variant %q which does not exist in any configured gallery", entry.Name, v.Model)
		}
		if target.HasVariants() {
			return nil, fmt.Errorf("model %q references variant %q which declares variants of its own; resolution is a single pass, so those would be silently ignored", entry.Name, v.Model)
		}

		option := VariantOption{Variant: v, Backend: target.Backend}
		if env.ProbeMemory != nil {
			option.ProbedMemory = env.ProbeMemory(target)
		}
		options = append(options, option)
	}

	// The base is never filtered nor ranked, so its size would change nothing
	// and probing it would spend a round trip on a discarded answer.
	return append(options, VariantOption{
		Variant: Variant{Model: entry.Name},
		Backend: entry.Backend,
		IsBase:  true,
	}), nil
}

// probeContextLength is the context size the footprint estimate is taken at.
// Selection compares whole models against whole hosts, so this only has to be a
// consistent, realistic default rather than the context a user will eventually
// configure.
const probeContextLength = 8192

// probeTimeout bounds a single entry's probe. An entry with several variants
// probes each of them, so an unreachable host has to give up quickly: an
// unknown size only costs a variant its ranking, whereas a stalled probe costs
// the user the whole install.
const probeTimeout = 5 * time.Second

// probeEntryMemory measures what a gallery entry will occupy, without
// downloading it.
//
// pkg/vram does the work and caches its results across calls: it range-fetches
// the GGUF header for a real estimate, falls back to an HTTP HEAD for the
// content length, and finally to any declared size:. A HuggingFace repo listing
// is deliberately NOT used as a further fallback, unlike the gallery UI's
// estimator: a quantization entry's urls routinely point at the base model's
// repo, and summing every weight file there would overstate one variant badly
// enough to wrongly filter it out.
//
// Zero means the size could not be determined. Callers must read that as
// unknown rather than as zero, so an unreachable network downgrades a variant's
// ranking instead of failing the install.
func probeEntryMemory(ctx context.Context, entry *GalleryModel) uint64 {
	input := vram.ModelEstimateInput{Size: entry.Size}
	for _, f := range entry.AdditionalFiles {
		if vram.IsWeightFile(f.URI) {
			input.Files = append(input.Files, vram.FileInput{URI: f.URI})
		}
	}
	if len(input.Files) == 0 && input.Size == "" {
		return 0
	}

	ctx, cancel := context.WithTimeout(ctx, probeTimeout)
	defer cancel()

	estimate, err := vram.EstimateModelMultiContext(ctx, input, []uint32{probeContextLength})
	if err != nil {
		xlog.Debug("Could not probe a model variant's size; treating it as unknown", "model", entry.Name, "error", err)
		return 0
	}
	return estimate.VRAMForContext(probeContextLength)
}

// ResolveVariant picks the gallery entry to install for a host, from the
// variants an entry declares plus the entry itself, and returns it renamed to
// the entry's name and carrying the entry's presentation metadata.
//
// Why the metadata split: the payload (url, config_file, files, overrides)
// must come from the chosen variant because that is what actually gets
// downloaded, while the presentation (name, description, icon, tags) must come
// from the entry the user asked for, so the installed model presents as that
// model rather than as one of its variants.
//
// The base entry always resolves, whatever the host has. It is a complete entry
// that every older LocalAI release installs unconditionally, so refusing it here
// would make the gallery behave worse the newer the client is. That is also why
// the base declares no memory requirement of its own: a floor that can only
// warn cannot change any outcome.
func ResolveVariant(models []*GalleryModel, entry *GalleryModel, env ResolveEnv, pin string) (*GalleryModel, Variant, error) {
	options, err := variantOptions(models, entry, env)
	if err != nil {
		return nil, Variant{}, err
	}

	selection, err := SelectVariant(options, env, pin)
	if err != nil {
		return nil, Variant{}, fmt.Errorf("resolving variant for model %q: %w", entry.Name, err)
	}
	selected := selection.Option

	if selection.FellBackToBase && len(entry.Variants) > 0 {
		xlog.Warn("No declared variant of this model fits this system; installing the entry's own build",
			"model", entry.Name, "available_memory", env.AvailableMemory, "reasons", strings.Join(selection.Reasons, "; "))
	}

	// A pin is an operator override and deliberately bypasses the hardware
	// checks, but a silent bypass makes a later out-of-memory failure
	// impossible to trace back to the pin, so it is recorded loudly here.
	if pin != "" {
		if need, known := selected.EffectiveMemory(); known && env.AvailableMemory < need {
			xlog.Warn("Pinned model variant declares more memory than this system reports; installing anyway because the pin overrides hardware resolution",
				"model", entry.Name, "variant", selected.Variant.Model, "required_memory", need, "available_memory", env.AvailableMemory)
		}
	}

	// Selecting the base means installing the entry's own payload, which is the
	// entry itself; there is no second entry to look up.
	source := entry
	if !selected.IsBase {
		// variantOptions already proved this lookup succeeds.
		source = FindGalleryElement(models, selected.Variant.Model)
	}

	resolved := *source
	resolved.Name = entry.Name
	resolved.Description = entry.Description
	resolved.Icon = entry.Icon
	resolved.License = entry.License
	// The resolved entry is a concrete install target, so it must not carry the
	// variant list any more; leaving it would let a second pass resolve the
	// already-resolved entry all over again.
	resolved.Variants = nil

	// The struct copy above is shallow, so every reference-typed field still
	// aliases the gallery's own entries. The install path mutates Overrides in
	// place (mergo merges the caller's request overrides into it) and appends to
	// the URL and tag slices, which would write the caller's request into the
	// gallery catalog itself and leak between installs the moment this path
	// reads from a cached, long-lived gallery listing. Detach them here.
	//
	// Overrides and ConfigFile are copied all the way down rather than cloned at
	// the top level only: gallery overrides are nested in practice (a
	// parameters.model map is near-universal), and mergo recurses into nested
	// maps and overwrites them in place, so a top-level clone would still hand
	// the caller the gallery's own inner maps. The slices below hold value types,
	// so cloning them once fully detaches them.
	resolved.Overrides = deepCopyStringMap(source.Overrides)
	resolved.ConfigFile = deepCopyStringMap(source.ConfigFile)
	resolved.AdditionalFiles = slices.Clone(source.AdditionalFiles)
	resolved.URLs = slices.Clone(entry.URLs)
	resolved.Tags = slices.Clone(entry.Tags)

	return &resolved, selected.Variant, nil
}

// HostResolveEnv describes this machine to variant selection.
//
// It exists so the install path and every read-only surface that reports what
// selection WOULD choose derive the host from one place. Two copies of this
// wiring would eventually disagree, and a picker that shows a different answer
// than the installer produces is worse than no picker at all.
func HostResolveEnv(ctx context.Context, systemState *system.SystemState) ResolveEnv {
	return ResolveEnv{
		AvailableMemory: availableModelMemory(systemState),
		// The whole hardware gate. IsBackendCompatible already derives
		// Darwin-only, NVIDIA-only, ROCm-only and SYCL-only from the backend
		// name, so a gallery author never has to describe hardware.
		//
		// The uri argument is deliberately empty: it exists for backend OCI
		// images, and passing a model's gallery url here would let an unrelated
		// substring in a download link decide hardware compatibility.
		BackendCompatible: func(backend string) bool {
			return systemState.IsBackendCompatible(backend, "")
		},
		ProbeMemory: func(target *GalleryModel) uint64 {
			return probeEntryMemory(ctx, target)
		},
	}
}

// noGPUDetected is what SystemState.DetectedCapability() reports when no
// usable GPU was found. The constant is unexported in pkg/system.
const noGPUDetected = "default"

// availableModelMemory reports how much memory a model may occupy on this host.
//
// With a usable GPU the model lives in VRAM. Without one it lives in system
// RAM, read through xsysinfo because that path is cgroup-aware: under
// Kubernetes the container's limit, not the node's physical RAM, is what the
// model actually gets.
//
// An unreadable RAM figure yields 0, which drops every variant with a known
// requirement and installs the base. That is the safe direction: an unknown
// host should not be talked into a larger download.
func availableModelMemory(systemState *system.SystemState) uint64 {
	if systemState.DetectedCapability() != noGPUDetected {
		return systemState.VRAM
	}
	ram, err := xsysinfo.GetSystemRAMInfo()
	if err != nil || ram == nil {
		xlog.Warn("Could not read system RAM; treating this host as having no memory budget for variant selection", "error", err)
		return 0
	}
	return ram.Total
}

// deepCopyStringMap copies a decoded YAML map so no part of the result, at any
// depth, is reachable from the original.
func deepCopyStringMap(m map[string]any) map[string]any {
	if m == nil {
		return nil
	}
	out := make(map[string]any, len(m))
	for k, v := range m {
		out[k] = deepCopyYAMLValue(v)
	}
	return out
}

// deepCopyYAMLValue recurses through the only container shapes a YAML decoder
// produces. Scalars are returned as-is because they cannot be mutated through
// the copy. map[any]any is handled as well because gopkg.in/yaml.v2 decodes
// non-string keys into it, and gallery documents are not guaranteed to have
// passed through the v3 decoder.
func deepCopyYAMLValue(v any) any {
	switch t := v.(type) {
	case map[string]any:
		return deepCopyStringMap(t)
	case map[any]any:
		out := make(map[any]any, len(t))
		for k, val := range t {
			out[k] = deepCopyYAMLValue(val)
		}
		return out
	case []any:
		out := make([]any, len(t))
		for i, val := range t {
			out[i] = deepCopyYAMLValue(val)
		}
		return out
	default:
		return v
	}
}

// Installs a model from the gallery
func InstallModelFromGallery(
	ctx context.Context,
	modelGalleries, backendGalleries []lconfig.Gallery,
	systemState *system.SystemState,
	modelLoader *model.ModelLoader,
	name string, req GalleryModel, downloadStatus func(string, string, string, float64), enforceScan, automaticallyInstallBackend, requireBackendIntegrity bool, options ...InstallOption) error {

	installOpts := applyInstallOptions(options...)

	applyModel := func(model *GalleryModel, record *ModelConfig) error {
		name = strings.ReplaceAll(name, string(os.PathSeparator), "__")

		var config ModelConfig

		if len(model.URL) > 0 {
			var err error
			config, err = GetGalleryConfigFromURLWithContext[ModelConfig](ctx, model.URL, systemState.Model.ModelsPath)
			if err != nil {
				return err
			}
			config.Description = model.Description
			config.License = model.License
		} else if len(model.ConfigFile) > 0 {
			// TODO: is this worse than using the override method with a blank cfg yaml?
			reYamlConfig, err := yaml.Marshal(model.ConfigFile)
			if err != nil {
				return err
			}
			config = ModelConfig{
				ConfigFile:  string(reYamlConfig),
				Description: model.Description,
				License:     model.License,
				Name:        model.Name,
				Files:       make([]File, 0), // Real values get added below, must be blank
				// URLs are deliberately not seeded here: they are appended once
				// below for both branches, and seeding them too would write every
				// URL twice into the persisted gallery file.
				// Prompt Template Skipped for now - I expect in this mode that they will be delivered as files.
			}
		} else {
			return fmt.Errorf("invalid gallery model %+v", model)
		}

		if record != nil {
			config.EntryName = record.EntryName
			config.ResolvedVariant = record.ResolvedVariant
			config.PinnedVariant = record.PinnedVariant
			// The variant's own name would otherwise be persisted here, which
			// contradicts the whole point of variants: the model is known by
			// the entry's stable name regardless of which variant backs it.
			config.Name = record.EntryName
		}

		installName := model.Name
		if req.Name != "" {
			installName = req.Name
		}

		// Copy the model configuration from the request schema
		config.URLs = append(config.URLs, model.URLs...)
		config.Icon = model.Icon
		config.Files = append(config.Files, req.AdditionalFiles...)
		config.Files = append(config.Files, model.AdditionalFiles...)

		// TODO model.Overrides could be merged with user overrides (not defined yet)
		if req.Overrides != nil {
			if err := mergo.Merge(&model.Overrides, req.Overrides, mergo.WithOverride); err != nil {
				return err
			}
		}

		installedModel, err := InstallModel(ctx, systemState, installName, &config, model.Overrides, downloadStatus, enforceScan, options...)
		if err != nil {
			return err
		}
		xlog.Debug("Installed model", "model", installedModel.Name)
		if automaticallyInstallBackend && installedModel.Backend != "" {
			xlog.Debug("Installing backend", "backend", installedModel.Backend)

			if err := InstallBackendFromGallery(ctx, backendGalleries, systemState, modelLoader, installedModel.Backend, downloadStatus, false, requireBackendIntegrity); err != nil {
				return err
			}
		}

		return nil
	}

	models, err := AvailableGalleryModels(modelGalleries, systemState)
	if err != nil {
		return err
	}

	model := FindGalleryElement(models, name)
	if model == nil {
		return fmt.Errorf("no model found with name %q", name)
	}

	// An entry without variants is installed directly. An entry with them is
	// still installable as-is; selection below only decides whether one of its
	// declared alternatives suits this host better.
	if !model.HasVariants() {
		// A caller who named a variant asked for something this entry cannot
		// give. Installing the entry anyway would look like success while
		// quietly ignoring the choice, so it is refused by name instead.
		if installOpts.variant != "" {
			return fmt.Errorf("%w: %q was requested but model %q declares no variants", ErrPinNotFound, installOpts.variant, model.Name)
		}
		return applyModel(model, nil)
	}

	pin := installOpts.variant
	// A previously recorded pin survives reinstalls and upgrades, so a user
	// who deliberately chose a variant is not silently re-resolved onto a
	// different one by a hardware or gallery change.
	//
	// The record is keyed by the name the model was installed under, not by the
	// gallery entry name: applyModel writes it to ._gallery_<installName>.yaml,
	// where installName is req.Name whenever the caller supplied one. Reading it
	// back under the entry's own name would miss the record for every custom-named
	// install and silently re-resolve a deliberately pinned model onto a
	// different variant, possibly swapping its backend.
	if pin == "" {
		installName := model.Name
		if req.Name != "" {
			installName = req.Name
		}
		if previous, err := GetLocalModelConfiguration(systemState.Model.ModelsPath, installName); err == nil && previous != nil {
			pin = previous.PinnedVariant
		}
	}

	env := HostResolveEnv(ctx, systemState)

	resolved, variant, err := ResolveVariant(models, model, env, pin)
	if err != nil {
		return err
	}

	xlog.Info("Resolved model to variant",
		"model", model.Name, "variant", variant.Model,
		"available_memory", env.AvailableMemory, "pinned", pin != "")

	return applyModel(resolved, &ModelConfig{
		EntryName:       model.Name,
		ResolvedVariant: variant.Model,
		PinnedVariant:   pin,
	})
}

func InstallModel(ctx context.Context, systemState *system.SystemState, nameOverride string, galleryConfig *ModelConfig, configOverrides map[string]any, downloadStatus func(string, string, string, float64), enforceScan bool, options ...InstallOption) (*lconfig.ModelConfig, error) {
	installOptions := applyInstallOptions(options...)
	config := galleryConfig
	basePath := systemState.Model.ModelsPath
	// Create base path if it doesn't exist
	err := os.MkdirAll(basePath, 0750)
	if err != nil {
		return nil, fmt.Errorf("failed to create base path: %v", err)
	}

	if len(configOverrides) > 0 {
		xlog.Debug("Config overrides", "overrides", configOverrides)
	}

	name := config.Name
	if nameOverride != "" {
		name = nameOverride
	}

	if err := utils.VerifyPath(name+".yaml", basePath); err != nil {
		return nil, err
	}

	modelConfig := lconfig.ModelConfig{}
	configFilePath := filepath.Join(basePath, name+".yaml")
	var updatedConfigYAML []byte
	writeConfig := len(configOverrides) != 0 || len(config.ConfigFile) != 0

	if writeConfig {
		// Read and update config file as map[string]interface{}
		configMap := make(map[string]any)
		err = yaml.Unmarshal([]byte(config.ConfigFile), &configMap)
		if err != nil {
			return nil, fmt.Errorf("failed to unmarshal config YAML: %v", err)
		}

		configMap["name"] = name

		if configOverrides != nil {
			if err := mergo.Merge(&configMap, configOverrides, mergo.WithOverride); err != nil {
				return nil, err
			}
		}

		// Marshal the merged map for typed validation and default application.
		updatedConfigYAML, err = yaml.Marshal(configMap)
		if err != nil {
			return nil, fmt.Errorf("failed to marshal updated config YAML: %v", err)
		}

		err = yaml.Unmarshal(updatedConfigYAML, &modelConfig)
		if err != nil {
			return nil, fmt.Errorf("failed to unmarshal updated config YAML: %v", err)
		}

		// Apply model-family-specific inference defaults so they are persisted in the config YAML.
		// Apply to the typed struct for validation, and merge into configMap for serialization
		// (configMap preserves unknown fields that ModelConfig would drop).
		lconfig.ApplyInferenceDefaults(&modelConfig, name, modelConfig.Model)

		// Merge inference defaults into configMap so they are persisted without losing unknown fields.
		if modelConfig.Temperature != nil {
			if _, exists := configMap["temperature"]; !exists {
				configMap["temperature"] = *modelConfig.Temperature
			}
		}
		if modelConfig.TopP != nil {
			if _, exists := configMap["top_p"]; !exists {
				configMap["top_p"] = *modelConfig.TopP
			}
		}
		if modelConfig.TopK != nil {
			if _, exists := configMap["top_k"]; !exists {
				configMap["top_k"] = *modelConfig.TopK
			}
		}
		if modelConfig.MinP != nil {
			if _, exists := configMap["min_p"]; !exists {
				configMap["min_p"] = *modelConfig.MinP
			}
		}
		if modelConfig.RepeatPenalty != 0 {
			if _, exists := configMap["repeat_penalty"]; !exists {
				configMap["repeat_penalty"] = modelConfig.RepeatPenalty
			}
		}
		if modelConfig.PresencePenalty != 0 {
			if _, exists := configMap["presence_penalty"]; !exists {
				configMap["presence_penalty"] = modelConfig.PresencePenalty
			}
		}

		if valid, err := modelConfig.Validate(); !valid {
			return nil, fmt.Errorf("failed to validate updated config YAML: %v", err)
		}

		if len(config.Files) == 0 {
			artifactSpec, inferred, hasArtifact, err := modelConfig.PrimaryArtifactSpec(basePath)
			if err != nil {
				return nil, err
			}
			if hasArtifact && enforceScan {
				artifact, err := artifactSpec.Normalize()
				if err != nil {
					return nil, err
				}
				probe := downloader.URI(fmt.Sprintf("%s/%s/resolve/%s/.gitattributes", strings.TrimRight(downloader.HF_ENDPOINT, "/"), artifact.Source.Repo, url.PathEscape(artifact.Source.Revision)))
				if _, err := downloader.HuggingFaceScan(probe); errors.Is(err, downloader.ErrUnsafeFilesFound) {
					return nil, err
				}
			}

			if hasArtifact {
				applied, err := bindPrimaryArtifact(ctx, basePath, &modelConfig, configMap, installOptions.materializer, artifactSpec, inferred)
				if err != nil {
					return nil, err
				}
				if applied {
					// Re-marshal from configMap after artifact binding to preserve unknown fields.
					updatedConfigYAML, err = yaml.Marshal(configMap)
					if err != nil {
						return nil, fmt.Errorf("failed to marshal config with inference defaults: %v", err)
					}
				}
			}
		}
	}

	// Download files and verify their SHA
	tasks := make([]downloader.FileTask, 0, len(config.Files))
	for i, file := range config.Files {
		// Check for cancellation before each file
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
		}

		xlog.Debug("Checking file exists and matches SHA", "filename", file.Filename)

		if err := utils.VerifyPath(file.Filename, basePath); err != nil {
			return nil, err
		}

		// Create file path
		filePath := filepath.Join(basePath, file.Filename)

		if enforceScan {
			scanResults, err := downloader.HuggingFaceScan(downloader.URI(file.URI))
			if err != nil && errors.Is(err, downloader.ErrUnsafeFilesFound) {
				xlog.Error("Contains unsafe file(s)!", "model", config.Name, "clamAV", scanResults.ClamAVInfectedFiles, "pickles", scanResults.DangerousPickles)
				return nil, err
			}
		}
		tasks = append(tasks, downloader.FileTask{
			URI:         downloader.URI(file.URI),
			Destination: filePath,
			SHA256:      file.SHA256,
			FileIndex:   i,
			TotalFiles:  len(config.Files),
		})
	}
	if err := downloader.DownloadFilesWithContext(ctx, tasks, downloadStatus); err != nil {
		return nil, err
	}

	// Write prompt template contents to separate files
	for _, template := range config.PromptTemplates {
		if err := utils.VerifyPath(template.Name+".tmpl", basePath); err != nil {
			return nil, err
		}
		// Create file path
		filePath := filepath.Join(basePath, template.Name+".tmpl")

		// Create parent directory
		err := os.MkdirAll(filepath.Dir(filePath), 0750)
		if err != nil {
			return nil, fmt.Errorf("failed to create parent directory for prompt template %q: %v", template.Name, err)
		}
		// Create and write file content
		err = os.WriteFile(filePath, []byte(template.Content), 0644)
		if err != nil {
			return nil, fmt.Errorf("failed to write prompt template %q: %v", template.Name, err)
		}

		xlog.Debug("Prompt template written", "template", template.Name)
	}

	if writeConfig {
		if len(modelConfig.Artifacts) > 0 {
			modelartifacts.ReportProgress(ctx, modelartifacts.ProgressEvent{
				Phase: modelartifacts.PhasePersisting, Artifact: modelConfig.Artifacts[0].Name,
			})
		}
		if err := writeModelConfigAtomic(configFilePath, updatedConfigYAML); err != nil {
			return nil, fmt.Errorf("failed to atomically write updated config file: %w", err)
		}

		xlog.Debug("Written config file", "file", configFilePath)
	}

	// Save the model gallery file for further reference
	modelFile := filepath.Join(basePath, galleryFileName(name))
	data, err := yaml.Marshal(config)
	if err != nil {
		return nil, err
	}

	xlog.Debug("Written gallery file", "file", modelFile)

	return &modelConfig, os.WriteFile(modelFile, data, 0644)
}

func galleryFileName(name string) string {
	return "._gallery_" + name + ".yaml"
}

// GalleryFileName returns the on-disk filename of the gallery metadata file
// for a given installed model name (e.g. "._gallery_<name>.yaml").
func GalleryFileName(name string) string {
	return galleryFileName(name)
}

func GetLocalModelConfiguration(basePath string, name string) (*ModelConfig, error) {
	name = strings.ReplaceAll(name, string(os.PathSeparator), "__")
	galleryFile := filepath.Join(basePath, galleryFileName(name))
	return ReadConfigFile[ModelConfig](galleryFile)
}

func listModelFiles(systemState *system.SystemState, name string) ([]string, error) {

	configFile := filepath.Join(systemState.Model.ModelsPath, fmt.Sprintf("%s.yaml", name))
	if err := utils.VerifyPath(configFile, systemState.Model.ModelsPath); err != nil {
		return nil, fmt.Errorf("failed to verify path %s: %w", configFile, err)
	}

	// os.PathSeparator is not allowed in model names. Replace them with "__" to avoid conflicts with file paths.
	name = strings.ReplaceAll(name, string(os.PathSeparator), "__")

	galleryFile := filepath.Join(systemState.Model.ModelsPath, galleryFileName(name))
	if err := utils.VerifyPath(galleryFile, systemState.Model.ModelsPath); err != nil {
		return nil, fmt.Errorf("failed to verify path %s: %w", galleryFile, err)
	}

	additionalFiles := []string{}
	allFiles := []string{}

	// Galleryname is the name of the model in this case
	dat, err := os.ReadFile(configFile)
	if err == nil {
		modelConfig := &lconfig.ModelConfig{}

		err = yaml.Unmarshal(dat, &modelConfig)
		if err != nil {
			return nil, err
		}
		if modelConfig.Model != "" && len(modelConfig.Artifacts) == 0 {
			additionalFiles = append(additionalFiles, modelConfig.ModelFileName())
		}

		if modelConfig.MMProj != "" {
			additionalFiles = append(additionalFiles, modelConfig.MMProjFileName())
		}
	}

	// read the model config
	galleryconfig, err := ReadConfigFile[ModelConfig](galleryFile)
	if err == nil && galleryconfig != nil {
		for _, f := range galleryconfig.Files {
			fullPath := filepath.Join(systemState.Model.ModelsPath, f.Filename)
			if err := utils.VerifyPath(fullPath, systemState.Model.ModelsPath); err != nil {
				return allFiles, fmt.Errorf("failed to verify path %s: %w", fullPath, err)
			}
			allFiles = append(allFiles, fullPath)
		}
	} else {
		xlog.Error("failed to read gallery file", "error", err, "file", configFile)
	}

	for _, f := range additionalFiles {
		fullPath := filepath.Join(filepath.Join(systemState.Model.ModelsPath, f))
		if err := utils.VerifyPath(fullPath, systemState.Model.ModelsPath); err != nil {
			return allFiles, fmt.Errorf("failed to verify path %s: %w", fullPath, err)
		}
		allFiles = append(allFiles, fullPath)
	}

	allFiles = append(allFiles, galleryFile)

	// skip duplicates
	allFiles = utils.Unique(allFiles)

	return allFiles, nil
}

func DeleteModelFromSystem(systemState *system.SystemState, name string) error {
	configFile := filepath.Join(systemState.Model.ModelsPath, fmt.Sprintf("%s.yaml", name))

	filesToRemove, err := listModelFiles(systemState, name)
	if err != nil {
		return err
	}

	allOtherFiles := []string{}
	// Get all files of all other models
	fi, err := os.ReadDir(systemState.Model.ModelsPath)
	if err != nil {
		return err
	}
	for _, f := range fi {
		if f.IsDir() {
			continue
		}
		if strings.HasPrefix(f.Name(), "._gallery_") {
			continue
		}
		if !strings.HasSuffix(f.Name(), ".yaml") && !strings.HasSuffix(f.Name(), ".yml") {
			continue
		}
		if f.Name() == fmt.Sprintf("%s.yaml", name) || f.Name() == fmt.Sprintf("%s.yml", name) {
			continue
		}

		name := strings.TrimSuffix(f.Name(), ".yaml")
		name = strings.TrimSuffix(name, ".yml")

		xlog.Debug("Checking file", "file", f.Name())
		files, err := listModelFiles(systemState, name)
		if err != nil {
			xlog.Debug("failed to list files for model", "error", err, "model", f.Name())
			continue
		}
		allOtherFiles = append(allOtherFiles, files...)
	}

	xlog.Debug("Files to remove", "files", filesToRemove)
	xlog.Debug("All other files", "files", allOtherFiles)

	// Removing files
	for _, f := range filesToRemove {
		if slices.Contains(allOtherFiles, f) {
			xlog.Debug("Skipping file because it is part of another model", "file", f)
			continue
		}
		if e := os.Remove(f); e != nil {
			xlog.Error("failed to remove file", "error", e, "file", f)
		}
	}

	return os.Remove(configFile)
}

// This is ***NEVER*** going to be perfect or finished.
// This is a BEST EFFORT function to surface known-vulnerable models to users.
func SafetyScanGalleryModels(galleries []lconfig.Gallery, systemState *system.SystemState) error {
	galleryModels, err := AvailableGalleryModels(galleries, systemState)
	if err != nil {
		return err
	}
	for _, gM := range galleryModels {
		if gM.Installed {
			err = errors.Join(err, SafetyScanGalleryModel(gM))
		}
	}
	return err
}

func SafetyScanGalleryModel(galleryModel *GalleryModel) error {
	for _, file := range galleryModel.AdditionalFiles {
		scanResults, err := downloader.HuggingFaceScan(downloader.URI(file.URI))
		if err != nil && errors.Is(err, downloader.ErrUnsafeFilesFound) {
			xlog.Error("Contains unsafe file(s)!", "model", galleryModel.Name, "clamAV", scanResults.ClamAVInfectedFiles, "pickles", scanResults.DangerousPickles)
			return err
		}
	}
	return nil
}