// IPFilter middleware blocks requests from unwanted IPs based on GeoIP, ASN, and other rules.
// Package middleware provides IP and User-Agent filtering to protect against abuse.
// Features:
// - GeoIP country/continent blocking
// - ASN blocking
// - Data center detection
// - Bot detection
// - Cache-optimized lookups
package middleware

import (
	"log"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/cloudflare/ahocorasick"
	"github.com/gofiber/fiber/v2"
	geoip2 "github.com/oschwald/geoip2-golang"
)

// --- IP Block Cache --- //
// blockCacheEntry stores the result of a block check for an IP
type blockCacheEntry struct {
	blocked    bool
	blockType  string
	blockValue string
	customPage string
	asnOrgName string
	expiresAt  time.Time
}

// Cache for IP block lookups (TTL set via config)
var (
	ipBlockCache      map[string]blockCacheEntry
	ipBlockCacheMutex sync.RWMutex
	ipBlockCacheTTL   time.Duration
)

// --- GeoIP Configuration ---
var (
	geoIPCountryDBPath = "./data/GeoLite2-Country.mmdb"
	geoIPASNDBPath     = "./data/GeoLite2-ASN.mmdb"
	geoipCountryReader *geoip2.Reader
	geoipASNReader     *geoip2.Reader

	// Lists and pages (populated from ipfilter.toml)
	blockedCountryCodes   map[string]bool
	blockedContinentCodes map[string]bool
	blockedASNs           map[string][]uint
	blockedASNNames       map[string][]string
	countryBlockPages     map[string]string
	continentBlockPages   map[string]string
	asnGroupBlockPages    map[string]string
	defaultBlockPage      string

	// Aho-Corasick matchers for ASN name groups
	asnNameMatchers      map[string]*ahocorasick.Matcher
	asnNameMatchersMutex sync.RWMutex
)

// init loads ipfilter.toml, sets up all block lists, and hooks this into our plugin registry
func init() {
	// Load full configuration from ipfilter.toml
	type ipfilterConfig struct {
		BlockedCountryCodes   []string            `toml:"blockedCountryCodes"`
		BlockedContinentCodes []string            `toml:"blockedContinentCodes"`
		BlockedASNs           map[string][]int    `toml:"blockedASNs"`
		BlockedASNNames       map[string][]string `toml:"blockedASNNames"`
		CountryBlockPages     map[string]string   `toml:"countryBlockPages"`
		ContinentBlockPages   map[string]string   `toml:"continentBlockPages"`
		ASNGroupBlockPages    map[string]string   `toml:"asnGroupBlockPages"`
		DefaultBlockPage      string              `toml:"defaultBlockPage"`
		IPBlockCacheTTLSec    int                 `toml:"ipBlockCacheTTLSec"`
	}
	var cfg ipfilterConfig
	if err := LoadConfig("ipfilter", &cfg); err != nil {
		log.Fatalf("Failed to load ipfilter config: %v", err)
	}
	// override blockedCountryCodes
	blockedCountryCodes = make(map[string]bool)
	for _, c := range cfg.BlockedCountryCodes {
		blockedCountryCodes[c] = true
	}
	// override blockedContinentCodes
	blockedContinentCodes = make(map[string]bool)
	for _, c := range cfg.BlockedContinentCodes {
		blockedContinentCodes[c] = true
	}
	// override blockedASNs
	newASNs := make(map[string][]uint)
	for group, arr := range cfg.BlockedASNs {
		uintArr := make([]uint, len(arr))
		for i, v := range arr {
			uintArr[i] = uint(v)
		}
		newASNs[group] = uintArr
	}
	blockedASNs = newASNs
	// override blockedASNNames
	blockedASNNames = cfg.BlockedASNNames
	// override page maps
	countryBlockPages = cfg.CountryBlockPages
	continentBlockPages = cfg.ContinentBlockPages
	asnGroupBlockPages = cfg.ASNGroupBlockPages
	// override default page
	defaultBlockPage = cfg.DefaultBlockPage
	// override cache TTL
	ipBlockCacheTTL = time.Duration(cfg.IPBlockCacheTTLSec) * time.Second

	// Initialize the cache
	ipBlockCache = make(map[string]blockCacheEntry)

	// Initialize GeoIP databases
	initGeoIP()

	// Build Aho-Corasick matchers for ASN name matching
	buildASNNameMatchers()

	// Register IP-block plugin
	RegisterPlugin("ipfilter", IPBlockMiddleware)
}

