package sipguardian

import (
	"sort"
	"strconv"
	"sync"
	"time"

	"go.uber.org/zap"
)

// EnumerationConfig holds configuration for extension enumeration detection
type EnumerationConfig struct {
	// MaxExtensions is the maximum unique extensions an IP can probe before ban
	MaxExtensions int `json:"max_extensions,omitempty"`

	// ExtensionWindow is the time window for counting unique extensions
	ExtensionWindow time.Duration `json:"extension_window,omitempty"`

	// SequentialThreshold triggers on N+ consecutive extensions (100,101,102...)
	SequentialThreshold int `json:"sequential_threshold,omitempty"`

	// RapidFireCount triggers on N extensions in RapidFireWindow
	RapidFireCount int `json:"rapid_fire_count,omitempty"`

	// RapidFireWindow is the time window for rapid-fire detection
	RapidFireWindow time.Duration `json:"rapid_fire_window,omitempty"`

	// EnumBanTime is ban duration for enumeration attacks (longer than normal)
	EnumBanTime time.Duration `json:"enum_ban_time,omitempty"`

	// ExemptExtensions are extensions that don't count toward enumeration
	ExemptExtensions []string `json:"exempt_extensions,omitempty"`
}

// DefaultEnumerationConfig returns sensible defaults
func DefaultEnumerationConfig() EnumerationConfig {
	return EnumerationConfig{
		MaxExtensions:       20,
		ExtensionWindow:     5 * time.Minute,
		SequentialThreshold: 5,
		RapidFireCount:      10,
		RapidFireWindow:     30 * time.Second,
		EnumBanTime:         2 * time.Hour,
		ExemptExtensions:    []string{},
	}
}

// EnumerationDetector detects extension scanning attacks
type EnumerationDetector struct {
	config    EnumerationConfig
	tracker   map[string]*ExtensionAttempts
	exemptSet map[string]bool
	logger    *zap.Logger
	mu        sync.RWMutex
}

// ExtensionAttempts tracks extension attempts for a single IP
type ExtensionAttempts struct {
	extensions map[string]time.Time // extension -> last attempt time
	ordered    []extAttempt         // time-ordered for pattern detection
	firstSeen  time.Time
	lastSeen   time.Time
	flaggedSeq bool // already detected sequential pattern
	mu         sync.Mutex
}

type extAttempt struct {
	extension string
	timestamp time.Time
	numeric   int // -1 if non-numeric
}

// EnumerationResult contains detection results
type EnumerationResult struct {
	Detected    bool
	Reason      string
	UniqueCount int
	SeqStart    int // start of sequential range (if detected)
	SeqEnd      int // end of sequential range (if detected)
	Extensions  []string
}

// Global detector instance
var (
	globalEnumDetector *EnumerationDetector
	enumDetectorMu     sync.Mutex
)

// GetEnumerationDetector returns the global enumeration detector
func GetEnumerationDetector(logger *zap.Logger) *EnumerationDetector {
	enumDetectorMu.Lock()
	defer enumDetectorMu.Unlock()

	if globalEnumDetector == nil {
		globalEnumDetector = NewEnumerationDetector(logger, DefaultEnumerationConfig())
	}
	return globalEnumDetector
}

// SetEnumerationConfig updates the global detector configuration
func SetEnumerationConfig(config EnumerationConfig) {
	enumDetectorMu.Lock()
	defer enumDetectorMu.Unlock()

	if globalEnumDetector != nil {
		globalEnumDetector.config = config
		// Rebuild exempt set
		globalEnumDetector.exemptSet = make(map[string]bool)
		for _, ext := range config.ExemptExtensions {
			globalEnumDetector.exemptSet[ext] = true
		}
	}
}

// NewEnumerationDetector creates a new detector
func NewEnumerationDetector(logger *zap.Logger, config EnumerationConfig) *EnumerationDetector {
	exemptSet := make(map[string]bool)
	for _, ext := range config.ExemptExtensions {
		exemptSet[ext] = true
	}

	return &EnumerationDetector{
		config:    config,
		tracker:   make(map[string]*ExtensionAttempts),
		exemptSet: exemptSet,
		logger:    logger,
	}
}

// RecordAttempt records an extension attempt and returns detection result
func (ed *EnumerationDetector) RecordAttempt(ip, extension string) EnumerationResult {
	// Skip exempt extensions
	if ed.exemptSet[extension] {
		return EnumerationResult{Detected: false}
	}

	ed.mu.Lock()
	attempts, exists := ed.tracker[ip]
	if !exists {
		attempts = &ExtensionAttempts{
			extensions: make(map[string]time.Time),
			ordered:    make([]extAttempt, 0),
			firstSeen:  time.Now(),
		}
		ed.tracker[ip] = attempts
	}
	ed.mu.Unlock()

	return attempts.record(extension, ed.config)
}

