package app

import (
	"encoding/json"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"
)

// ThreatIntel tracks malicious IPs and their activities
type ThreatIntel struct {
	mu           sync.RWMutex
	maliciousIPs map[string]*IPThreatInfo
	configPath   string
	autoSave     bool
}

// IPThreatInfo contains information about a potentially malicious IP
type IPThreatInfo struct {
	IP                string                 `json:"ip"`
	FirstSeen         time.Time              `json:"first_seen"`
	LastSeen          time.Time              `json:"last_seen"`
	TotalConnections  int                    `json:"total_connections"`
	Services          map[string]int         `json:"services"`          // service -> connection count
	AuthAttempts      int                    `json:"auth_attempts"`     // total authentication attempts
	UniqueUsernames   map[string]int         `json:"unique_usernames"`  // username -> attempt count
	UniquePasswords   map[string]int         `json:"unique_passwords"`  // password -> attempt count
	Countries         map[string]int         `json:"countries"`         // country -> count (if GeoIP available)
	ThreatScore       int                    `json:"threat_score"`      // calculated threat score
	IsBlacklisted     bool                   `json:"is_blacklisted"`
	Notes             []string               `json:"notes"`
	RawPayloads       []string               `json:"raw_payloads,omitempty"` // sample payloads
}

// NewThreatIntel creates a new threat intelligence tracker
func NewThreatIntel(configPath string, autoSave bool) *ThreatIntel {
	ti := &ThreatIntel{
		maliciousIPs: make(map[string]*IPThreatInfo),
		configPath:   configPath,
		autoSave:     autoSave,
	}
	
	// Try to load existing data
	ti.Load()
	
	// Start auto-save routine if enabled
	if autoSave {
		go ti.autoSaveRoutine()
	}
	
	return ti
}

// RecordActivity records activity from an IP address
func (ti *ThreatIntel) RecordActivity(record Record) {
	ti.mu.Lock()
	defer ti.mu.Unlock()
	
	ip := record.RemoteAddr
	if ip == "" {
		return
	}
	
	// Get or create IP info
	info, exists := ti.maliciousIPs[ip]
	if !exists {
		info = &IPThreatInfo{
			IP:              ip,
			FirstSeen:       record.Timestamp,
			Services:        make(map[string]int),
			UniqueUsernames: make(map[string]int),
			UniquePasswords: make(map[string]int),
			Countries:       make(map[string]int),
			Notes:           []string{},
			RawPayloads:     []string{},
		}
		ti.maliciousIPs[ip] = info
	}
	
	// Update basic info
	info.LastSeen = record.Timestamp
	info.TotalConnections++
	info.Services[record.Service]++
	
	// Process authentication attempts
	if record.Details != nil {
		if username, ok := record.Details["username"]; ok && username != "" {
			info.UniqueUsernames[username]++
			info.AuthAttempts++
		}
		if password, ok := record.Details["password"]; ok && password != "" {
			info.UniquePasswords[password]++
		}
		
		// Check for specific attack patterns
		ti.analyzeAttackPatterns(info, record)
	}
	
	// Store sample payloads (limit to 10)
	if record.RawPayload != "" && len(info.RawPayloads) < 10 {
		info.RawPayloads = append(info.RawPayloads, record.RawPayload)
	}
	
	// Recalculate threat score
	info.ThreatScore = ti.calculateThreatScore(info)
	
	// Auto-blacklist based on threat score
	if info.ThreatScore >= 100 && !info.IsBlacklisted {
		info.IsBlacklisted = true
		info.Notes = append(info.Notes, fmt.Sprintf("Auto-blacklisted at %s (threat score: %d)", time.Now().Format(time.RFC3339), info.ThreatScore))
	}
}

// analyzeAttackPatterns looks for specific attack patterns and adds notes
func (ti *ThreatIntel) analyzeAttackPatterns(info *IPThreatInfo, record Record) {
	// Check for brute force patterns
	if info.AuthAttempts > 10 {
		if len(info.Notes) == 0 || info.Notes[len(info.Notes)-1] != "Brute force attack detected" {
			info.Notes = append(info.Notes, "Brute force attack detected")
		}
	}
	
	// Check for credential stuffing (many unique usernames)
	if len(info.UniqueUsernames) > 20 {
		if len(info.Notes) == 0 || info.Notes[len(info.Notes)-1] != "Credential stuffing attack detected" {
			info.Notes = append(info.Notes, "Credential stuffing attack detected")
		}
	}
	
	// Check for service scanning (multiple services)
	if len(info.Services) > 5 {
		if len(info.Notes) == 0 || info.Notes[len(info.Notes)-1] != "Port/service scanning detected" {
			info.Notes = append(info.Notes, "Port/service scanning detected")
		}
	}
	
	// Check for common attack usernames
	commonAttackUsernames := []string{"admin", "root", "administrator", "user", "test", "guest", "oracle", "postgres", "mysql"}
	if username, ok := record.Details["username"]; ok {
		for _, attackUser := range commonAttackUsernames {
			if username == attackUser {
				info.Notes = append(info.Notes, fmt.Sprintf("Used common attack username: %s", username))
				break
			}
		}
	}
	
	// Check for common attack passwords
	commonAttackPasswords := []string{"password", "123456", "admin", "root", "toor", "password123", "qwerty"}
	if password, ok := record.Details["password"]; ok {
		for _, attackPass := range commonAttackPasswords {
			if password == attackPass {
				info.Notes = append(info.Notes, fmt.Sprintf("Used common attack password: %s", password))
				break
			}
		}
	}
}