// IPBlockMiddleware returns a handler that stops requests coming from blocked IPs
// Use this early so bad traffic gets rejected fast
func IPBlockMiddleware() fiber.Handler {
	return func(c *fiber.Ctx) error {
		// Get the client IP
		clientIP := getRealIP(c)

		// Check IP blocklist (GeoIP/ASN)
		if blocked, blockType, blockValue, customPage, asnOrgName := isBlockedIPExtended(clientIP); blocked {
			// For API requests, return JSON response
			if strings.HasPrefix(c.Path(), "/api") {
				return c.Status(fiber.StatusForbidden).JSON(fiber.Map{
					"error":   "Access denied from your location or network.",
					"reason":  "geoip",
					"type":    blockType,
					"value":   blockValue,
					"asn_org": asnOrgName,
				})
			}

			// For HTML requests, serve a block page
			// Log block information
			log.Printf("Blocking access: type=%s, value=%s, custom_page=%s, asn_org=%s",
				blockType, blockValue, customPage, asnOrgName)

			// If no custom page specified, use the default
			if customPage == "" {
				customPage = defaultBlockPage
			}

			// Try to serve the block page
			possiblePaths := []string{
				"./public/static/" + customPage,
				"./public/html/" + customPage,
				"./public/" + customPage,
				"./develop/static/" + customPage,
			}

			var htmlContentBytes []byte
			var err error
			foundPath := ""

			for _, path := range possiblePaths {
				if htmlContentBytes, err = os.ReadFile(path); err == nil {
					// Found the file, serve it
					foundPath = path
					break
				}
			}

			if foundPath != "" {
				// Found the file, serve it
				c.Status(fiber.StatusForbidden)
				c.Set("Content-Type", "text/html; charset=utf-8")

				// Replace placeholder with ASN Org Name if available
				htmlContent := string(htmlContentBytes)
				replaceValue := "Blocked Network" // Default fallback text
				if asnOrgName != "" {
					replaceValue = asnOrgName
				}
				htmlContent = strings.Replace(htmlContent, "{{.ASNName}}", replaceValue, -1)

				return c.SendString(htmlContent)
			}

			// If no custom or default page was found, fall back to default message
			return c.Status(fiber.StatusForbidden).SendString("Access denied from your location or network.")
		}

		// Not blocked, continue to next middleware
		return c.Next()
	}
}

// initGeoIP loads the GeoLite2 databases
func initGeoIP() {
	var err error
	geoipCountryReader, err = geoip2.Open(geoIPCountryDBPath)
	if err != nil {
		log.Printf("WARNING: Could not open GeoLite2 Country database at %s: %v. GeoIP country/continent blocking disabled.", geoIPCountryDBPath, err)
	}
	geoipASNReader, err = geoip2.Open(geoIPASNDBPath)
	if err != nil {
		log.Printf("WARNING: Could not open GeoLite2 ASN database at %s: %v. GeoIP ASN blocking disabled.", geoIPASNDBPath, err)
	}
}

// isBlockedIP checks if the IP address is blocked based on GeoIP or ASN.
func isBlockedIP(ipStr string) bool {
	blocked, _, _, _, _ := isBlockedIPExtended(ipStr)
	return blocked
}

