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

import com.pablotj.ai.chat.domain.exception.AiServiceUnavailableException;
import com.pablotj.ai.chat.domain.model.ConversationMessage;
import com.pablotj.ai.chat.domain.model.MessageContent;
import com.pablotj.ai.chat.domain.service.AiConversationService;
import com.pablotj.ai.chat.infrastructure.ai.prompt.ConversationPromptBuilder;
import de.kherud.llama.InferenceParameters;
import de.kherud.llama.LlamaOutput;
import io.micrometer.core.instrument.Counter;
import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.core.instrument.Timer;
import java.util.List;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

/**
 * Llama-based implementation of the AI conversation service.
 * Provides offline AI capabilities using the Llama model.
 */
@Service
public class LlamaAiConversationService implements AiConversationService {

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

    private final LlamaModelManager modelManager;
    private final ConversationPromptBuilder promptBuilder;
    private final Timer responseTimer;
    private final Counter requestCounter;
    private final Counter errorCounter;

    @Value("${ai.model.inference.max-tokens:2048}")
    private int maxTokens;

    @Value("${ai.model.inference.temperature:0.7}")
    private float temperature;

    @Value("${ai.model.inference.top-p:0.9}")
    private float topP;

    @Value("${ai.model.inference.top-k:40}")
    private int topK;

    public LlamaAiConversationService(
            LlamaModelManager modelManager,
            ConversationPromptBuilder promptBuilder,
            MeterRegistry meterRegistry) {
        this.modelManager = modelManager;
        this.promptBuilder = promptBuilder;
        this.responseTimer = Timer.builder("ai.response.duration")
                .description("Time taken to generate AI responses")
                .register(meterRegistry);
        this.requestCounter = Counter.builder("ai.requests.total")
                .description("Total number of AI requests")
                .register(meterRegistry);
        this.errorCounter = Counter.builder("ai.errors.total")
                .description("Total number of AI errors")
                .register(meterRegistry);
    }

    @Override
    public MessageContent generateResponse(List<ConversationMessage> conversationHistory, ConversationMessage userMessage) {
        requestCounter.increment();

        try {
            return responseTimer.recordCallable(() -> {
                try {
                    logger.debug("Generating AI response for conversation with {} messages", conversationHistory.size());

                    if (!isAvailable()) {
                        throw new AiServiceUnavailableException("AI model is not available");
                    }

                    String prompt = promptBuilder.buildConversationPrompt(conversationHistory, userMessage);
                    String response = generateResponseInternal(prompt);

                    logger.debug("Successfully generated AI response with {} characters", response.length());

                    return MessageContent.of(response);

                } catch (Exception e) {
                    errorCounter.increment();
                    logger.error("Error generating AI response", e);
                    throw new AiServiceUnavailableException("Failed to generate AI response: " + e.getMessage(), e);
                }
            });
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private String generateResponseInternal(String prompt) {
        InferenceParameters parameters = new InferenceParameters(prompt)
                .setNPredict(maxTokens)
                .setTemperature(temperature)
                .setTopP(topP)
                .setTopK(topK)
                .setUseChatTemplate(false);

        StringBuilder responseBuilder = new StringBuilder();

        for (LlamaOutput output : modelManager.getModel().generate(parameters)) {
            responseBuilder.append(output.text);
        }

        return cleanResponse(responseBuilder.toString());
    }

    private String cleanResponse(String response) {
        return response
                .replace("<|end_of_turn|>", "")
                .replace("<|im_end|>", "")
                .trim();
    }

    @Override
    public boolean isAvailable() {
        try {
            return modelManager.isModelLoaded();
        } catch (Exception e) {
            logger.warn("Error checking AI service availability", e);
            return false;
        }
    }

    @Override
    public String getModelInfo() {
        return modelManager.getModelInfo();
    }
}
