// Package sipguardian provides a Caddy module for SIP-aware rate limiting and IP banning.
// It integrates with caddy-l4 for Layer 4 proxying and caddy-ratelimit for rate limiting.
package sipguardian

import (
	"fmt"
	"net"
	"sync"
	"time"

	"github.com/caddyserver/caddy/v2"
	"github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
	"go.uber.org/zap"
)

func init() {
	caddy.RegisterModule(SIPGuardian{})
}

// BanEntry represents a banned IP with metadata
type BanEntry struct {
	IP        string    `json:"ip"`
	Reason    string    `json:"reason"`
	BannedAt  time.Time `json:"banned_at"`
	ExpiresAt time.Time `json:"expires_at"`
	HitCount  int       `json:"hit_count"`
}

// SIPGuardian implements intelligent SIP protection at Layer 4
type SIPGuardian struct {
	// Configuration
	MaxFailures   int           `json:"max_failures,omitempty"`
	FindTime      caddy.Duration `json:"find_time,omitempty"`
	BanTime       caddy.Duration `json:"ban_time,omitempty"`
	WhitelistCIDR []string       `json:"whitelist_cidr,omitempty"`

	// Runtime state
	logger       *zap.Logger
	bannedIPs    map[string]*BanEntry
	failureCounts map[string]*failureTracker
	whitelistNets []*net.IPNet
	mu           sync.RWMutex
}

type failureTracker struct {
	count     int
	firstSeen time.Time
	lastSeen  time.Time
}

// CaddyModule returns the Caddy module information.
func (SIPGuardian) CaddyModule() caddy.ModuleInfo {
	return caddy.ModuleInfo{
		ID:  "sip_guardian",
		New: func() caddy.Module { return new(SIPGuardian) },
	}
}

// Provision sets up the module.
func (g *SIPGuardian) Provision(ctx caddy.Context) error {
	g.logger = ctx.Logger()
	g.bannedIPs = make(map[string]*BanEntry)
	g.failureCounts = make(map[string]*failureTracker)

	// Set defaults
	if g.MaxFailures == 0 {
		g.MaxFailures = 5
	}
	if g.FindTime == 0 {
		g.FindTime = caddy.Duration(10 * time.Minute)
	}
	if g.BanTime == 0 {
		g.BanTime = caddy.Duration(1 * time.Hour)
	}

	// Parse whitelist CIDRs
	for _, cidr := range g.WhitelistCIDR {
		_, network, err := net.ParseCIDR(cidr)
		if err != nil {
			return fmt.Errorf("invalid whitelist CIDR %s: %v", cidr, err)
		}
		g.whitelistNets = append(g.whitelistNets, network)
	}

	// Start cleanup goroutine
	go g.cleanupLoop(ctx)

	g.logger.Info("SIP Guardian initialized",
		zap.Int("max_failures", g.MaxFailures),
		zap.Duration("find_time", time.Duration(g.FindTime)),
		zap.Duration("ban_time", time.Duration(g.BanTime)),
		zap.Int("whitelist_count", len(g.whitelistNets)),
	)

	return nil
}

// IsWhitelisted checks if an IP is in the whitelist
func (g *SIPGuardian) IsWhitelisted(ip string) bool {
	parsedIP := net.ParseIP(ip)
	if parsedIP == nil {
		return false
	}
	for _, network := range g.whitelistNets {
		if network.Contains(parsedIP) {
			return true
		}
	}
	return false
}

// IsBanned checks if an IP is currently banned
func (g *SIPGuardian) IsBanned(ip string) bool {
	g.mu.RLock()
	defer g.mu.RUnlock()

	if entry, exists := g.bannedIPs[ip]; exists {
		if time.Now().Before(entry.ExpiresAt) {
			return true
		}
	}
	return false
}

// RecordFailure records a failed authentication attempt
func (g *SIPGuardian) RecordFailure(ip, reason string) bool {
	if g.IsWhitelisted(ip) {
		return false
	}

	g.mu.Lock()
	defer g.mu.Unlock()

	now := time.Now()
	findWindow := time.Duration(g.FindTime)

	tracker, exists := g.failureCounts[ip]
	if !exists || now.Sub(tracker.firstSeen) > findWindow {
		// Start new tracking window
		tracker = &failureTracker{
			count:     1,
			firstSeen: now,
			lastSeen:  now,
		}
		g.failureCounts[ip] = tracker
	} else {
		tracker.count++
		tracker.lastSeen = now
	}

	g.logger.Debug("Failure recorded",
		zap.String("ip", ip),
		zap.String("reason", reason),
		zap.Int("count", tracker.count),
	)

	// Check if we should ban
	if tracker.count >= g.MaxFailures {
		g.banIP(ip, reason)
		return true
	}

	return false
}