// isBlockedIPExtended checks if the IP address is blocked and returns detailed information.
// Returns:
// - blocked: true if IP is blocked
// - blockType: "country", "continent", "asn_number_group", or "asn_name_group"
// - blockValue: the country code, continent code, or ASN group name (that triggered the block)
// - customPage: path to custom block page if defined, or empty string
// - asnOrgName: the specific ASN Organization name if blockType is "asn_name_group"
func isBlockedIPExtended(ipStr string) (blocked bool, blockType string, blockValue string, customPage string, asnOrgName string) {
	// Safeguard against empty IPs causing issues
	if ipStr == "" {
		log.Printf("ERROR: Empty IP passed to isBlockedIPExtended")
		return false, "", "", "", ""
	}

	ip := net.ParseIP(ipStr)
	if ip == nil {
		log.Printf("WARNING: Could not parse IP address for GeoIP check: %s", ipStr)
		return false, "", "", "", "" // Cannot check invalid IP
	}

	// --- Check Cache First --- //
	ipBlockCacheMutex.RLock()
	entry, found := ipBlockCache[ipStr]
	ipBlockCacheMutex.RUnlock() // Unlock immediately after reading

	if found && time.Now().Before(entry.expiresAt) {
		// Return cached result
		return entry.blocked, entry.blockType, entry.blockValue, entry.customPage, entry.asnOrgName
	}
	// --- End Cache Check --- //

	var countryRecord *geoip2.Country
	var asnRecord *geoip2.ASN
	var countryErr, asnErr error
	var wg sync.WaitGroup

	// Concurrently perform GeoIP lookups if readers are available
	if geoipCountryReader != nil {
		wg.Add(1)
		go func() {
			defer wg.Done()
			countryRecord, countryErr = geoipCountryReader.Country(ip)
		}()
	}

	if geoipASNReader != nil {
		wg.Add(1)
		go func() {
			defer wg.Done()
			asnRecord, asnErr = geoipASNReader.ASN(ip)
		}()
	}

	// Wait for lookups to complete
	wg.Wait()

	// --- Process Results --- //

	// 1. Check Country/Continent
	if countryRecord != nil && countryErr == nil {
		// Check country first
		if blockedCountryCodes[countryRecord.Country.IsoCode] {
			blockType = "country"
			blockValue = countryRecord.Country.IsoCode
			log.Printf("INFO: Blocking IP %s based on %s: %s", ipStr, blockType, blockValue)
			customPage = countryBlockPages[blockValue]
			return cacheAndReturnBlockResult(ipStr, blockType, blockValue, customPage, "")
		}

		// Then check continent
		if blockedContinentCodes[countryRecord.Continent.Code] {
			blockType = "continent"
			blockValue = countryRecord.Continent.Code
			log.Printf("INFO: Blocking IP %s based on %s: %s", ipStr, blockType, blockValue)
			customPage = continentBlockPages[blockValue]
			return cacheAndReturnBlockResult(ipStr, blockType, blockValue, customPage, "")
		}
	} else if countryErr != nil && !strings.Contains(countryErr.Error(), "cannot be found in the database") {
		// Log errors other than "not found"
		log.Printf("WARNING: GeoIP country lookup error for IP %s: %v", ipStr, countryErr)
	}

	// 2. Check ASN (Number and Name)
	if asnRecord != nil && asnErr == nil {
		clientASN := asnRecord.AutonomousSystemNumber
		asnOrg := asnRecord.AutonomousSystemOrganization // Keep org name for potential use

		// Check ASN Number Groups
		for groupName, asnList := range blockedASNs {
			for _, blockedASN := range asnList {
				if clientASN == blockedASN {
					blockType = "asn_number_group"
					blockValue = groupName
					log.Printf("INFO: Blocking IP %s based on %s: %s (ASN: %d - %s)", ipStr, blockType, blockValue, clientASN, asnOrg)
					customPage = asnGroupBlockPages[groupName]
					return cacheAndReturnBlockResult(ipStr, blockType, blockValue, customPage, "")
				}
			}
		}

		// Check ASN Name Groups (Case-Insensitive Substring Match)
		if asnOrg != "" {
			lowerASNOrg := strings.ToLower(asnOrg)

			// Acquire read lock for accessing the global matchers map
			asnNameMatchersMutex.RLock()
			defer asnNameMatchersMutex.RUnlock()

			for groupName, matcher := range asnNameMatchers {
				if matcher == nil { // Check added during previous fix, keep it
					log.Printf("WARNING: Nil matcher found for group: %s during check", groupName)
					continue
				}

				// Use Aho-Corasick matcher - protect against any potential panic
				var matches []int
				func() {
					defer func() {
						if r := recover(); r != nil {
							log.Printf("RECOVERED from panic in Aho-Corasick: %v", r)
							matches = nil // Ensure it's empty if we panic
						}
					}()
					// Match against the globally shared matcher (assumed read-safe)
					matches = matcher.Match([]byte(lowerASNOrg))
				}()

				if len(matches) > 0 {
					blockType = "asn_name_group"
					blockValue = groupName
					log.Printf("INFO: Blocking IP %s based on %s: %s (ASN: %d, Org: '%s')", ipStr, blockType, blockValue, clientASN, asnOrg)
					customPage = asnGroupBlockPages[groupName]
					// No need to unlock here, defer handles it
					return cacheAndReturnBlockResult(ipStr, blockType, blockValue, customPage, asnOrg)
				}
			}
			// RUnlock happens via defer
		}
	} else if asnErr != nil && !strings.Contains(asnErr.Error(), "cannot be found in the database") {
		// Log errors other than "not found"
		log.Printf("WARNING: GeoIP ASN lookup error for IP %s: %v", ipStr, asnErr)
	}

	// --- Cache the result before returning --- //
	computedEntry := blockCacheEntry{
		blocked:   false,
		expiresAt: time.Now().Add(ipBlockCacheTTL),
	}
	ipBlockCacheMutex.Lock()
	ipBlockCache[ipStr] = computedEntry
	ipBlockCacheMutex.Unlock()
	return false, "", "", "", "" // Not blocked
}

