package dev.sonpx.loyalty.agent.core.memory;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import dev.sonpx.loyalty.agent.core.config.AgentProperties;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;

import java.sql.Timestamp;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.UUID;

/**
 * JDBC-backed ChatMemory for Spring AI with PostgreSQL persistence,
 * Vietnamese token estimation, tool result truncation, and polymorphic message serialization.
 */
@Slf4j
@Component
public class JdbcChatMemory implements ChatMemory {

    private static final int MAX_TOOL_RESULT_LENGTH = 4000;

    private final JdbcTemplate jdbcTemplate;
    private final AgentProperties agentProperties;
    private final ObjectMapper objectMapper;

    public JdbcChatMemory(JdbcTemplate jdbcTemplate, AgentProperties agentProperties, ObjectMapper objectMapper) {
        this.jdbcTemplate = jdbcTemplate;
        this.agentProperties = agentProperties;
        this.objectMapper = objectMapper;
    }

    @Override
    @Transactional
    public void add(String conversationId, List<Message> messages) {
        if (messages == null || messages.isEmpty()) {
            return;
        }

        Long currentMaxSeq = jdbcTemplate.queryForObject(
                "SELECT COALESCE(MAX(sequence_num), 0) FROM chat_memory_messages WHERE conversation_id = ?",
                Long.class,
                conversationId
        );
        long seq = (currentMaxSeq != null ? currentMaxSeq : 0);

        for (Message msg : messages) {
            seq++;
            saveMessage(conversationId, seq, msg);
        }

        // Apply token-aware truncation
        truncateIfExceedsTokenLimit(conversationId);
    }

    private void saveMessage(String conversationId, long sequenceNum, Message msg) {
        String id = UUID.randomUUID().toString();
        String messageType = msg.getMessageType().name();
        String content = msg.getText();
        String metadataJson = null;

        if (msg instanceof AssistantMessage assistantMsg && assistantMsg.hasToolCalls()) {
            try {
                List<ToolCallDto> dtos = assistantMsg.getToolCalls().stream()
                        .map(tc -> new ToolCallDto(tc.id(), tc.type(), tc.name(), tc.arguments()))
                        .toList();
                metadataJson = objectMapper.writeValueAsString(Map.of("toolCalls", dtos));
            } catch (Exception e) {
                log.error("Failed to serialize AssistantMessage toolCalls for conversation: {}", conversationId, e);
            }
        } else if (msg instanceof ToolResponseMessage toolMsg) {
            try {
                List<ToolResponseDto> dtos = toolMsg.getResponses().stream()
                        .map(tr -> {
                            String resData = tr.responseData();
                            if (resData != null && resData.length() > MAX_TOOL_RESULT_LENGTH) {
                                resData = resData.substring(0, MAX_TOOL_RESULT_LENGTH) + "\n... [truncated]";
                            }
                            return new ToolResponseDto(tr.id(), tr.name(), resData);
                        })
                        .toList();
                metadataJson = objectMapper.writeValueAsString(Map.of("responses", dtos));
            } catch (Exception e) {
                log.error("Failed to serialize ToolResponseMessage responses for conversation: {}", conversationId, e);
            }
        }

        jdbcTemplate.update(
                "INSERT INTO chat_memory_messages (id, conversation_id, sequence_num, message_type, content, metadata, created_at) " +
                        "VALUES (?, ?, ?, ?, ?, ?::jsonb, ?)",
                id, conversationId, sequenceNum, messageType, content, metadataJson, Timestamp.from(Instant.now())
        );
    }

    private void truncateIfExceedsTokenLimit(String conversationId) {
        List<MessageWrapper> history = loadMessageWrappers(conversationId);
        int maxTokens = agentProperties.getMaxMemoryTokens();
        int charsPerToken = agentProperties.getCharsPerToken();

        while (history.size() > 1 && estimateTokensWrappers(history, charsPerToken) > maxTokens) {
            int removeIndex = -1;
            for (int i = 0; i < history.size(); i++) {
                if (!(history.get(i).message instanceof SystemMessage)) {
                    removeIndex = i;
                    break;
                }
            }
            if (removeIndex == -1) {
                break; // Only system messages remain
            }
            MessageWrapper toRemove = history.remove(removeIndex);
            jdbcTemplate.update("DELETE FROM chat_memory_messages WHERE id = ?", toRemove.id);
        }
    }

