package org.dromara.ai.security;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.time.LocalDateTime;
import java.time.temporal.ChronoUnit;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

/**
 * AI访问控制服务
 * 提供用户权限管理、API访问限制和安全控制功能
 *
 * @author pengles
 * @date 2024-01-01
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class AiAccessControlService {

    /**
     * 访问控制结果
     */
    public static class AccessResult {
        private boolean allowed;
        private String reason;
        private long remainingQuota;
        private long resetTime;
        private Map<String, Object> details;

        public AccessResult(boolean allowed, String reason) {
            this.allowed = allowed;
            this.reason = reason;
            this.remainingQuota = 0;
            this.resetTime = 0;
            this.details = new HashMap<>();
        }

        // Getters and Setters
        public boolean isAllowed() {
            return allowed;
        }

        public void setAllowed(boolean allowed) {
            this.allowed = allowed;
        }

        public String getReason() {
            return reason;
        }

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

        public long getRemainingQuota() {
            return remainingQuota;
        }

        public void setRemainingQuota(long remainingQuota) {
            this.remainingQuota = remainingQuota;
        }

        public long getResetTime() {
            return resetTime;
        }

        public void setResetTime(long resetTime) {
            this.resetTime = resetTime;
        }

        public Map<String, Object> getDetails() {
            return details;
        }

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

    /**
     * 用户权限级别
     */
    public enum PermissionLevel {
        GUEST(1, "访客", 10, 100),
        BASIC(2, "基础用户", 100, 1000),
        PREMIUM(3, "高级用户", 1000, 10000),
        VIP(4, "VIP用户", 5000, 50000),
        ADMIN(5, "管理员", -1, -1);

        private final int level;
        private final String description;
        private final long dailyQuota;
        private final long monthlyQuota;

        PermissionLevel(int level, String description, long dailyQuota, long monthlyQuota) {
            this.level = level;
            this.description = description;
            this.dailyQuota = dailyQuota;
            this.monthlyQuota = monthlyQuota;
        }

        public int getLevel() {
            return level;
        }

        public String getDescription() {
            return description;
        }

        public long getDailyQuota() {
            return dailyQuota;
        }

        public long getMonthlyQuota() {
            return monthlyQuota;
        }
    }

    /**
     * 限流类型
     */
    public enum RateLimitType {
        PER_SECOND,     // 每秒限制
        PER_MINUTE,     // 每分钟限制
        PER_HOUR,       // 每小时限制
        PER_DAY,        // 每天限制
        PER_MONTH       // 每月限制
    }

    /**
     * 用户配额信息
     */
    public static class UserQuota {
        private String userId;
        private PermissionLevel permissionLevel;
        private long dailyUsed;
        private long monthlyUsed;
        private LocalDateTime lastResetDaily;
        private LocalDateTime lastResetMonthly;
        private boolean blocked;
        private LocalDateTime blockedUntil;
        private String blockReason;
        private Map<String, Long> customLimits;

        public UserQuota(String userId, PermissionLevel permissionLevel) {
            this.userId = userId;
            this.permissionLevel = permissionLevel;
            this.dailyUsed = 0;
            this.monthlyUsed = 0;
            this.lastResetDaily = LocalDateTime.now();
            this.lastResetMonthly = LocalDateTime.now();
            this.blocked = false;
            this.customLimits = new HashMap<>();
        }

        // Getters and Setters
        public String getUserId() {
            return userId;
        }

        public void setUserId(String userId) {
            this.userId = userId;
        }

        public PermissionLevel getPermissionLevel() {
            return permissionLevel;
        }

        public void setPermissionLevel(PermissionLevel permissionLevel) {
            this.permissionLevel = permissionLevel;
        }

        public long getDailyUsed() {
            return dailyUsed;
        }

        public void setDailyUsed(long dailyUsed) {
            this.dailyUsed = dailyUsed;
        }

        public long getMonthlyUsed() {
            return monthlyUsed;
        }

        public void setMonthlyUsed(long monthlyUsed) {
            this.monthlyUsed = monthlyUsed;
        }

        public LocalDateTime getLastResetDaily() {
            return lastResetDaily;
        }

        public void setLastResetDaily(LocalDateTime lastResetDaily) {
            this.lastResetDaily = lastResetDaily;
        }

        public LocalDateTime getLastResetMonthly() {
            return lastResetMonthly;
        }

        public void setLastResetMonthly(LocalDateTime lastResetMonthly) {
            this.lastResetMonthly = lastResetMonthly;
        }

        public boolean isBlocked() {
            return blocked;
        }

        public void setBlocked(boolean blocked) {
            this.blocked = blocked;
        }

        public LocalDateTime getBlockedUntil() {
            return blockedUntil;
        }

        public void setBlockedUntil(LocalDateTime blockedUntil) {
            this.blockedUntil = blockedUntil;
        }

        public String getBlockReason() {
            return blockReason;
        }

        public void setBlockReason(String blockReason) {
            this.blockReason = blockReason;
        }

        public Map<String, Long> getCustomLimits() {
            return customLimits;
        }

        public void setCustomLimits(Map<String, Long> customLimits) {
            this.customLimits = customLimits;
        }
    }

    /**
     * 限流计数器
     */
    public static class RateLimitCounter {
        private final AtomicLong count = new AtomicLong(0);
        private volatile long windowStart;
        private final long windowSizeMs;

        public RateLimitCounter(long windowSizeMs) {
            this.windowSizeMs = windowSizeMs;
            this.windowStart = System.currentTimeMillis();
        }

        public boolean tryAcquire(long limit) {
            long now = System.currentTimeMillis();

            // 检查是否需要重置窗口
            if (now - windowStart >= windowSizeMs) {
                synchronized (this) {
                    if (now - windowStart >= windowSizeMs) {
                        count.set(0);
                        windowStart = now;
                    }
                }
            }

            return count.incrementAndGet() <= limit;
        }

        public long getCurrentCount() {
            return count.get();
        }

        public long getRemainingTime() {
            long now = System.currentTimeMillis();
            return Math.max(0, windowSizeMs - (now - windowStart));
        }
    }

    // 用户配额缓存
    private final Map<String, UserQuota> userQuotaCache = new ConcurrentHashMap<>();

    // 限流计数器缓存
    private final Map<String, Map<RateLimitType, RateLimitCounter>> rateLimitCounters = new ConcurrentHashMap<>();

    // IP黑名单
    private final Set<String> ipBlacklist = ConcurrentHashMap.newKeySet();

    // 用户黑名单
    private final Set<String> userBlacklist = ConcurrentHashMap.newKeySet();

    // API密钥缓存
    private final Map<String, String> apiKeyCache = new ConcurrentHashMap<>();

    // 违规计数器
    private final Map<String, AtomicInteger> violationCounters = new ConcurrentHashMap<>();

    /**
     * 检查用户访问权限
     *
     * @param userId       用户ID
     * @param operation    操作类型
     * @param resourceCost 资源消耗
     * @return 访问控制结果
     */
    public AccessResult checkAccess(String userId, String operation, long resourceCost) {
        try {
            // 1. 检查用户是否在黑名单中
            if (userBlacklist.contains(userId)) {
                return new AccessResult(false, "User is blacklisted");
            }

            // 2. 获取用户配额信息
            UserQuota quota = getUserQuota(userId);

            // 3. 检查用户是否被临时封禁
            if (quota.isBlocked()) {
                if (quota.getBlockedUntil() != null && LocalDateTime.now().isAfter(quota.getBlockedUntil())) {
                    // 解除封禁
                    quota.setBlocked(false);
                    quota.setBlockedUntil(null);
                    quota.setBlockReason(null);
                } else {
                    return new AccessResult(false, "User is temporarily blocked: " + quota.getBlockReason());
                }
            }

            // 4. 重置配额（如果需要）
            resetQuotaIfNeeded(quota);

            // 5. 检查权限级别
            if (!hasPermission(quota.getPermissionLevel(), operation)) {
                return new AccessResult(false, "Insufficient permission level");
            }

            // 6. 检查配额限制
            AccessResult quotaResult = checkQuotaLimit(quota, resourceCost);
            if (!quotaResult.isAllowed()) {
                return quotaResult;
            }

            // 7. 检查限流
            AccessResult rateLimitResult = checkRateLimit(userId, operation);
            if (!rateLimitResult.isAllowed()) {
                return rateLimitResult;
            }

            // 8. 更新使用量
            updateUsage(quota, resourceCost);

            // 9. 构建成功结果
            AccessResult result = new AccessResult(true, "Access granted");
            result.setRemainingQuota(Math.max(0, quota.getPermissionLevel().getDailyQuota() - quota.getDailyUsed()));

            Map<String, Object> details = new HashMap<>();
            details.put("permissionLevel", quota.getPermissionLevel().name());
            details.put("dailyUsed", quota.getDailyUsed());
            details.put("monthlyUsed", quota.getMonthlyUsed());
            details.put("operation", operation);
            details.put("resourceCost", resourceCost);
            result.setDetails(details);

            log.debug("Access granted for user: {}, operation: {}, cost: {}", userId, operation, resourceCost);
            return result;

        } catch (Exception e) {
            log.error("Error checking access for user: {}", userId, e);
            return new AccessResult(false, "Internal error during access check");
        }
    }

    /**
     * 检查IP访问权限
     *
     * @param ipAddress IP地址
     * @return 访问控制结果
     */
    public AccessResult checkIpAccess(String ipAddress) {
        if (ipBlacklist.contains(ipAddress)) {
            return new AccessResult(false, "IP address is blacklisted");
        }

        // 检查IP限流
        String key = "ip:" + ipAddress;
        Map<RateLimitType, RateLimitCounter> counters = rateLimitCounters.computeIfAbsent(key,
                k -> new ConcurrentHashMap<>());

        RateLimitCounter perMinuteCounter = counters.computeIfAbsent(RateLimitType.PER_MINUTE,
                k -> new RateLimitCounter(60 * 1000)); // 1分钟窗口

        if (!perMinuteCounter.tryAcquire(100)) { // 每分钟最多100次请求
            AccessResult result = new AccessResult(false, "IP rate limit exceeded");
            result.setResetTime(perMinuteCounter.getRemainingTime());
            return result;
        }

        return new AccessResult(true, "IP access granted");
    }

    /**
     * 验证API密钥
     *
     * @param apiKey API密钥
     * @return 用户ID，如果验证失败返回null
     */
    public String validateApiKey(String apiKey) {
        if (apiKey == null || apiKey.trim().isEmpty()) {
            return null;
        }

        // 从缓存中查找
        String userId = apiKeyCache.get(apiKey);
        if (userId != null) {
            return userId;
        }

        // 这里应该从数据库中验证API密钥
        // 为了演示，我们使用简单的格式："user_" + userId
        if (apiKey.startsWith("ak_")) {
            userId = apiKey.substring(3);
            apiKeyCache.put(apiKey, userId);
            return userId;
        }

        return null;
    }

    /**
     * 记录违规行为
     *
     * @param userId        用户ID
     * @param violationType 违规类型
     * @param description   违规描述
     */
    public void recordViolation(String userId, String violationType, String description) {
        AtomicInteger counter = violationCounters.computeIfAbsent(userId, k -> new AtomicInteger(0));
        int violationCount = counter.incrementAndGet();

        log.warn("Violation recorded for user: {}, type: {}, description: {}, total: {}",
                userId, violationType, description, violationCount);

        // 根据违规次数采取措施
        if (violationCount >= 10) {
            blockUser(userId, "Excessive violations", 24); // 封禁24小时
        } else if (violationCount >= 5) {
            blockUser(userId, "Multiple violations", 1); // 封禁1小时
        }
    }

    /**
     * 封禁用户
     *
     * @param userId 用户ID
     * @param reason 封禁原因
     * @param hours  封禁小时数
     */
    public void blockUser(String userId, String reason, int hours) {
        UserQuota quota = getUserQuota(userId);
        quota.setBlocked(true);
        quota.setBlockReason(reason);
        quota.setBlockedUntil(LocalDateTime.now().plusHours(hours));

        log.warn("User blocked: {}, reason: {}, until: {}", userId, reason, quota.getBlockedUntil());
    }

    /**
     * 解除用户封禁
     *
     * @param userId 用户ID
     */
    public void unblockUser(String userId) {
        UserQuota quota = getUserQuota(userId);
        quota.setBlocked(false);
        quota.setBlockedUntil(null);
        quota.setBlockReason(null);

        // 重置违规计数
        violationCounters.remove(userId);

        log.info("User unblocked: {}", userId);
    }

    /**
     * 添加IP到黑名单
     *
     * @param ipAddress IP地址
     */
    public void addIpToBlacklist(String ipAddress) {
        ipBlacklist.add(ipAddress);
        log.warn("IP added to blacklist: {}", ipAddress);
    }

    /**
     * 从黑名单移除IP
     *
     * @param ipAddress IP地址
     */
    public void removeIpFromBlacklist(String ipAddress) {
        ipBlacklist.remove(ipAddress);
        log.info("IP removed from blacklist: {}", ipAddress);
    }

    /**
     * 设置用户权限级别
     *
     * @param userId 用户ID
     * @param level  权限级别
     */
    public void setUserPermissionLevel(String userId, PermissionLevel level) {
        UserQuota quota = getUserQuota(userId);
        quota.setPermissionLevel(level);
        log.info("User permission level updated: {}, level: {}", userId, level);
    }

    /**
     * 获取用户配额信息
     */
    private UserQuota getUserQuota(String userId) {
        return userQuotaCache.computeIfAbsent(userId,
                k -> new UserQuota(userId, PermissionLevel.BASIC));
    }

    /**
     * 重置配额（如果需要）
     */
    private void resetQuotaIfNeeded(UserQuota quota) {
        LocalDateTime now = LocalDateTime.now();

        // 重置日配额
        if (ChronoUnit.DAYS.between(quota.getLastResetDaily(), now) >= 1) {
            quota.setDailyUsed(0);
            quota.setLastResetDaily(now);
        }

        // 重置月配额
        if (ChronoUnit.MONTHS.between(quota.getLastResetMonthly(), now) >= 1) {
            quota.setMonthlyUsed(0);
            quota.setLastResetMonthly(now);
        }
    }

    /**
     * 检查权限级别
     */
    private boolean hasPermission(PermissionLevel userLevel, String operation) {
        // 管理员拥有所有权限
        if (userLevel == PermissionLevel.ADMIN) {
            return true;
        }

        // 根据操作类型检查权限
        switch (operation.toLowerCase()) {
            case "chat":
                return userLevel.getLevel() >= PermissionLevel.GUEST.getLevel();
            case "image_generation":
                return userLevel.getLevel() >= PermissionLevel.BASIC.getLevel();
            case "advanced_analysis":
                return userLevel.getLevel() >= PermissionLevel.PREMIUM.getLevel();
            case "batch_processing":
                return userLevel.getLevel() >= PermissionLevel.VIP.getLevel();
            default:
                return userLevel.getLevel() >= PermissionLevel.BASIC.getLevel();
        }
    }

    /**
     * 检查配额限制
     */
    private AccessResult checkQuotaLimit(UserQuota quota, long resourceCost) {
        PermissionLevel level = quota.getPermissionLevel();

        // 管理员无限制
        if (level == PermissionLevel.ADMIN) {
            return new AccessResult(true, "Admin unlimited access");
        }

        // 检查日配额
        if (level.getDailyQuota() > 0 && quota.getDailyUsed() + resourceCost > level.getDailyQuota()) {
            AccessResult result = new AccessResult(false, "Daily quota exceeded");
            result.setRemainingQuota(Math.max(0, level.getDailyQuota() - quota.getDailyUsed()));
            return result;
        }

        // 检查月配额
        if (level.getMonthlyQuota() > 0 && quota.getMonthlyUsed() + resourceCost > level.getMonthlyQuota()) {
            AccessResult result = new AccessResult(false, "Monthly quota exceeded");
            result.setRemainingQuota(Math.max(0, level.getMonthlyQuota() - quota.getMonthlyUsed()));
            return result;
        }

        return new AccessResult(true, "Quota check passed");
    }

    /**
     * 检查限流
     */
    private AccessResult checkRateLimit(String userId, String operation) {
        String key = "user:" + userId + ":" + operation;
        Map<RateLimitType, RateLimitCounter> counters = rateLimitCounters.computeIfAbsent(key,
                k -> new ConcurrentHashMap<>());

        // 每秒限制
        RateLimitCounter perSecondCounter = counters.computeIfAbsent(RateLimitType.PER_SECOND,
                k -> new RateLimitCounter(1000)); // 1秒窗口

        if (!perSecondCounter.tryAcquire(5)) { // 每秒最多5次请求
            AccessResult result = new AccessResult(false, "Rate limit exceeded (per second)");
            result.setResetTime(perSecondCounter.getRemainingTime());
            return result;
        }

        // 每分钟限制
        RateLimitCounter perMinuteCounter = counters.computeIfAbsent(RateLimitType.PER_MINUTE,
                k -> new RateLimitCounter(60 * 1000)); // 1分钟窗口

        if (!perMinuteCounter.tryAcquire(100)) { // 每分钟最多100次请求
            AccessResult result = new AccessResult(false, "Rate limit exceeded (per minute)");
            result.setResetTime(perMinuteCounter.getRemainingTime());
            return result;
        }

        return new AccessResult(true, "Rate limit check passed");
    }

    /**
     * 更新使用量
     */
    private void updateUsage(UserQuota quota, long resourceCost) {
        quota.setDailyUsed(quota.getDailyUsed() + resourceCost);
        quota.setMonthlyUsed(quota.getMonthlyUsed() + resourceCost);
    }

    /**
     * 获取访问控制统计信息
     */
    public Map<String, Object> getAccessControlStats() {
        Map<String, Object> stats = new HashMap<>();
        stats.put("totalUsers", userQuotaCache.size());
        stats.put("blockedUsers", userQuotaCache.values().stream().mapToLong(q -> q.isBlocked() ? 1 : 0).sum());
        stats.put("blacklistedIps", ipBlacklist.size());
        stats.put("blacklistedUsers", userBlacklist.size());
        stats.put("activeApiKeys", apiKeyCache.size());
        stats.put("totalViolations", violationCounters.values().stream().mapToInt(AtomicInteger::get).sum());

        // 按权限级别统计用户数
        Map<String, Long> usersByLevel = new HashMap<>();
        for (PermissionLevel level : PermissionLevel.values()) {
            long count = userQuotaCache.values().stream()
                    .mapToLong(q -> q.getPermissionLevel() == level ? 1 : 0).sum();
            usersByLevel.put(level.name(), count);
        }
        stats.put("usersByPermissionLevel", usersByLevel);

        return stats;
    }

    /**
     * 清理过期数据
     */
    public void cleanupExpiredData() {
        LocalDateTime now = LocalDateTime.now();

        // 清理过期的封禁
        userQuotaCache.values().removeIf(quota -> {
            if (quota.isBlocked() && quota.getBlockedUntil() != null && now.isAfter(quota.getBlockedUntil())) {
                quota.setBlocked(false);
                quota.setBlockedUntil(null);
                quota.setBlockReason(null);
            }
            return false;
        });

        // 清理过期的限流计数器
        rateLimitCounters.entrySet().removeIf(entry -> {
            Map<RateLimitType, RateLimitCounter> counters = entry.getValue();
            counters.entrySet().removeIf(counterEntry -> {
                RateLimitCounter counter = counterEntry.getValue();
                return counter.getRemainingTime() == 0 && counter.getCurrentCount() == 0;
            });
            return counters.isEmpty();
        });

        log.debug("Expired access control data cleaned up");
    }
}