package main

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"net/netip"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/redis/go-redis/v9"
)

var (
	checkRequests = prometheus.NewCounter(
		prometheus.CounterOpts{
			Name: "ipcheck_requests_total",
			Help: "Total number of IP check requests",
		},
	)
	checkBlocked = prometheus.NewCounter(
		prometheus.CounterOpts{
			Name: "ipcheck_blocked_total",
			Help: "Total number of blocked IPs",
		},
	)
	checkWhitelist = prometheus.NewCounter(
		prometheus.CounterOpts{
			Name: "ipcheck_whitelisted_total",
			Help: "Total number of whitelisted IPs",
		},
	)
)

func init() {
	prometheus.MustRegister(checkRequests, checkBlocked, checkWhitelist)
}

func handleCheck(w http.ResponseWriter, r *http.Request) {
	checkRequests.Inc()

	ipStr := strings.TrimPrefix(r.URL.Path, "/check/")
	ip, err := netip.ParseAddr(ipStr)
	if err != nil {
		http.Error(w, "invalid IP", http.StatusBadRequest)
		return
	}

	cats := []string{"generic", "test"}
	if q := r.URL.Query().Get("cats"); q != "" {
		cats = strings.Split(q, ",")
	}

	matches, err := checkIP(ip, cats)
	if err != nil {
		http.Error(w, "internal error", http.StatusInternalServerError)
		return
	}

	if len(matches) == 0 {
		wl, _ := rdb.Exists(ctx, "wl:"+ip.String()).Result()
		if wl > 0 {
			checkWhitelist.Inc()
		}
	} else {
		checkBlocked.Inc()
	}

	writeJSON(w, map[string]any{
		"ip":         ipStr,
		"blocked":    len(matches) > 0,
		"categories": matches,
	})
}

// Redis Client und Context
var ctx = context.Background()
var rdb = redis.NewClient(&redis.Options{
	Addr: "10.10.5.249:6379",
})

// Präfix-Cache (pro Kategorie)
type prefixCacheEntry struct {
	prefixes []netip.Prefix
	expireAt time.Time
}

var (
	prefixCache   = map[string]prefixCacheEntry{}
	prefixCacheMu sync.Mutex
)

// Prüfen der IP
func checkIP(ip netip.Addr, cats []string) ([]string, error) {
	// Whitelist
	wl, err := rdb.Exists(ctx, "wl:"+ip.String()).Result()
	if err != nil {
		return nil, err
	}
	if wl > 0 {
		return []string{}, nil
	}

	matches := []string{}
	for _, cat := range cats {
		prefixes, err := loadCategoryPrefixes(cat)
		if err != nil {
			return nil, err
		}

		for _, pfx := range prefixes {
			if pfx.Contains(ip) {
				fmt.Printf("💡 MATCH: %s in %s (%s)\n", ip, cat, pfx)
				matches = append(matches, cat)
				break
			}
		}
	}
	return matches, nil
}

func loadCategoryPrefixes(cat string) ([]netip.Prefix, error) {
	prefixCacheMu.Lock()
	defer prefixCacheMu.Unlock()

	entry, ok := prefixCache[cat]
	if ok && time.Now().Before(entry.expireAt) {
		return entry.prefixes, nil
	}

	// Redis HKEYS holen
	keys, err := rdb.HKeys(ctx, "bl:"+cat).Result()
	if err != nil {
		return nil, err
	}

	var prefixes []netip.Prefix
	for _, k := range keys {
		k = strings.TrimSpace(k) // spaces entfernen!
		pfx, err := netip.ParsePrefix(k)
		if err == nil {
			prefixes = append(prefixes, pfx)
		} else {
			fmt.Printf("⚠️ Ungültiger Prefix in Redis (%s): %s\n", cat, k)
		}
	}

	prefixCache[cat] = prefixCacheEntry{
		prefixes: prefixes,
		expireAt: time.Now().Add(1 * time.Second),
	}

	return prefixes, nil
}

// JSON-Antwort
func writeJSON(w http.ResponseWriter, v any) {
	w.Header().Set("Content-Type", "application/json")
	enc := json.NewEncoder(w)
	enc.SetIndent("", "  ")
	enc.Encode(v)
}

