package middleware

import (
	"net"
	"net/http"

	"github.com/gin-gonic/gin"
	"github.com/oschwald/geoip2-golang"
	"github.com/rs/zerolog/log"
)

// GeoIPMiddleware blocks requests from outside North America
type GeoIPMiddleware struct {
	db *geoip2.Reader
}

// NewGeoIPMiddleware creates a new GeoIP middleware
func NewGeoIPMiddleware(dbPath string) (*GeoIPMiddleware, error) {
	db, err := geoip2.Open(dbPath)
	if err != nil {
		return nil, err
	}
	
	return &GeoIPMiddleware{db: db}, nil
}

// Close closes the GeoIP database
func (g *GeoIPMiddleware) Close() error {
	return g.db.Close()
}

// Middleware returns the Gin middleware function
func (g *GeoIPMiddleware) Middleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		// Get client IP
		clientIP := c.ClientIP()
		
		// Parse IP
		ip := net.ParseIP(clientIP)
		if ip == nil {
			// Invalid IP, block it
			log.Warn().Str("ip", clientIP).Msg("Invalid IP address, blocking request")
			c.JSON(http.StatusForbidden, gin.H{"error": "Access denied"})
			c.Abort()
			return
		}
		
		// Skip for private/local IPs
		if ip.IsPrivate() || ip.IsLoopback() || ip.IsLinkLocalUnicast() {
			c.Next()
			return
		}
		
		// Look up country
		record, err := g.db.Country(ip)
		if err != nil {
			log.Warn().Str("ip", clientIP).Err(err).Msg("Failed to lookup country, blocking request")
			c.JSON(http.StatusForbidden, gin.H{"error": "Access denied"})
			c.Abort()
			return
		}
		
		// Check if country is in North America
		countryCode := record.Country.IsoCode
		if !g.isNorthAmericanCountry(countryCode) {
			log.Info().Str("ip", clientIP).Str("country", countryCode).Msg("Blocking request from outside North America")
			c.JSON(http.StatusForbidden, gin.H{"error": "Access denied - region not supported"})
			c.Abort()
			return
		}
		
		c.Next()
	}
}

// isNorthAmericanCountry checks if the country code is from North America
func (g *GeoIPMiddleware) isNorthAmericanCountry(countryCode string) bool {
	northAmericanCountries := map[string]bool{
		"US": true, // United States
		"CA": true, // Canada
		"MX": true, // Mexico
		"GT": true, // Guatemala
		"BZ": true, // Belize
		"SV": true, // El Salvador
		"HN": true, // Honduras
		"NI": true, // Nicaragua
		"CR": true, // Costa Rica
		"PA": true, // Panama
	}
	
	return northAmericanCountries[countryCode]
}