// banIP adds an IP to the ban list (must hold lock)
func (g *SIPGuardian) banIP(ip, reason string) {
	now := time.Now()
	banDuration := time.Duration(g.BanTime)

	g.bannedIPs[ip] = &BanEntry{
		IP:        ip,
		Reason:    reason,
		BannedAt:  now,
		ExpiresAt: now.Add(banDuration),
		HitCount:  g.failureCounts[ip].count,
	}

	// Clear failure counter
	delete(g.failureCounts, ip)

	g.logger.Warn("IP banned",
		zap.String("ip", ip),
		zap.String("reason", reason),
		zap.Duration("duration", banDuration),
	)
}

// UnbanIP manually removes an IP from the ban list
func (g *SIPGuardian) UnbanIP(ip string) bool {
	g.mu.Lock()
	defer g.mu.Unlock()

	if _, exists := g.bannedIPs[ip]; exists {
		delete(g.bannedIPs, ip)
		g.logger.Info("IP unbanned", zap.String("ip", ip))
		return true
	}
	return false
}

// GetBannedIPs returns a list of currently banned IPs
func (g *SIPGuardian) GetBannedIPs() []BanEntry {
	g.mu.RLock()
	defer g.mu.RUnlock()

	var entries []BanEntry
	now := time.Now()
	for _, entry := range g.bannedIPs {
		if now.Before(entry.ExpiresAt) {
			entries = append(entries, *entry)
		}
	}
	return entries
}

// GetStats returns current statistics
func (g *SIPGuardian) GetStats() map[string]interface{} {
	g.mu.RLock()
	defer g.mu.RUnlock()

	activeBans := 0
	now := time.Now()
	for _, entry := range g.bannedIPs {
		if now.Before(entry.ExpiresAt) {
			activeBans++
		}
	}

	return map[string]interface{}{
		"active_bans":      activeBans,
		"tracked_failures": len(g.failureCounts),
		"whitelist_count":  len(g.whitelistNets),
	}
}

// cleanupLoop periodically removes expired entries
func (g *SIPGuardian) cleanupLoop(ctx caddy.Context) {
	ticker := time.NewTicker(1 * time.Minute)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			g.cleanup()
		}
	}
}

func (g *SIPGuardian) cleanup() {
	g.mu.Lock()
	defer g.mu.Unlock()

	now := time.Now()
	findWindow := time.Duration(g.FindTime)

	// Cleanup expired bans
	for ip, entry := range g.bannedIPs {
		if now.After(entry.ExpiresAt) {
			delete(g.bannedIPs, ip)
			g.logger.Debug("Ban expired", zap.String("ip", ip))
		}
	}

	// Cleanup old failure trackers
	for ip, tracker := range g.failureCounts {
		if now.Sub(tracker.firstSeen) > findWindow {
			delete(g.failureCounts, ip)
		}
	}
}

// UnmarshalCaddyfile implements caddyfile.Unmarshaler.
func (g *SIPGuardian) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
	for d.Next() {
		for d.NextBlock(0) {
			switch d.Val() {
			case "max_failures":
				if !d.NextArg() {
					return d.ArgErr()
				}
				var val int
				if _, err := fmt.Sscanf(d.Val(), "%d", &val); err != nil {
					return d.Errf("invalid max_failures: %v", err)
				}
				g.MaxFailures = val

			case "find_time":
				if !d.NextArg() {
					return d.ArgErr()
				}
				dur, err := caddy.ParseDuration(d.Val())
				if err != nil {
					return d.Errf("invalid find_time: %v", err)
				}
				g.FindTime = caddy.Duration(dur)

			case "ban_time":
				if !d.NextArg() {
					return d.ArgErr()
				}
				dur, err := caddy.ParseDuration(d.Val())
				if err != nil {
					return d.Errf("invalid ban_time: %v", err)
				}
				g.BanTime = caddy.Duration(dur)

			case "whitelist":
				for d.NextArg() {
					g.WhitelistCIDR = append(g.WhitelistCIDR, d.Val())
				}

			default:
				return d.Errf("unknown directive: %s", d.Val())
			}
		}
	}
	return nil
}

// Interface guards
var (
	_ caddy.Module          = (*SIPGuardian)(nil)
	_ caddy.Provisioner     = (*SIPGuardian)(nil)
	_ caddyfile.Unmarshaler = (*SIPGuardian)(nil)
)