// Helper function to cache block results
func cacheAndReturnBlockResult(ipStr string, blockType string, blockValue string, customPage string, asnOrgName string) (bool, string, string, string, string) {
	// Create the cache entry
	computedEntry := blockCacheEntry{
		blocked:    true,
		blockType:  blockType,
		blockValue: blockValue,
		customPage: customPage,
		asnOrgName: asnOrgName,
		expiresAt:  time.Now().Add(ipBlockCacheTTL),
	}

	// Use a separate defer+recover to ensure we don't crash the entire server
	// if there's any issue with the cache
	func() {
		defer func() {
			if r := recover(); r != nil {
				log.Printf("RECOVERED from panic while caching result: %v", r)
			}
		}()

		ipBlockCacheMutex.Lock()
		defer ipBlockCacheMutex.Unlock() // Use defer to ensure unlock happens
		ipBlockCache[ipStr] = computedEntry
	}()

	return true, blockType, blockValue, customPage, asnOrgName
}

// buildASNNameMatchers creates Aho-Corasick matchers for faster ASN name checking
func buildASNNameMatchers() {
	// Acquire write lock before modifying the global map
	asnNameMatchersMutex.Lock()
	defer asnNameMatchersMutex.Unlock()

	// Clear any existing matchers first
	asnNameMatchers = make(map[string]*ahocorasick.Matcher)

	for groupName, nameList := range blockedASNNames {
		// Skip if the name list is empty
		if len(nameList) == 0 {
			log.Printf("Skipping matcher build for empty group: %s", groupName)
			continue
		}

		// Convert names to lowercase byte slices for case-insensitive matching
		dict := make([][]byte, 0, len(nameList))
		for _, name := range nameList {
			if name != "" {
				dict = append(dict, []byte(strings.ToLower(name)))
			}
		}

		// Only create a matcher if we have patterns
		if len(dict) > 0 {
			// Use a recovery mechanism in case the matcher creation fails
			func() {
				defer func() {
					if r := recover(); r != nil {
						log.Printf("PANIC while building Aho-Corasick matcher for group %s: %v", groupName, r)
						// Ensure the entry for this group is nil if creation failed
						asnNameMatchers[groupName] = nil
					}
				}()

				// This assignment happens under the write lock
				asnNameMatchers[groupName] = ahocorasick.NewMatcher(dict)
				log.Printf("Built Aho-Corasick matcher for ASN name group: %s (%d patterns)", groupName, len(dict))
			}()
		} else {
			log.Printf("No valid patterns found for ASN name group: %s", groupName)
		}
	}
	// Unlock happens via defer
}

// ReloadGeoIPDatabases closes and reopens the GeoIP database readers
// to load updated database files. Safe to call while the server is running.
func ReloadGeoIPDatabases() {
	// Close existing readers if they're open
	if geoipCountryReader != nil {
		geoipCountryReader.Close()
		geoipCountryReader = nil
	}
	if geoipASNReader != nil {
		geoipASNReader.Close()
		geoipASNReader = nil
	}

	// Re-initialize the readers
	initGeoIP()
	log.Printf("GeoIP databases reloaded")
}

// getRealIP gets the real client IP when behind a reverse proxy
// It checks X-Forwarded-For header first, then falls back to c.IP()
func getRealIP(c *fiber.Ctx) string {
	// Check X-Forwarded-For header first
	if xff := c.Get("X-Forwarded-For"); xff != "" {
		// X-Forwarded-For can contain multiple IPs (client, proxy1, proxy2, ...)
		// The first one is the original client IP
		ips := strings.Split(xff, ",")
		if len(ips) > 0 {
			// Get the first IP and trim whitespace
			clientIP := strings.TrimSpace(ips[0])
			// Validate it's a real IP
			if net.ParseIP(clientIP) != nil {
				log.Printf("Using X-Forwarded-For IP: %s (original: %s)", clientIP, c.IP())
				return clientIP
			}
		}
	}

	// Also check for custom Remote-Addr header that might be set by some proxies
	if remoteAddr := c.Get("$remote_addr"); remoteAddr != "" {
		// Validate it's a real IP
		if net.ParseIP(remoteAddr) != nil {
			log.Printf("Using $remote_addr IP: %s (original: %s)", remoteAddr, c.IP())
			return remoteAddr
		}
	}

	// Fallback to default IP
	return c.IP()
}