package com.pablotj.ai.chat.infrastructure.ai;

import com.pablotj.ai.chat.domain.exception.AiServiceUnavailableException;
import de.kherud.llama.LlamaModel;
import de.kherud.llama.ModelParameters;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

/**
 * Manages the Llama model lifecycle and configuration.
 * Handles model loading, initialization, and cleanup.
 */
@Component
public class LlamaModelManager {

    private static final Logger logger = LoggerFactory.getLogger(LlamaModelManager.class);

    @Value("${ai.model.name}")
    private String modelName;

    @Value("${ai.model.path}")
    private String modelPath;

    @Value("${ai.model.inference.threads:8}")
    private int threads;

    @Value("${ai.model.gpu.enabled:true}")
    private boolean gpuEnabled;

    @Value("${ai.model.gpu.layers:35}")
    private int gpuLayers;

    @Value("${ai.model.gpu.main-gpu:0}")
    private int mainGpu;

    @Value("${ai.model.context.size:4096}")
    private int contextSize;

    private LlamaModel model;
    private volatile boolean modelLoaded = false;

    @PostConstruct
    public void initializeModel() {
        try {
            logger.info("Initializing Llama model: {}", modelName);

            validateModelFile();
            loadModel();

            modelLoaded = true;
            logger.info("Successfully initialized Llama model: {}", getModelInfo());

        } catch (Exception e) {
            logger.error("Failed to initialize Llama model", e);
            throw new AiServiceUnavailableException("Failed to initialize AI model: " + e.getMessage(), e);
        }
    }

    private void validateModelFile() {
        Path modelFilePath = Paths.get(modelPath);

        if (!Files.exists(modelFilePath)) {
            throw new AiServiceUnavailableException(
                    String.format("Model file not found: %s", modelFilePath.toAbsolutePath())
            );
        }

        if (!Files.isReadable(modelFilePath)) {
            throw new AiServiceUnavailableException(
                    String.format("Model file is not readable: %s", modelFilePath.toAbsolutePath())
            );
        }

        logger.debug("Model file validation successful: {}", modelFilePath.toAbsolutePath());
    }

    private void loadModel() {
        ModelParameters parameters = new ModelParameters()
                .setModel(modelPath)
                .setSeed(42)
                .setThreads(threads)
                .setMainGpu(gpuEnabled ? mainGpu : -1)
                .setGpuLayers(gpuEnabled ? gpuLayers : 0);
        //.setContextSize(contextSize);

        logger.debug("Loading model with parameters: threads={}, gpu={}, layers={}, context={}",
                threads, gpuEnabled, gpuLayers, contextSize);

        model = new LlamaModel(parameters);
    }

    public LlamaModel getModel() {
        if (!modelLoaded || model == null) {
            throw new AiServiceUnavailableException("AI model is not loaded or available");
        }
        return model;
    }

    public boolean isModelLoaded() {
        return modelLoaded && model != null;
    }

    public String getModelInfo() {
        return String.format("Llama Model: %s (GPU: %s, Layers: %d, Context: %d)",
                modelName, gpuEnabled ? "enabled" : "disabled", gpuLayers, contextSize);
    }

    @PreDestroy
    public void cleanup() {
        if (model != null) {
            logger.info("Cleaning up Llama model resources");
            try {
                model.close();
                modelLoaded = false;
                logger.info("Successfully cleaned up Llama model resources");
            } catch (Exception e) {
                logger.warn("Error during model cleanup", e);
            }
        }
    }
}
