package com.pablotj.ia.chat.boot.infraestructure.llm;

import com.pablotj.ia.chat.boot.domain.exception.BusinessLogicException;
import de.kherud.llama.LlamaModel;
import de.kherud.llama.ModelParameters;
import jakarta.annotation.PostConstruct;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

@Component
public class LlmModelLoader implements AutoCloseable {

    @Value(value = "${llama.model.name}")
    private String modelName;

    @Value(value = "${llama.model.gpu.enabled}")
    private boolean gpuEnabled;

    @Value(value = "${llama.model.gpu.layers}")
    private int gpuLayers;

    @Value(value = "${llama.model.tokens}")
    private int tokens;

    private LlamaModel model;

    @PostConstruct
    public void init() {
        try {
            ModelParameters params = new ModelParameters()
                    .setModel(String.format("models/%s.gguf", modelName))
                    .setSeed(42)
                    .setThreads(8)
                    .setMainGpu(gpuEnabled ? 0 : -1)
                    .setGpuLayers(gpuEnabled ? gpuLayers : -1)
                    .setPredict(tokens);
            model = new LlamaModel(params);
        } catch (Exception e) {
            throw new BusinessLogicException("Error loading model", e);
        }
    }

    public LlamaModel getModel() {
        return model;
    }

    @Override
    public void close() {
        if (model != null) model.close();
    }
}