    private int estimateTokensWrappers(List<MessageWrapper> history, int charsPerToken) {
        int totalChars = history.stream()
                .mapToInt(w -> w.message.getText() != null ? w.message.getText().length() : 0)
                .sum();
        return totalChars / (charsPerToken > 0 ? charsPerToken : 2);
    }

    @Override
    public List<Message> get(String conversationId) {
        List<MessageWrapper> wrappers = loadMessageWrappers(conversationId);
        return wrappers.stream().map(w -> w.message).toList();
    }

    private List<MessageWrapper> loadMessageWrappers(String conversationId) {
        String sql = "SELECT id, message_type, content, metadata FROM chat_memory_messages " +
                "WHERE conversation_id = ? ORDER BY sequence_num ASC";

        return jdbcTemplate.query(sql, (rs, rowNum) -> {
            String id = rs.getString("id");
            String messageTypeStr = rs.getString("message_type");
            String content = rs.getString("content");
            String metadataJson = rs.getString("metadata");

            MessageType type = MessageType.valueOf(messageTypeStr);
            Message message = deserializeMessage(type, content, metadataJson);
            return new MessageWrapper(id, message);
        }, conversationId);
    }

    private Message deserializeMessage(MessageType type, String content, String metadataJson) {
        return switch (type) {
            case USER -> new UserMessage(content != null ? content : "");
            case SYSTEM -> new SystemMessage(content != null ? content : "");
            case ASSISTANT -> {
                if (metadataJson != null && !metadataJson.isBlank()) {
                    try {
                        Map<String, Object> map = objectMapper.readValue(metadataJson, new TypeReference<>() {});
                        if (map.containsKey("toolCalls")) {
                            List<ToolCallDto> dtos = objectMapper.convertValue(map.get("toolCalls"), new TypeReference<>() {});
                            List<AssistantMessage.ToolCall> toolCalls = dtos.stream()
                                    .map(d -> new AssistantMessage.ToolCall(d.id(), d.type(), d.name(), d.arguments()))
                                    .toList();
                            yield AssistantMessage.builder()
                                    .content(content != null ? content : "")
                                    .toolCalls(toolCalls)
                                    .build();
                        }
                    } catch (Exception e) {
                        log.error("Error deserializing AssistantMessage metadata: {}", metadataJson, e);
                    }
                }
                yield new AssistantMessage(content != null ? content : "");
            }
            case TOOL -> {
                if (metadataJson != null && !metadataJson.isBlank()) {
                    try {
                        Map<String, Object> map = objectMapper.readValue(metadataJson, new TypeReference<>() {});
                        if (map.containsKey("responses")) {
                            List<ToolResponseDto> dtos = objectMapper.convertValue(map.get("responses"), new TypeReference<>() {});
                            List<ToolResponseMessage.ToolResponse> responses = dtos.stream()
                                    .map(d -> new ToolResponseMessage.ToolResponse(d.id(), d.name(), d.responseData()))
                                    .toList();
                            yield ToolResponseMessage.builder()
                                    .responses(responses)
                                    .build();
                        }
                    } catch (Exception e) {
                        log.error("Error deserializing ToolResponseMessage metadata: {}", metadataJson, e);
                    }
                }
                yield ToolResponseMessage.builder().responses(List.of()).build();
            }
        };
    }

    @Override
    @Transactional
    public void clear(String conversationId) {
        jdbcTemplate.update("DELETE FROM chat_memory_messages WHERE conversation_id = ?", conversationId);
    }

    private record MessageWrapper(String id, Message message) {}
    private record ToolCallDto(String id, String type, String name, String arguments) {}
    private record ToolResponseDto(String id, String name, String responseData) {}
}