// calculateThreatScore calculates a threat score based on various factors
func (ti *ThreatIntel) calculateThreatScore(info *IPThreatInfo) int {
	score := 0
	
	// Base score for any connection
	score += 1
	
	// Score based on number of connections
	if info.TotalConnections > 100 {
		score += 50
	} else if info.TotalConnections > 50 {
		score += 30
	} else if info.TotalConnections > 10 {
		score += 15
	} else if info.TotalConnections > 5 {
		score += 5
	}
	
	// Score based on authentication attempts
	if info.AuthAttempts > 50 {
		score += 40
	} else if info.AuthAttempts > 20 {
		score += 25
	} else if info.AuthAttempts > 10 {
		score += 15
	} else if info.AuthAttempts > 5 {
		score += 10
	}
	
	// Score based on unique usernames (credential stuffing)
	if len(info.UniqueUsernames) > 50 {
		score += 30
	} else if len(info.UniqueUsernames) > 20 {
		score += 20
	} else if len(info.UniqueUsernames) > 10 {
		score += 10
	}
	
	// Score based on service diversity (scanning)
	if len(info.Services) > 10 {
		score += 25
	} else if len(info.Services) > 5 {
		score += 15
	} else if len(info.Services) > 3 {
		score += 10
	}
	
	// Score based on time span (persistent attacker)
	timeSpan := info.LastSeen.Sub(info.FirstSeen)
	if timeSpan > 24*time.Hour {
		score += 20
	} else if timeSpan > 12*time.Hour {
		score += 15
	} else if timeSpan > 6*time.Hour {
		score += 10
	} else if timeSpan > 1*time.Hour {
		score += 5
	}
	
	return score
}

// IsBlacklisted checks if an IP is blacklisted
func (ti *ThreatIntel) IsBlacklisted(ip string) bool {
	ti.mu.RLock()
	defer ti.mu.RUnlock()
	
	if info, exists := ti.maliciousIPs[ip]; exists {
		return info.IsBlacklisted
	}
	return false
}

// GetThreatInfo returns threat information for an IP
func (ti *ThreatIntel) GetThreatInfo(ip string) (*IPThreatInfo, bool) {
	ti.mu.RLock()
	defer ti.mu.RUnlock()
	
	info, exists := ti.maliciousIPs[ip]
	if !exists {
		return nil, false
	}
	
	// Return a copy to avoid race conditions
	infoCopy := *info
	infoCopy.Services = make(map[string]int)
	infoCopy.UniqueUsernames = make(map[string]int)
	infoCopy.UniquePasswords = make(map[string]int)
	infoCopy.Countries = make(map[string]int)
	
	for k, v := range info.Services {
		infoCopy.Services[k] = v
	}
	for k, v := range info.UniqueUsernames {
		infoCopy.UniqueUsernames[k] = v
	}
	for k, v := range info.UniquePasswords {
		infoCopy.UniquePasswords[k] = v
	}
	for k, v := range info.Countries {
		infoCopy.Countries[k] = v
	}
	
	infoCopy.Notes = make([]string, len(info.Notes))
	copy(infoCopy.Notes, info.Notes)
	
	infoCopy.RawPayloads = make([]string, len(info.RawPayloads))
	copy(infoCopy.RawPayloads, info.RawPayloads)
	
	return &infoCopy, true
}

// GetTopThreats returns the top N threats by score
func (ti *ThreatIntel) GetTopThreats(n int) []*IPThreatInfo {
	ti.mu.RLock()
	defer ti.mu.RUnlock()
	
	var threats []*IPThreatInfo
	for _, info := range ti.maliciousIPs {
		threats = append(threats, info)
	}
	
	// Sort by threat score (simple bubble sort for small datasets)
	for i := 0; i < len(threats)-1; i++ {
		for j := 0; j < len(threats)-i-1; j++ {
			if threats[j].ThreatScore < threats[j+1].ThreatScore {
				threats[j], threats[j+1] = threats[j+1], threats[j]
			}
		}
	}
	
	if n > len(threats) {
		n = len(threats)
	}
	
	return threats[:n]
}

// GetBlacklistedIPs returns all blacklisted IPs
func (ti *ThreatIntel) GetBlacklistedIPs() []string {
	ti.mu.RLock()
	defer ti.mu.RUnlock()
	
	var blacklisted []string
	for ip, info := range ti.maliciousIPs {
		if info.IsBlacklisted {
			blacklisted = append(blacklisted, ip)
		}
	}
	
	return blacklisted
}

