package middleware

import (
	"net/http"
	"sync"
	"time"

	"simple_portal/models"

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

// FailRecord stores the login failure count and first failure time for an IP.
type FailRecord struct {
	Count   int
	FirstAt time.Time
}

// IPBanGuard maintains an in-memory counter of login failures per IP.
// When 5 failures occur within 5 minutes, the IP is automatically banned.
type IPBanGuard struct {
	failMap sync.Map // map[string]*FailRecord
}

// NewIPBanGuard creates a new IPBanGuard instance.
func NewIPBanGuard() *IPBanGuard {
	return &IPBanGuard{}
}

// RecordFail records a login failure for the given IP.
// It returns true if the IP should be auto-banned (5 failures within 5 minutes).
// The caller is responsible for creating the actual ban record in the database.
func (g *IPBanGuard) RecordFail(ip string) bool {
	now := time.Now()

	// Try to load existing record
	if val, ok := g.failMap.Load(ip); ok {
		rec := val.(*FailRecord)
		// Check if within 5-minute window
		if now.Sub(rec.FirstAt) > 5*time.Minute {
			// Reset — outside the window
			rec.Count = 1
			rec.FirstAt = now
			return false
		}
		rec.Count++
		if rec.Count >= 5 {
			// Auto-ban: reset the counter after creating the ban
			g.failMap.Delete(ip)
			return true
		}
		return false
	}

	// New record
	g.failMap.Store(ip, &FailRecord{Count: 1, FirstAt: now})
	return false
}

// ResetFail resets the failure counter for the given IP (on successful login).
func (g *IPBanGuard) ResetFail(ip string) {
	g.failMap.Delete(ip)
}

// IPWhitelistRequired returns a Gin middleware that checks if the visitor's IP
// is in the whitelist. Only when the whitelist table has records does it restrict
// access; when the whitelist is empty, no IP is restricted.
//
// 登录页面路由（/admin/login）始终放行。
// 已通过 session 认证的用户也放行，避免管理员添加白名单后自己被锁定。
// sessionCheck 是一个函数，接收 sessionID 返回是否有效。
func IPWhitelistRequired(sessionCheck func(sessionID string) bool) gin.HandlerFunc {
	return func(c *gin.Context) {
		// 登录页面始终放行
		path := c.Request.URL.Path
		if path == "/admin/login" {
			c.Next()
			return
		}

		hasWhitelist, err := models.HasWhitelist()
		if err != nil {
			c.Next()
			return
		}
		if !hasWhitelist {
			c.Next()
			return
		}

		clientIP := c.ClientIP()
		allowed, err := models.IsIPWhitelisted(clientIP)
		if err != nil {
			c.Next()
			return
		}
		if allowed {
			c.Next()
			return
		}

		// IP不在白名单中，但如果有有效session则放行
		if sessionID, err := c.Cookie("session_id"); err == nil && sessionID != "" {
			if sessionCheck(sessionID) {
				c.Next()
				return
			}
		}

		c.HTML(http.StatusForbidden, "admin/403.html", gin.H{
			"Title": "访问被拒绝",
			"IP":    clientIP,
		})
		c.Abort()
	}
}