package middleware

import (
	"net/http"
	"strings"
	"sync"

	"github.com/gin-gonic/gin"
)

// IPBlacklist IP黑名单管理
var IPBlacklist = &ipBlacklist{
	blockedIPs:  make(map[string]bool),
	allowedIPs:  make(map[string]bool),
	blockedCIDR: []string{},
}

type ipBlacklist struct {
	mu          sync.RWMutex
	blockedIPs  map[string]bool
	allowedIPs  map[string]bool
	blockedCIDR []string // 预留CIDR支持
	enabled     bool
}

// SetEnabled 启用/禁用IP黑名单
func (bl *ipBlacklist) SetEnabled(enabled bool) {
	bl.mu.Lock()
	defer bl.mu.Unlock()
	bl.enabled = enabled
}

// IsEnabled 是否启用
func (bl *ipBlacklist) IsEnabled() bool {
	bl.mu.RLock()
	defer bl.mu.RUnlock()
	return bl.enabled
}

// AddBlockedIP 添加封禁IP
func (bl *ipBlacklist) AddBlockedIP(ip string) {
	bl.mu.Lock()
	defer bl.mu.Unlock()
	bl.blockedIPs[ip] = true
}

// RemoveBlockedIP 移除封禁IP
func (bl *ipBlacklist) RemoveBlockedIP(ip string) {
	bl.mu.Lock()
	defer bl.mu.Unlock()
	delete(bl.blockedIPs, ip)
}

// IsBlocked 检查IP是否被封禁
func (bl *ipBlacklist) IsBlocked(ip string) bool {
	bl.mu.RLock()
	defer bl.mu.RUnlock()

	if !bl.enabled {
		return false
	}

	// 白名单优先
	if bl.allowedIPs[ip] {
		return false
	}

	return bl.blockedIPs[ip]
}

// GetBlockedIPs 获取所有封禁IP
func (bl *ipBlacklist) GetBlockedIPs() []string {
	bl.mu.RLock()
	defer bl.mu.RUnlock()

	ips := make([]string, 0, len(bl.blockedIPs))
	for ip := range bl.blockedIPs {
		ips = append(ips, ip)
	}
	return ips
}

// SetBlockedIPs 批量设置封禁IP
func (bl *ipBlacklist) SetBlockedIPs(ips []string) {
	bl.mu.Lock()
	defer bl.mu.Unlock()

	bl.blockedIPs = make(map[string]bool)
	for _, ip := range ips {
		ip = strings.TrimSpace(ip)
		if ip != "" {
			bl.blockedIPs[ip] = true
		}
	}
}

// AddAllowedIP 添加白名单IP
func (bl *ipBlacklist) AddAllowedIP(ip string) {
	bl.mu.Lock()
	defer bl.mu.Unlock()
	bl.allowedIPs[ip] = true
}

// IPBlacklistMiddleware IP黑名单中间件
func IPBlacklistMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		if !IPBlacklist.IsEnabled() {
			c.Next()
			return
		}

		clientIP := c.ClientIP()

		if IPBlacklist.IsBlocked(clientIP) {
			c.AbortWithStatusJSON(http.StatusForbidden, gin.H{
				"code":    403,
				"message": "您的IP已被封禁",
			})
			return
		}

		c.Next()
	}
}