// record adds an attempt and checks for enumeration patterns
func (ea *ExtensionAttempts) record(extension string, config EnumerationConfig) EnumerationResult {
	ea.mu.Lock()
	defer ea.mu.Unlock()

	now := time.Now()
	ea.lastSeen = now

	// Clean old entries
	ea.cleanup(config.ExtensionWindow)

	// Parse as numeric if possible
	numeric := -1
	if n, err := strconv.Atoi(extension); err == nil {
		numeric = n
	}

	// Record the attempt (update timestamp if already seen)
	if _, exists := ea.extensions[extension]; !exists {
		ea.extensions[extension] = now
		ea.ordered = append(ea.ordered, extAttempt{
			extension: extension,
			timestamp: now,
			numeric:   numeric,
		})
	} else {
		ea.extensions[extension] = now
	}

	// Check detection rules
	result := EnumerationResult{
		UniqueCount: len(ea.extensions),
		Extensions:  ea.getExtensionList(),
	}

	// Rule 1: Too many unique extensions
	if len(ea.extensions) >= config.MaxExtensions {
		result.Detected = true
		result.Reason = "extension_count_exceeded"
		return result
	}

	// Rule 2: Sequential pattern detection
	if !ea.flaggedSeq && len(ea.extensions) >= config.SequentialThreshold {
		if detected, start, end := ea.detectSequentialPattern(config.SequentialThreshold); detected {
			ea.flaggedSeq = true
			result.Detected = true
			result.Reason = "sequential_enumeration"
			result.SeqStart = start
			result.SeqEnd = end
			return result
		}
	}

	// Rule 3: Rapid-fire detection
	if detected := ea.detectRapidFire(config.RapidFireCount, config.RapidFireWindow); detected {
		result.Detected = true
		result.Reason = "rapid_fire_enumeration"
		return result
	}

	return result
}

// detectSequentialPattern finds consecutive numeric sequences
func (ea *ExtensionAttempts) detectSequentialPattern(threshold int) (bool, int, int) {
	var nums []int
	for ext := range ea.extensions {
		if n, err := strconv.Atoi(ext); err == nil {
			nums = append(nums, n)
		}
	}

	if len(nums) < threshold {
		return false, 0, 0
	}

	sort.Ints(nums)

	// Find longest consecutive sequence
	maxRun, runStart := 1, nums[0]
	currentRun, currentStart := 1, nums[0]

	for i := 1; i < len(nums); i++ {
		if nums[i] == nums[i-1]+1 {
			currentRun++
			if currentRun > maxRun {
				maxRun = currentRun
				runStart = currentStart
			}
		} else {
			currentRun = 1
			currentStart = nums[i]
		}
	}

	return maxRun >= threshold, runStart, runStart + maxRun - 1
}

// detectRapidFire checks for many extensions in a short time
func (ea *ExtensionAttempts) detectRapidFire(count int, window time.Duration) bool {
	cutoff := time.Now().Add(-window)
	seen := make(map[string]bool)

	for _, attempt := range ea.ordered {
		if attempt.timestamp.After(cutoff) && !seen[attempt.extension] {
			seen[attempt.extension] = true
		}
	}

	return len(seen) >= count
}

// cleanup removes old entries outside the window
func (ea *ExtensionAttempts) cleanup(window time.Duration) {
	cutoff := time.Now().Add(-window)

	// Clean extensions map
	for ext, ts := range ea.extensions {
		if ts.Before(cutoff) {
			delete(ea.extensions, ext)
		}
	}

	// Clean ordered slice
	newOrdered := make([]extAttempt, 0, len(ea.ordered))
	for _, a := range ea.ordered {
		if a.timestamp.After(cutoff) {
			newOrdered = append(newOrdered, a)
		}
	}
	ea.ordered = newOrdered

	// Reset sequential flag if we've cleaned enough entries
	if len(ea.extensions) < 3 {
		ea.flaggedSeq = false
	}
}

func (ea *ExtensionAttempts) getExtensionList() []string {
	exts := make([]string, 0, len(ea.extensions))
	for ext := range ea.extensions {
		exts = append(exts, ext)
	}
	return exts
}

// Cleanup removes stale IP entries from the detector
func (ed *EnumerationDetector) Cleanup() {
	ed.mu.Lock()
	defer ed.mu.Unlock()

	cutoff := time.Now().Add(-ed.config.ExtensionWindow * 2)

	for ip, attempts := range ed.tracker {
		attempts.mu.Lock()
		if attempts.lastSeen.Before(cutoff) {
			delete(ed.tracker, ip)
		}
		attempts.mu.Unlock()
	}
}

// GetStats returns detector statistics
func (ed *EnumerationDetector) GetStats() map[string]interface{} {
	ed.mu.RLock()
	defer ed.mu.RUnlock()

	// Count total tracked extensions across all IPs
	totalExtensions := 0
	for _, attempts := range ed.tracker {
		attempts.mu.Lock()
		totalExtensions += len(attempts.extensions)
		attempts.mu.Unlock()
	}

	return map[string]interface{}{
		"tracked_ips":          len(ed.tracker),
		"total_extensions":     totalExtensions,
		"max_extensions":       ed.config.MaxExtensions,
		"sequential_threshold": ed.config.SequentialThreshold,
		"rapid_fire_count":     ed.config.RapidFireCount,
		"extension_window":     ed.config.ExtensionWindow.String(),
		"rapid_fire_window":    ed.config.RapidFireWindow.String(),
	}
}

// GetIPAttempts returns the extension attempts for a specific IP (for debugging/admin)
func (ed *EnumerationDetector) GetIPAttempts(ip string) *EnumerationResult {
	ed.mu.RLock()
	attempts, exists := ed.tracker[ip]
	ed.mu.RUnlock()

	if !exists {
		return nil
	}

	attempts.mu.Lock()
	defer attempts.mu.Unlock()

	return &EnumerationResult{
		Detected:    false,
		UniqueCount: len(attempts.extensions),
		Extensions:  attempts.getExtensionList(),
	}
}

// ResetIP clears tracking for a specific IP (for admin use)
func (ed *EnumerationDetector) ResetIP(ip string) {
	ed.mu.Lock()
	defer ed.mu.Unlock()
	delete(ed.tracker, ip)
}