// ManualBlacklist manually blacklists an IP
func (ti *ThreatIntel) ManualBlacklist(ip string, reason string) {
	ti.mu.Lock()
	defer ti.mu.Unlock()
	
	info, exists := ti.maliciousIPs[ip]
	if !exists {
		info = &IPThreatInfo{
			IP:              ip,
			FirstSeen:       time.Now(),
			Services:        make(map[string]int),
			UniqueUsernames: make(map[string]int),
			UniquePasswords: make(map[string]int),
			Countries:       make(map[string]int),
			Notes:           []string{},
			RawPayloads:     []string{},
		}
		ti.maliciousIPs[ip] = info
	}
	
	info.IsBlacklisted = true
	info.LastSeen = time.Now()
	note := fmt.Sprintf("Manually blacklisted at %s", time.Now().Format(time.RFC3339))
	if reason != "" {
		note += fmt.Sprintf(" - Reason: %s", reason)
	}
	info.Notes = append(info.Notes, note)
}

// RemoveFromBlacklist removes an IP from the blacklist
func (ti *ThreatIntel) RemoveFromBlacklist(ip string) {
	ti.mu.Lock()
	defer ti.mu.Unlock()
	
	if info, exists := ti.maliciousIPs[ip]; exists {
		info.IsBlacklisted = false
		info.Notes = append(info.Notes, fmt.Sprintf("Removed from blacklist at %s", time.Now().Format(time.RFC3339)))
	}
}

// Save saves threat intelligence data to file
func (ti *ThreatIntel) Save() error {
	ti.mu.RLock()
	defer ti.mu.RUnlock()
	
	if ti.configPath == "" {
		return nil
	}
	
	// Create directory if it doesn't exist
	if err := os.MkdirAll(filepath.Dir(ti.configPath), 0755); err != nil {
		return fmt.Errorf("create threat intel dir: %w", err)
	}
	
	data, err := json.MarshalIndent(ti.maliciousIPs, "", "  ")
	if err != nil {
		return fmt.Errorf("marshal threat intel data: %w", err)
	}
	
	return os.WriteFile(ti.configPath, data, 0644)
}

// Load loads threat intelligence data from file
func (ti *ThreatIntel) Load() error {
	ti.mu.Lock()
	defer ti.mu.Unlock()
	
	if ti.configPath == "" {
		return nil
	}
	
	data, err := os.ReadFile(ti.configPath)
	if err != nil {
		if os.IsNotExist(err) {
			return nil // File doesn't exist yet, that's okay
		}
		return fmt.Errorf("read threat intel file: %w", err)
	}
	
	return json.Unmarshal(data, &ti.maliciousIPs)
}

// autoSaveRoutine periodically saves threat intelligence data
func (ti *ThreatIntel) autoSaveRoutine() {
	ticker := time.NewTicker(5 * time.Minute)
	defer ticker.Stop()
	
	for range ticker.C {
		if err := ti.Save(); err != nil {
			// Log error but don't stop the routine
			fmt.Printf("Error auto-saving threat intel: %v\n", err)
		}
	}
}

// GetStats returns overall statistics
func (ti *ThreatIntel) GetStats() map[string]interface{} {
	ti.mu.RLock()
	defer ti.mu.RUnlock()
	
	totalIPs := len(ti.maliciousIPs)
	blacklistedIPs := 0
	totalConnections := 0
	totalAuthAttempts := 0
	serviceStats := make(map[string]int)
	
	for _, info := range ti.maliciousIPs {
		if info.IsBlacklisted {
			blacklistedIPs++
		}
		totalConnections += info.TotalConnections
		totalAuthAttempts += info.AuthAttempts
		
		for service, count := range info.Services {
			serviceStats[service] += count
		}
	}
	
	return map[string]interface{}{
		"total_ips":           totalIPs,
		"blacklisted_ips":     blacklistedIPs,
		"total_connections":   totalConnections,
		"total_auth_attempts": totalAuthAttempts,
		"service_stats":       serviceStats,
	}
}

// ShouldBlock determines if a connection should be blocked based on IP reputation
func (ti *ThreatIntel) ShouldBlock(ip string) bool {
	// Check if IP is blacklisted
	if ti.IsBlacklisted(ip) {
		return true
	}
	
	// Additional checks could be added here:
	// - Rate limiting
	// - Temporary blocks for high-frequency connections
	// - Integration with external threat feeds
	
	return false
}

// IsPrivateIP checks if an IP address is private/internal
func IsPrivateIP(ip string) bool {
	parsedIP := net.ParseIP(ip)
	if parsedIP == nil {
		return false
	}
	
	// Check for private IP ranges
	privateRanges := []string{
		"10.0.0.0/8",
		"172.16.0.0/12",
		"192.168.0.0/16",
		"127.0.0.0/8",
		"169.254.0.0/16",
		"::1/128",
		"fc00::/7",
		"fe80::/10",
	}
	
	for _, cidr := range privateRanges {
		_, network, err := net.ParseCIDR(cidr)
		if err != nil {
			continue
		}
		if network.Contains(parsedIP) {
			return true
		}
	}
	
	return false
}