package caddyhole

import (
	"fmt"
	"io"
	"math/rand"
	"net"
	"net/http"
	"os"
	"strings"

	"github.com/caddyserver/caddy/v2"
	"github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
	"github.com/caddyserver/caddy/v2/caddyconfig/httpcaddyfile"
	"github.com/caddyserver/caddy/v2/modules/caddyhttp"
	"github.com/oschwald/geoip2-golang"
)

func init() {
	caddy.RegisterModule(CaddyHole{})
	httpcaddyfile.RegisterHandlerDirective("caddyhole", parseCaddyfile)
}

// CaddyHole implements a Caddy module that blocks clients based on country
// and sends crawlers/bots random data from /dev/random.
type CaddyHole struct {
	// Path to the GeoIP2 database file
	DatabasePath string `json:"database_path,omitempty"`

	// List of country ISO codes to block
	BlockedCountries []string `json:"blocked_countries,omitempty"`

	// Minimum bytes to send to bots (default: 1MB)
	MinBotBytes int64 `json:"min_bot_bytes,omitempty"`

	// Maximum bytes to send to bots (default: 100MB)
	MaxBotBytes int64 `json:"max_bot_bytes,omitempty"`

	db *geoip2.Reader
}

// CaddyModule returns the Caddy module information.
func (CaddyHole) CaddyModule() caddy.ModuleInfo {
	return caddy.ModuleInfo{
		ID:  "http.handlers.caddyhole",
		New: func() caddy.Module { return new(CaddyHole) },
	}
}

// Provision sets up the CaddyHole module.
func (c *CaddyHole) Provision(_ caddy.Context) error {
	// Set defaults
	if c.MinBotBytes == 0 {
		c.MinBotBytes = 1024 * 1024 // 1MB
	}
	if c.MaxBotBytes == 0 {
		c.MaxBotBytes = 100 * 1024 * 1024 // 100MB
	}

	// Open GeoIP2 database if path is provided
	if c.DatabasePath != "" {
		db, err := geoip2.Open(c.DatabasePath)
		if err != nil {
			return fmt.Errorf("failed to open GeoIP2 database: %v", err)
		}
		c.db = db
	}

	return nil
}

// Cleanup closes the GeoIP2 database.
func (c *CaddyHole) Cleanup() error {
	if c.db != nil {
		return c.db.Close()
	}
	return nil
}

// ServeHTTP implements caddyhttp.MiddlewareHandler.
func (c *CaddyHole) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyhttp.Handler) error {
	// Check if the request is from a bot/crawler
	if c.isBot(r) {
		c.feedBot(w)
		return nil
	}

	// Check if the request should be blocked based on country
	if c.shouldBlock(r) {
		w.WriteHeader(http.StatusForbidden)
		w.Write([]byte("Access denied"))
		return nil
	}

	// Continue to the next handler
	return next.ServeHTTP(w, r)
}

// isBot checks if the request is from a bot/crawler based on User-Agent.
func (c *CaddyHole) isBot(r *http.Request) bool {
	userAgent := strings.ToLower(r.Header.Get("User-Agent"))

	// Common bot/crawler identifiers
	botSignatures := []string{
		"bot", "crawler", "spider", "scraper", "curl", "wget",
		"python-requests", "python-urllib", "go-http-client",
		"java", "perl", "ruby", "php", "http_request",
	}

	for _, sig := range botSignatures {
		if strings.Contains(userAgent, sig) {
			return true
		}
	}

	return false
}

// feedBot sends random data from /dev/random to the bot.
func (c *CaddyHole) feedBot(w http.ResponseWriter) {
	// Calculate random amount of bytes to send
	bytesToSend := c.MinBotBytes
	if c.MaxBotBytes > c.MinBotBytes {
		bytesToSend += rand.Int63n(c.MaxBotBytes - c.MinBotBytes)
	}

	// Open /dev/random
	devRandom, err := os.Open("/dev/random")
	if err != nil {
		// Fallback to /dev/urandom if /dev/random is not available
		devRandom, err = os.Open("/dev/urandom")
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
	}
	defer devRandom.Close()

	// Set headers to make it look like a legitimate response
	w.Header().Set("Content-Type", "application/octet-stream")
	w.Header().Set("Content-Length", fmt.Sprintf("%d", bytesToSend))
	w.WriteHeader(http.StatusOK)

	// Copy random data to the response
	io.CopyN(w, devRandom, bytesToSend)
}

// shouldBlock checks if the request should be blocked based on country.
func (c *CaddyHole) shouldBlock(r *http.Request) bool {
	if c.db == nil || len(c.BlockedCountries) == 0 {
		return false
	}

	// Get the client IP address
	ip := c.getClientIP(r)
	if ip == nil {
		return false
	}

	// Look up the country for the IP
	record, err := c.db.Country(ip)
	if err != nil {
		return false
	}

	// Check if the country is in the blocked list
	for _, blocked := range c.BlockedCountries {
		if strings.EqualFold(record.Country.IsoCode, blocked) {
			return true
		}
	}

	return false
}

// getClientIP extracts the client IP address from the request.
func (c *CaddyHole) getClientIP(r *http.Request) net.IP {
	// Check X-Forwarded-For header first
	xff := r.Header.Get("X-Forwarded-For")
	if xff != "" {
		ips := strings.Split(xff, ",")
		if len(ips) > 0 {
			ipStr := strings.TrimSpace(ips[0])
			if ip := net.ParseIP(ipStr); ip != nil {
				return ip
			}
		}
	}

	// Check X-Real-IP header
	xri := r.Header.Get("X-Real-IP")
	if xri != "" {
		if ip := net.ParseIP(xri); ip != nil {
			return ip
		}
	}

	// Fall back to RemoteAddr
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return nil
	}

	return net.ParseIP(host)
}

// UnmarshalCaddyfile implements caddyfile.Unmarshaler.
func (c *CaddyHole) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
	for d.Next() {
		for d.NextBlock(0) {
			switch d.Val() {
			case "database":
				if !d.NextArg() {
					return d.ArgErr()
				}
				c.DatabasePath = d.Val()

			case "block_countries":
				c.BlockedCountries = d.RemainingArgs()
				if len(c.BlockedCountries) == 0 {
					return d.ArgErr()
				}

			case "min_bot_bytes":
				if !d.NextArg() {
					return d.ArgErr()
				}
				var err error
				_, err = fmt.Sscanf(d.Val(), "%d", &c.MinBotBytes)
				if err != nil {
					return d.Errf("invalid min_bot_bytes value: %v", err)
				}

			case "max_bot_bytes":
				if !d.NextArg() {
					return d.ArgErr()
				}
				var err error
				_, err = fmt.Sscanf(d.Val(), "%d", &c.MaxBotBytes)
				if err != nil {
					return d.Errf("invalid max_bot_bytes value: %v", err)
				}

			default:
				return d.Errf("unrecognized subdirective: %s", d.Val())
			}
		}
	}
	return nil
}

// parseCaddyfile unmarshals tokens from h into a new Middleware.
func parseCaddyfile(h httpcaddyfile.Helper) (caddyhttp.MiddlewareHandler, error) {
	var c CaddyHole
	err := c.UnmarshalCaddyfile(h.Dispenser)
	return &c, err
}

// Interface guards
var (
	_ caddy.Provisioner           = (*CaddyHole)(nil)
	_ caddy.CleanerUpper          = (*CaddyHole)(nil)
	_ caddyhttp.MiddlewareHandler = (*CaddyHole)(nil)
	_ caddyfile.Unmarshaler       = (*CaddyHole)(nil)
)