package main

import (
	"bufio"
	"context"
	"fmt"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"
)

const (
	repoDir = "cloud-provider-ip-addresses"
	port    = 8080
)

// IPGenerator generates IPs from CIDR ranges lazily
type IPGenerator struct {
	mu              sync.Mutex
	cidrFiles       []string
	currentFile     int
	currentCIDRs    []string
	activeGens      []*hostGenerator
	rng             *rand.Rand
	totalCIDRsCount int
}

type hostGenerator struct {
	network *net.IPNet
	current net.IP
	done    bool
}

func newHostGenerator(cidr string) (*hostGenerator, error) {
	_, network, err := net.ParseCIDR(cidr)
	if err != nil {
		return nil, err
	}

	// Only IPv4
	if network.IP.To4() == nil {
		return nil, fmt.Errorf("not IPv4")
	}

	// Check if multicast
	if network.IP.IsMulticast() {
		return nil, fmt.Errorf("multicast network")
	}

	ones, bits := network.Mask.Size()
	
	hg := &hostGenerator{
		network: network,
		current: make(net.IP, len(network.IP)),
	}
	copy(hg.current, network.IP)

	// For /32, just use the single address
	if ones == bits {
		return hg, nil
	}

	// For other networks, skip network address (start at .1)
	hg.increment()
	
	return hg, nil
}

func (hg *hostGenerator) increment() {
	for i := len(hg.current) - 1; i >= 0; i-- {
		hg.current[i]++
		if hg.current[i] != 0 {
			break
		}
	}
}

func (hg *hostGenerator) next() (string, bool) {
	if hg.done {
		return "", false
	}

	ones, bits := hg.network.Mask.Size()

	// Handle /32 specially
	if ones == bits {
		if !hg.current.Equal(hg.network.IP) {
			hg.done = true
			return "", false
		}
		ip := hg.current.String()
		hg.done = true
		return ip, true
	}

	// Check if we're still in the network
	if !hg.network.Contains(hg.current) {
		hg.done = true
		return "", false
	}

	// Check if this is the broadcast address (last IP in range)
	broadcast := make(net.IP, len(hg.network.IP))
	copy(broadcast, hg.network.IP)
	for i := range broadcast {
		broadcast[i] |= ^hg.network.Mask[i]
	}

	if hg.current.Equal(broadcast) {
		hg.done = true
		return "", false
	}

	// Skip multicast addresses
	if hg.current.IsMulticast() {
		hg.increment()
		return hg.next()
	}

	ip := hg.current.String()
	hg.increment()
	
	return ip, true
}

func newIPGenerator() (*IPGenerator, error) {
	gen := &IPGenerator{
		rng: rand.New(rand.NewSource(time.Now().UnixNano())),
	}

	// Find all IP files
	err := filepath.Walk(repoDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !info.IsDir() && strings.HasSuffix(path, ".txt") && strings.Contains(strings.ToLower(path), "ips") {
			gen.cidrFiles = append(gen.cidrFiles, path)
		}
		return nil
	})

	if err != nil {
		return nil, fmt.Errorf("failed to scan repo directory: %w", err)
	}

	if len(gen.cidrFiles) == 0 {
		return nil, fmt.Errorf("no IP files found in %s", repoDir)
	}

	// Load first batch of CIDRs
	if err := gen.loadNextFile(); err != nil {
		return nil, err
	}

	log.Printf("📁 Found %d IP files", len(gen.cidrFiles))
	log.Printf("📊 Total CIDRs discovered: %d", gen.totalCIDRsCount)

	return gen, nil
}

func (g *IPGenerator) loadNextFile() error {
	if g.currentFile >= len(g.cidrFiles) {
		// Wrap around and reshuffle
		g.currentFile = 0
		g.rng.Shuffle(len(g.cidrFiles), func(i, j int) {
			g.cidrFiles[i], g.cidrFiles[j] = g.cidrFiles[j], g.cidrFiles[i]
		})
	}

	filepath := g.cidrFiles[g.currentFile]
	g.currentFile++

	file, err := os.Open(filepath)
	if err != nil {
		return fmt.Errorf("failed to open %s: %w", filepath, err)
	}
	defer file.Close()

	g.currentCIDRs = g.currentCIDRs[:0] // Clear but keep capacity
	scanner := bufio.NewScanner(file)
	
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}

		fields := strings.Fields(line)
		for _, field := range fields {
			if field != "" {
				// Basic validation
				if strings.Contains(field, "/") || net.ParseIP(field) != nil {
					g.currentCIDRs = append(g.currentCIDRs, field)
					g.totalCIDRsCount++
				}
			}
		}
	}

	if err := scanner.Err(); err != nil {
		return fmt.Errorf("error reading %s: %w", filepath, err)
	}

	// Shuffle CIDRs from this file
	g.rng.Shuffle(len(g.currentCIDRs), func(i, j int) {
		g.currentCIDRs[i], g.currentCIDRs[j] = g.currentCIDRs[j], g.currentCIDRs[i]
	})

	// Initialize generators for this batch
	g.activeGens = make([]*hostGenerator, 0, len(g.currentCIDRs))
	for _, cidr := range g.currentCIDRs {
		// Ensure it has CIDR notation
		if !strings.Contains(cidr, "/") {
			cidr = cidr + "/32"
		}

		gen, err := newHostGenerator(cidr)
		if err != nil {
			// Skip invalid CIDRs silently
			continue
		}
		g.activeGens = append(g.activeGens, gen)
	}

	return nil
}

