package com.hertz.modules.ai.service.impl;

import com.hertz.modules.ai.service.AiService;
import com.hertz.modules.ai.service.ConversationService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;
import com.hertz.modules.knowledge.service.KnowledgeVectorStoreManager;

@Service
@Slf4j
@RequiredArgsConstructor
public class AiServiceImpl implements AiService {

    private final OllamaChatModel chatModel;
    private final KnowledgeVectorStoreManager vectorStoreManager;
    private final ConversationService conversationService;

    private static final String SYSTEM_PROMPT = "你是赫兹官网的AI助手。赫兹是一个基于Spring Boot的权限控制管理系统框架，用于构建高性能、可扩展的应用程序。" +
            "你可以回答关于赫兹的问题，如：赫兹的架构、性能优化、微服务设计模式等。" +
            "如果用户有其他问题，也请尽力回答。";

    private static final int CHAT_MEMORY_MAX_MESSAGES = 10;

    @Override
    public String chat(String message, Double temperature, Long knowledgeBaseId, Long conversationId, Long userId) {
        long startTime = System.currentTimeMillis();
        try {
            log.info("Starting AI chat request. Message length: {}", message.length());
            
            String systemPromptContent = SYSTEM_PROMPT;
            if (knowledgeBaseId != null) {
                systemPromptContent = getRAGSystemPrompt(message, knowledgeBaseId);
            }

            OllamaOptions options = OllamaOptions.builder()
                .temperature(temperature != null ? temperature : 0.7)
                .build();

            List<Message> messages = buildPromptMessages(systemPromptContent, conversationId, userId, message);
                
            String response = chatModel.call(new Prompt(messages, options)).getResult().getOutput().getContent();
            
            log.info("AI chat request completed in {}ms", System.currentTimeMillis() - startTime);
            return response;
        } catch (Exception e) {
            log.error("AI chat request failed", e);
            throw new RuntimeException("AI service unavailable: " + e.getMessage());
        }
    }

    @Override
    public Flux<String> streamChat(String message, Double temperature, Long knowledgeBaseId, Long conversationId, Long userId) {
         OllamaOptions options = OllamaOptions.builder()
                .temperature(temperature != null ? temperature : 0.7)
                .build();

        String systemPromptContent = SYSTEM_PROMPT;
        if (knowledgeBaseId != null) {
            systemPromptContent = getRAGSystemPrompt(message, knowledgeBaseId);
        }

        List<Message> messages = buildPromptMessages(systemPromptContent, conversationId, userId, message);
                
        return chatModel.stream(new Prompt(messages, options))
                .map(response -> {
                    String content = response.getResult().getOutput().getContent();
                    return content != null ? content : "";
                })
                .doOnError(e -> log.error("Error in AI stream", e))
                .onErrorResume(e -> Flux.just(" [Error: " + e.getMessage() + "]"));
    }

    private List<Message> buildPromptMessages(String systemPromptContent, Long conversationId, Long userId, String userMessage) {
        List<com.hertz.modules.ai.entity.Message> history = List.of();
        if (conversationId != null && userId != null) {
            history = conversationService.getRecentMessages(conversationId, userId, CHAT_MEMORY_MAX_MESSAGES);
        }
        return assemblePromptMessages(systemPromptContent, history, userMessage);
    }

    List<Message> assemblePromptMessages(String systemPromptContent, List<com.hertz.modules.ai.entity.Message> history, String userMessage) {
        List<Message> promptMessages = new ArrayList<>();
        promptMessages.add(new SystemMessage(systemPromptContent));

        if (history != null && !history.isEmpty()) {
            for (com.hertz.modules.ai.entity.Message historyMessage : history) {
                Message mapped = mapStoredMessage(historyMessage);
                if (mapped != null) {
                    promptMessages.add(mapped);
                }
            }
        }

        if (!isDuplicateLastUserMessage(history, userMessage)) {
            promptMessages.add(new UserMessage(userMessage));
        }

        return promptMessages;
    }

    boolean isDuplicateLastUserMessage(List<com.hertz.modules.ai.entity.Message> history, String userMessage) {
        if (history == null || history.isEmpty()) {
            return false;
        }
        com.hertz.modules.ai.entity.Message last = history.get(history.size() - 1);
        if (last == null) {
            return false;
        }
        if (!"user".equalsIgnoreCase(last.getRole())) {
            return false;
        }
        return Objects.equals(last.getContent(), userMessage);
    }

    Message mapStoredMessage(com.hertz.modules.ai.entity.Message stored) {
        if (stored == null || stored.getContent() == null || stored.getContent().isBlank()) {
            return null;
        }
        String role = stored.getRole() == null ? "" : stored.getRole().trim().toLowerCase();
        return switch (role) {
            case "assistant" -> new AssistantMessage(stored.getContent());
            case "system" -> new SystemMessage(stored.getContent());
            default -> new UserMessage(stored.getContent());
        };
    }

    private String getRAGSystemPrompt(String message, Long knowledgeBaseId) {
        try {
            SearchRequest request = SearchRequest.builder().query(message).topK(3).build();
            List<Document> similarDocuments = vectorStoreManager.get(knowledgeBaseId).similaritySearch(request);
            log.info("RAG retrieved {} documents for kbId={}", similarDocuments.size(), knowledgeBaseId);
            
            if (similarDocuments.isEmpty()) {
                return SYSTEM_PROMPT;
            }
            
            String context = similarDocuments.stream()
                    .map(doc -> doc.getText())
                    .collect(Collectors.joining("\n\n"));
            
            return SYSTEM_PROMPT + "\n\n以下是相关背景知识：\n" + context
                + "\n\n请优先依据背景知识回答；如果背景知识不足以回答，请明确说明“知识库未检索到相关内容”，再给出你的通用建议。";
        } catch (Exception e) {
            log.error("RAG retrieval failed", e);
            return SYSTEM_PROMPT;
        }
    }
}