// Main + Beispiel-Setup
func main() {

	if 1 == 1 {
		err := importBlocklists()
		if err != nil {
			fmt.Println("Import-Fehler:", err)
			return
		}

		fmt.Println("Blocklisten-Import abgeschlossen.")
	}

	http.HandleFunc("/check/", handleCheck)
	http.Handle("/metrics", promhttp.Handler())

	fmt.Println("Server läuft auf :8080")
	http.ListenAndServe(":8080", nil)
}

// Tests

func TestCheckIP(t *testing.T) {
	// Setup Redis-Daten
	rdb.FlushDB(ctx)
	rdb.HSet(ctx, "bl:generic", map[string]string{
		"81.232.51.35/32": "1",
		"150.242.0.0/20":  "1",
	})
	rdb.HSet(ctx, "bl:test", map[string]string{
		"203.9.56.0/24": "1",
	})
	rdb.Set(ctx, "wl:8.8.8.8", "1", 0)

	tests := []struct {
		ip       string
		expected []string
	}{
		{"81.232.51.35", []string{"generic"}},
		{"150.242.5.10", []string{"generic"}},
		{"203.9.56.5", []string{"test"}},
		{"8.8.8.8", []string{}}, // Whitelisted
		{"1.1.1.1", []string{}},
	}

	for _, tc := range tests {
		addr := netip.MustParseAddr(tc.ip)
		cats := []string{"generic", "test"}
		res, err := checkIP(addr, cats)
		if err != nil {
			t.Errorf("unexpected error: %v", err)
		}
		if !equalStringSlices(res, tc.expected) {
			t.Errorf("for IP %s: expected %v, got %v", tc.ip, tc.expected, res)
		}
	}
}

func equalStringSlices(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	m := make(map[string]int)
	for _, v := range a {
		m[v]++
	}
	for _, v := range b {
		if m[v] == 0 {
			return false
		}
		m[v]--
	}
	return true
}

// Import

// URL-Liste
var blocklistURLs = map[string]string{
	"firehol": "https://raw.githubusercontent.com/firehol/blocklist-ipsets/master/firehol_level1.netset",
	"bitwire": "https://raw.githubusercontent.com/bitwire-it/ipblocklist/refs/heads/main/ip-list.txt",
	"RU":      "https://ipv64.net/blocklists/countries/ipv64_blocklist_RU.txt",
	"CN":      "https://ipv64.net/blocklists/countries/ipv64_blocklist_CN.txt",
}

// Importer
func importBlocklists() error {
	for category, url := range blocklistURLs {
		fmt.Printf("Lade %s (%s)...\n", category, url)
		resp, err := http.Get(url)
		if err != nil {
			return fmt.Errorf("Fehler beim Laden %s: %v", url, err)
		}
		defer resp.Body.Close()

		if resp.StatusCode != 200 {
			return fmt.Errorf("Fehler beim Laden %s: HTTP %d", url, resp.StatusCode)
		}

		scanner := bufio.NewScanner(resp.Body)
		count := 0

		pipe := rdb.Pipeline()

		for scanner.Scan() {
			line := strings.TrimSpace(scanner.Text())
			if line == "" || strings.HasPrefix(line, "#") {
				continue
			}

			// Prüfen ob IP oder Prefix
			if isValidPrefix(line) {
				pipe.HSet(ctx, "bl:"+category, line, 1)
				count++
			} else {
				fmt.Printf("⚠️ Ungültige Zeile (%s): %s\n", category, line)
			}
		}
		if err := scanner.Err(); err != nil {
			return fmt.Errorf("Lesefehler %s: %v", url, err)
		}

		// Commit Pipeline
		_, err = pipe.Exec(ctx)
		if err != nil {
			return fmt.Errorf("Redis-Fehler %s: %v", category, err)
		}

		fmt.Printf("✅ %d Einträge in Kategorie %s importiert.\n", count, category)
	}

	return nil
}

// Prüft ob eine Zeile ein valides IP-Präfix ist
func isValidPrefix(s string) bool {
	// Wenn es kein / enthält → vermutlich /32 oder /128 annehmen
	if !strings.Contains(s, "/") {
		if ip := net.ParseIP(s); ip != nil {
			if ip.To4() != nil {
				s = s + "/32"
			} else {
				s = s + "/128"
			}
		} else {
			return false
		}
	}

	_, err := netip.ParsePrefix(s)
	return err == nil
}