func (g *IPGenerator) Next() (string, error) {
	g.mu.Lock()
	defer g.mu.Unlock()

	for {
		// If no active generators, load next file
		if len(g.activeGens) == 0 {
			if err := g.loadNextFile(); err != nil {
				return "", fmt.Errorf("failed to load next file: %w", err)
			}
			if len(g.activeGens) == 0 {
				return "", fmt.Errorf("no more IPs available")
			}
		}

		// Pick a random generator
		idx := g.rng.Intn(len(g.activeGens))
		gen := g.activeGens[idx]

		ip, ok := gen.next()
		if !ok {
			// This generator is exhausted, remove it
			g.activeGens = append(g.activeGens[:idx], g.activeGens[idx+1:]...)
			continue
		}

		return ip, nil
	}
}

// Server holds per-consumer generators
type Server struct {
	generators map[string]*IPGenerator
	mu         sync.RWMutex
}

func newServer() *Server {
	return &Server{
		generators: make(map[string]*IPGenerator),
	}
}

func (s *Server) getGenerator(consumer string) (*IPGenerator, error) {
	s.mu.RLock()
	gen, exists := s.generators[consumer]
	s.mu.RUnlock()

	if exists {
		return gen, nil
	}

	// Create new generator for this consumer
	s.mu.Lock()
	defer s.mu.Unlock()

	// Double-check after acquiring write lock
	if gen, exists := s.generators[consumer]; exists {
		return gen, nil
	}

	newGen, err := newIPGenerator()
	if err != nil {
		return nil, err
	}

	s.generators[consumer] = newGen
	log.Printf("🆕 New consumer: %s", consumer)
	
	return newGen, nil
}

func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	consumer := r.RemoteAddr
	if host, _, err := net.SplitHostPort(consumer); err == nil {
		consumer = host
	}

	gen, err := s.getGenerator(consumer)
	if err != nil {
		log.Printf("❌ Failed to get generator for %s: %v", consumer, err)
		http.Error(w, "Internal server error", http.StatusInternalServerError)
		return
	}

	ip, err := gen.Next()
	if err != nil {
		log.Printf("❌ Failed to get IP for %s: %v", consumer, err)
		http.Error(w, "No more IPs available", http.StatusServiceUnavailable)
		return
	}

	log.Printf("📤 Serving IP to %s: %s", consumer, ip)

	w.Header().Set("Content-Type", "text/plain")
	fmt.Fprintf(w, "%s\n", ip)
}

func main() {
	// Check if repo directory exists
	if _, err := os.Stat(repoDir); os.IsNotExist(err) {
		log.Fatalf("❌ Error: Directory '%s' not found", repoDir)
	}

	server := newServer()

	mux := http.NewServeMux()
	mux.HandleFunc("/", server.handleRequest)

	httpServer := &http.Server{
		Addr:         fmt.Sprintf(":%d", port),
		Handler:      mux,
		ReadTimeout:  10 * time.Second,
		WriteTimeout: 10 * time.Second,
		IdleTimeout:  60 * time.Second,
	}

	// Graceful shutdown handling
	go func() {
		sigChan := make(chan os.Signal, 1)
		signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
		<-sigChan

		log.Println("\n🛑 Shutting down gracefully...")

		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()

		if err := httpServer.Shutdown(ctx); err != nil {
			log.Printf("❌ Error during shutdown: %v", err)
		}
	}()

	log.Printf("🌐 HTTP Input Server running on http://localhost:%d", port)
	log.Printf("   Serving individual IPv4 host addresses lazily")
	log.Printf("   In highly mixed random order per consumer")
	log.Printf("   Press Ctrl+C to stop")

	if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		log.Fatalf("❌ Server error: %v", err)
	}

	log.Println("✅ Server stopped cleanly")
}