package org.dromara.ai.security;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.dromara.ai.domain.dto.AiChatRequest;
import org.dromara.ai.domain.dto.AiChatResponse;
import org.springframework.stereotype.Component;

import java.time.LocalDateTime;
import java.util.HashMap;
import java.util.Map;

/**
 * AI安全管理器
 * 整合内容过滤和访问控制功能，提供统一的安全服务接口
 *
 * @author pengles
 * @date 2024-01-01
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class AiSecurityManager {

    private final AiContentFilterService contentFilterService;
    private final AiAccessControlService accessControlService;

    /**
     * 安全检查结果
     */
    public static class SecurityCheckResult {
        private boolean passed;
        private String reason;
        private AiChatRequest filteredRequest;
        private AiChatResponse filteredResponse;
        private Map<String, Object> securityDetails;

        public SecurityCheckResult(boolean passed, String reason) {
            this.passed = passed;
            this.reason = reason;
            this.securityDetails = new HashMap<>();
        }

        // Getters and Setters
        public boolean isPassed() {
            return passed;
        }

        public void setPassed(boolean passed) {
            this.passed = passed;
        }

        public String getReason() {
            return reason;
        }

        public void setReason(String reason) {
            this.reason = reason;
        }

        public AiChatRequest getFilteredRequest() {
            return filteredRequest;
        }

        public void setFilteredRequest(AiChatRequest filteredRequest) {
            this.filteredRequest = filteredRequest;
        }

        public AiChatResponse getFilteredResponse() {
            return filteredResponse;
        }

        public void setFilteredResponse(AiChatResponse filteredResponse) {
            this.filteredResponse = filteredResponse;
        }

        public Map<String, Object> getSecurityDetails() {
            return securityDetails;
        }

        public void setSecurityDetails(Map<String, Object> securityDetails) {
            this.securityDetails = securityDetails;
        }
    }

    /**
     * 安全策略配置
     */
    public static class SecurityPolicy {
        private boolean enableContentFilter;
        private boolean enableAccessControl;
        private AiContentFilterService.FilterStrategy filterStrategy;
        private boolean logViolations;
        private boolean blockOnViolation;
        private int maxViolationsPerHour;
        private boolean enableIpFiltering;
        private boolean enableRateLimit;

        public SecurityPolicy() {
            this.enableContentFilter = true;
            this.enableAccessControl = true;
            this.filterStrategy = AiContentFilterService.FilterStrategy.MODERATE;
            this.logViolations = true;
            this.blockOnViolation = true;
            this.maxViolationsPerHour = 10;
            this.enableIpFiltering = true;
            this.enableRateLimit = true;
        }

        // Getters and Setters
        public boolean isEnableContentFilter() {
            return enableContentFilter;
        }

        public void setEnableContentFilter(boolean enableContentFilter) {
            this.enableContentFilter = enableContentFilter;
        }

        public boolean isEnableAccessControl() {
            return enableAccessControl;
        }

        public void setEnableAccessControl(boolean enableAccessControl) {
            this.enableAccessControl = enableAccessControl;
        }

        public AiContentFilterService.FilterStrategy getFilterStrategy() {
            return filterStrategy;
        }

        public void setFilterStrategy(AiContentFilterService.FilterStrategy filterStrategy) {
            this.filterStrategy = filterStrategy;
        }

        public boolean isLogViolations() {
            return logViolations;
        }

        public void setLogViolations(boolean logViolations) {
            this.logViolations = logViolations;
        }

        public boolean isBlockOnViolation() {
            return blockOnViolation;
        }

        public void setBlockOnViolation(boolean blockOnViolation) {
            this.blockOnViolation = blockOnViolation;
        }

        public int getMaxViolationsPerHour() {
            return maxViolationsPerHour;
        }

        public void setMaxViolationsPerHour(int maxViolationsPerHour) {
            this.maxViolationsPerHour = maxViolationsPerHour;
        }

        public boolean isEnableIpFiltering() {
            return enableIpFiltering;
        }

        public void setEnableIpFiltering(boolean enableIpFiltering) {
            this.enableIpFiltering = enableIpFiltering;
        }

        public boolean isEnableRateLimit() {
            return enableRateLimit;
        }

        public void setEnableRateLimit(boolean enableRateLimit) {
            this.enableRateLimit = enableRateLimit;
        }
    }

    // 默认安全策略
    private SecurityPolicy defaultPolicy = new SecurityPolicy();

    /**
     * 检查请求安全性
     *
     * @param request   AI聊天请求
     * @param userId    用户ID
     * @param ipAddress IP地址
     * @param apiKey    API密钥
     * @return 安全检查结果
     */
    public SecurityCheckResult checkRequestSecurity(AiChatRequest request, String userId,
                                                    String ipAddress, String apiKey) {
        return checkRequestSecurity(request, userId, ipAddress, apiKey, defaultPolicy);
    }

    /**
     * 检查请求安全性（带策略）
     *
     * @param request   AI聊天请求
     * @param userId    用户ID
     * @param ipAddress IP地址
     * @param apiKey    API密钥
     * @param policy    安全策略
     * @return 安全检查结果
     */
    public SecurityCheckResult checkRequestSecurity(AiChatRequest request, String userId,
                                                    String ipAddress, String apiKey, SecurityPolicy policy) {
        SecurityCheckResult result = new SecurityCheckResult(true, "Security check passed");
        Map<String, Object> details = result.getSecurityDetails();
        details.put("checkTime", LocalDateTime.now());
        details.put("userId", userId);
        details.put("ipAddress", ipAddress);

        try {
            // 1. API密钥验证
            if (apiKey != null && !apiKey.isEmpty()) {
                String validatedUserId = accessControlService.validateApiKey(apiKey);
                if (validatedUserId == null) {
                    result.setPassed(false);
                    result.setReason("Invalid API key");
                    details.put("apiKeyValid", false);
                    return result;
                }
                // 如果没有提供userId，使用API密钥验证的用户ID
                if (userId == null || userId.isEmpty()) {
                    userId = validatedUserId;
                }
                details.put("apiKeyValid", true);
                details.put("validatedUserId", validatedUserId);
            }

            // 2. IP访问控制
            if (policy.isEnableIpFiltering() && ipAddress != null) {
                AiAccessControlService.AccessResult ipResult = accessControlService.checkIpAccess(ipAddress);
                if (!ipResult.isAllowed()) {
                    result.setPassed(false);
                    result.setReason("IP access denied: " + ipResult.getReason());
                    details.put("ipAccessDenied", true);
                    details.put("ipDenyReason", ipResult.getReason());
                    return result;
                }
                details.put("ipAccessGranted", true);
            }

            // 3. 用户访问控制
            if (policy.isEnableAccessControl() && userId != null) {
                // 计算资源消耗（基于消息长度和复杂度）
                long resourceCost = calculateResourceCost(request);

                AiAccessControlService.AccessResult accessResult = accessControlService.checkAccess(
                        userId, "chat", resourceCost);

                if (!accessResult.isAllowed()) {
                    result.setPassed(false);
                    result.setReason("Access denied: " + accessResult.getReason());
                    details.put("accessDenied", true);
                    details.put("accessDenyReason", accessResult.getReason());
                    details.put("remainingQuota", accessResult.getRemainingQuota());
                    return result;
                }
                details.put("accessGranted", true);
                details.put("resourceCost", resourceCost);
                details.put("remainingQuota", accessResult.getRemainingQuota());
                details.putAll(accessResult.getDetails());
            }

            // 4. 内容过滤
            if (policy.isEnableContentFilter()) {
                AiChatRequest filteredRequest = filterRequestContent(request, policy.getFilterStrategy());
                if (filteredRequest == null) {
                    result.setPassed(false);
                    result.setReason("Request content violates security policy");
                    details.put("contentFiltered", true);

                    // 记录违规
                    if (policy.isLogViolations() && userId != null) {
                        accessControlService.recordViolation(userId, "CONTENT_VIOLATION",
                                "Request content filtered");
                    }
                    return result;
                }
                result.setFilteredRequest(filteredRequest);
                details.put("contentFilterPassed", true);
            } else {
                result.setFilteredRequest(request);
            }

            log.debug("Request security check passed for user: {}, IP: {}", userId, ipAddress);
            return result;

        } catch (Exception e) {
            log.error("Error during request security check", e);
            result.setPassed(false);
            result.setReason("Security check failed due to internal error");
            details.put("error", e.getMessage());
            return result;
        }
    }

    /**
     * 检查响应安全性
     *
     * @param response AI聊天响应
     * @param userId   用户ID
     * @return 安全检查结果
     */
    public SecurityCheckResult checkResponseSecurity(AiChatResponse response, String userId) {
        return checkResponseSecurity(response, userId, defaultPolicy);
    }

    /**
     * 检查响应安全性（带策略）
     *
     * @param response AI聊天响应
     * @param userId   用户ID
     * @param policy   安全策略
     * @return 安全检查结果
     */
    public SecurityCheckResult checkResponseSecurity(AiChatResponse response, String userId, SecurityPolicy policy) {
        SecurityCheckResult result = new SecurityCheckResult(true, "Response security check passed");
        Map<String, Object> details = result.getSecurityDetails();
        details.put("checkTime", LocalDateTime.now());
        details.put("userId", userId);

        try {
            // 内容过滤
            if (policy.isEnableContentFilter()) {
                AiChatResponse filteredResponse = filterResponseContent(response, policy.getFilterStrategy());
                if (filteredResponse == null) {
                    result.setPassed(false);
                    result.setReason("Response content violates security policy");
                    details.put("contentFiltered", true);

                    // 记录违规
                    if (policy.isLogViolations() && userId != null) {
                        accessControlService.recordViolation(userId, "RESPONSE_CONTENT_VIOLATION",
                                "Response content filtered");
                    }
                    return result;
                }
                result.setFilteredResponse(filteredResponse);
                details.put("contentFilterPassed", true);
            } else {
                result.setFilteredResponse(response);
            }

            log.debug("Response security check passed for user: {}", userId);
            return result;

        } catch (Exception e) {
            log.error("Error during response security check", e);
            result.setPassed(false);
            result.setReason("Response security check failed due to internal error");
            details.put("error", e.getMessage());
            return result;
        }
    }

    /**
     * 过滤请求内容
     *
     * @param request  原始请求
     * @param strategy 过滤策略
     * @return 过滤后的请求，如果违规返回null
     */
    private AiChatRequest filterRequestContent(AiChatRequest request,
                                               AiContentFilterService.FilterStrategy strategy) {
        if (request == null || request.getMessages() == null) {
            return request;
        }

        // 创建请求副本
        AiChatRequest filteredRequest = new AiChatRequest();
        filteredRequest.setModel(request.getModel());
        filteredRequest.setStream(request.getStream());
        filteredRequest.setMaxTokens(request.getMaxTokens());
        filteredRequest.setTemperature(request.getTemperature());
        filteredRequest.setTopP(request.getTopP());
        filteredRequest.setFrequencyPenalty(request.getFrequencyPenalty());
        filteredRequest.setPresencePenalty(request.getPresencePenalty());
        filteredRequest.setStop(request.getStop());
        filteredRequest.setUserId(request.getUserId());
        filteredRequest.setMessages(new java.util.ArrayList<>());

        // 过滤每条消息
        for (AiChatRequest.AiMessage message : request.getMessages()) {
            if (message.getContent() != null) {
                AiContentFilterService.FilterResult filterResult =
                        contentFilterService.filterInput(message.getContent(), strategy);

                if (!filterResult.isPassed()) {
                    log.warn("Message content filtered: violations={}, riskScore={}",
                            filterResult.getViolations(), filterResult.getRiskScore());
                    return null; // 拒绝整个请求
                }

                // 使用过滤后的内容
                AiChatRequest.AiMessage filteredMessage = new AiChatRequest.AiMessage();
                filteredMessage.setRole(message.getRole());
                filteredMessage.setContent(filterResult.getFilteredContent());
                filteredMessage.setName(message.getName());
                filteredRequest.getMessages().add(filteredMessage);
            } else {
                filteredRequest.getMessages().add(message);
            }
        }

        return filteredRequest;
    }

    /**
     * 过滤响应内容
     *
     * @param response 原始响应
     * @param strategy 过滤策略
     * @return 过滤后的响应，如果违规返回null
     */
    private AiChatResponse filterResponseContent(AiChatResponse response,
                                                 AiContentFilterService.FilterStrategy strategy) {
        if (response == null || response.getChoices() == null) {
            return response;
        }

        // 创建响应副本
        AiChatResponse filteredResponse = new AiChatResponse();
        filteredResponse.setId(response.getId());
        filteredResponse.setObject(response.getObject());
        filteredResponse.setCreated(response.getCreated());
        filteredResponse.setModel(response.getModel());
        filteredResponse.setUsage(response.getUsage());
        filteredResponse.setSuccess(response.getSuccess());
        filteredResponse.setError(response.getError());
        filteredResponse.setChoices(new java.util.ArrayList<>());

        // 过滤每个选择
        for (AiChatResponse.Choice choice : response.getChoices()) {
            if (choice.getMessage() != null && choice.getMessage().getContent() != null) {
                AiContentFilterService.FilterResult filterResult =
                        contentFilterService.filterOutput(choice.getMessage().getContent(), strategy);

                if (!filterResult.isPassed()) {
                    log.warn("Response content filtered: violations={}, riskScore={}",
                            filterResult.getViolations(), filterResult.getRiskScore());
                    return null; // 拒绝整个响应
                }

                // 使用过滤后的内容
                AiChatResponse.Choice filteredChoice = new AiChatResponse.Choice();
                filteredChoice.setIndex(choice.getIndex());
                filteredChoice.setFinishReason(choice.getFinishReason());

                AiChatRequest.AiMessage filteredMessage = new AiChatRequest.AiMessage();
                filteredMessage.setRole(choice.getMessage().getRole());
                filteredMessage.setContent(filterResult.getFilteredContent());
                filteredChoice.setMessage(filteredMessage);

                filteredResponse.getChoices().add(filteredChoice);
            } else {
                filteredResponse.getChoices().add(choice);
            }
        }

        return filteredResponse;
    }

    /**
     * 计算资源消耗
     *
     * @param request AI聊天请求
     * @return 资源消耗值
     */
    private long calculateResourceCost(AiChatRequest request) {
        if (request == null || request.getMessages() == null) {
            return 1;
        }

        long cost = 0;

        // 基于消息数量
        cost += request.getMessages().size();

        // 基于内容长度
        for (AiChatRequest.AiMessage message : request.getMessages()) {
            if (message.getContent() != null) {
                cost += message.getContent().length() / 100; // 每100字符1个单位
            }
        }

        // 基于模型复杂度
        if (request.getModel() != null) {
            String model = request.getModel().toLowerCase();
            if (model.contains("gpt-4") || model.contains("claude-3")) {
                cost *= 3; // 高级模型消耗更多
            } else if (model.contains("gpt-3.5")) {
                cost *= 2;
            }
        }

        // 基于参数设置
        if (request.getMaxTokens() != null && request.getMaxTokens() > 1000) {
            cost += request.getMaxTokens() / 1000; // 每1000 token增加1个单位
        }

        return Math.max(1, cost); // 最少1个单位
    }

    /**
     * 设置默认安全策略
     *
     * @param policy 安全策略
     */
    public void setDefaultSecurityPolicy(SecurityPolicy policy) {
        this.defaultPolicy = policy;
        log.info("Default security policy updated");
    }

    /**
     * 获取默认安全策略
     *
     * @return 安全策略
     */
    public SecurityPolicy getDefaultSecurityPolicy() {
        return defaultPolicy;
    }

    /**
     * 获取安全统计信息
     *
     * @return 统计信息
     */
    public Map<String, Object> getSecurityStats() {
        Map<String, Object> stats = new HashMap<>();

        // 访问控制统计
        stats.putAll(accessControlService.getAccessControlStats());

        // 内容过滤统计
        stats.putAll(contentFilterService.getFilterStats());

        // 安全策略信息
        Map<String, Object> policyInfo = new HashMap<>();
        policyInfo.put("contentFilterEnabled", defaultPolicy.isEnableContentFilter());
        policyInfo.put("accessControlEnabled", defaultPolicy.isEnableAccessControl());
        policyInfo.put("filterStrategy", defaultPolicy.getFilterStrategy().name());
        policyInfo.put("ipFilteringEnabled", defaultPolicy.isEnableIpFiltering());
        policyInfo.put("rateLimitEnabled", defaultPolicy.isEnableRateLimit());
        stats.put("securityPolicy", policyInfo);

        return stats;
    }

    /**
     * 清理过期数据
     */
    public void cleanup() {
        accessControlService.cleanupExpiredData();
        log.debug("Security manager cleanup completed");
    }
}