package main

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

const (
	repoDir     = "cloud-provider-ip-addresses"
	port        = 8080
	stateDir    = "progress_state"
	saveInterval = 30 * time.Second
)

// GeneratorState represents the serializable state of a generator
type GeneratorState struct {
	CurrentFile     int                `json:"current_file"`
	CurrentCIDRs    []string           `json:"current_cidrs"`
	ActiveGenStates []HostGenState     `json:"active_gen_states"`
	CIDRFiles       []string           `json:"cidr_files"`
}

type HostGenState struct {
	CIDR    string `json:"cidr"`
	Current string `json:"current"`
	Done    bool   `json:"done"`
}

// 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
	consumer        string
	dirty           bool
}

type hostGenerator struct {
	cidr    string
	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{
		cidr:    cidr,
		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 (hg *hostGenerator) getState() HostGenState {
	return HostGenState{
		CIDR:    hg.cidr,
		Current: hg.current.String(),
		Done:    hg.done,
	}
}

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

	// Try to load existing state
	if err := gen.loadState(); err == nil {
		log.Printf("📂 Loaded saved state for consumer: %s", consumer)
		return gen, nil
	}

	// No saved state, initialize fresh
	// 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("🆕 New generator for %s: %d IP files, %d CIDRs", consumer, len(gen.cidrFiles), gen.totalCIDRsCount)
	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)
	}

	g.dirty = true
	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:]...)
			g.dirty = true
			continue
		}

		g.dirty = true
		return ip, nil
	}
}

func (g *IPGenerator) getState() GeneratorState {
	g.mu.Lock()
	defer g.mu.Unlock()

	activeStates := make([]HostGenState, len(g.activeGens))
	for i, gen := range g.activeGens {
		activeStates[i] = gen.getState()
	}

	return GeneratorState{
		CurrentFile:     g.currentFile,
		CurrentCIDRs:    g.currentCIDRs,
		ActiveGenStates: activeStates,
		CIDRFiles:       g.cidrFiles,
	}
}

func (g *IPGenerator) saveState() error {
	if !g.dirty {
		return nil
	}

	state := g.getState()

	// Ensure state directory exists
	if err := os.MkdirAll(stateDir, 0755); err != nil {
		return fmt.Errorf("failed to create state directory: %w", err)
	}

	// Use consumer as filename (sanitize for filesystem)
	filename := strings.ReplaceAll(g.consumer, ":", "_")
	filename = strings.ReplaceAll(filename, "/", "_")
	filepath := filepath.Join(stateDir, filename+".json")

	// Write to temp file first, then rename for atomic write
	tempPath := filepath + ".tmp"
	file, err := os.Create(tempPath)
	if err != nil {
		return fmt.Errorf("failed to create temp state file: %w", err)
	}

	encoder := json.NewEncoder(file)
	encoder.SetIndent("", "  ")
	if err := encoder.Encode(state); err != nil {
		file.Close()
		os.Remove(tempPath)
		return fmt.Errorf("failed to encode state: %w", err)
	}

	if err := file.Close(); err != nil {
		os.Remove(tempPath)
		return fmt.Errorf("failed to close temp state file: %w", err)
	}

	if err := os.Rename(tempPath, filepath); err != nil {
		os.Remove(tempPath)
		return fmt.Errorf("failed to rename state file: %w", err)
	}

	g.dirty = false
	return nil
}

func (g *IPGenerator) loadState() error {
	// Use consumer as filename (sanitize for filesystem)
	filename := strings.ReplaceAll(g.consumer, ":", "_")
	filename = strings.ReplaceAll(filename, "/", "_")
	filepath := filepath.Join(stateDir, filename+".json")

	file, err := os.Open(filepath)
	if err != nil {
		return err
	}
	defer file.Close()

	var state GeneratorState
	if err := json.NewDecoder(file).Decode(&state); err != nil {
		return fmt.Errorf("failed to decode state: %w", err)
	}

	// Restore state
	g.cidrFiles = state.CIDRFiles
	g.currentFile = state.CurrentFile
	g.currentCIDRs = state.CurrentCIDRs
	g.totalCIDRsCount = len(state.CurrentCIDRs)

	// Rebuild active generators from state
	g.activeGens = make([]*hostGenerator, 0, len(state.ActiveGenStates))
	for _, genState := range state.ActiveGenStates {
		gen, err := newHostGenerator(genState.CIDR)
		if err != nil {
			continue
		}

		// Restore current IP position
		gen.current = net.ParseIP(genState.Current)
		if gen.current == nil {
			continue
		}
		gen.done = genState.Done

		g.activeGens = append(g.activeGens, gen)
	}

	return nil
}

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

func newServer() *Server {
	s := &Server{
		generators: make(map[string]*IPGenerator),
		stopSaver:  make(chan struct{}),
	}
	
	s.startPeriodicSaver()
	return s
}

func (s *Server) startPeriodicSaver() {
	s.wg.Add(1)
	go func() {
		defer s.wg.Done()
		ticker := time.NewTicker(saveInterval)
		defer ticker.Stop()

		for {
			select {
			case <-ticker.C:
				s.saveAllStates()
			case <-s.stopSaver:
				// Final save before shutdown
				s.saveAllStates()
				return
			}
		}
	}()
}

func (s *Server) saveAllStates() {
	s.mu.RLock()
	generators := make([]*IPGenerator, 0, len(s.generators))
	for _, gen := range s.generators {
		generators = append(generators, gen)
	}
	s.mu.RUnlock()

	for _, gen := range generators {
		if err := gen.saveState(); err != nil {
			log.Printf("⚠️  Failed to save state for %s: %v", gen.consumer, err)
		}
	}
}

func (s *Server) shutdown() {
	close(s.stopSaver)
	s.wg.Wait()
	log.Println("💾 All states saved")
}

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(consumer)
	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)
}

type ConsumerStatus struct {
	Consumer        string `json:"consumer"`
	CurrentFile     int    `json:"current_file"`
	TotalFiles      int    `json:"total_files"`
	ActiveCIDRs     int    `json:"active_cidrs"`
	TotalCIDRs      int    `json:"total_cidrs_discovered"`
	CurrentFilePath string `json:"current_file_path,omitempty"`
}

type StatusResponse struct {
	TotalConsumers int              `json:"total_consumers"`
	Consumers      []ConsumerStatus `json:"consumers"`
	StateDirectory string           `json:"state_directory"`
	SaveInterval   string           `json:"save_interval"`
}

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

	s.mu.RLock()
	defer s.mu.RUnlock()

	response := StatusResponse{
		TotalConsumers: len(s.generators),
		Consumers:      make([]ConsumerStatus, 0, len(s.generators)),
		StateDirectory: stateDir,
		SaveInterval:   saveInterval.String(),
	}

	for consumer, gen := range s.generators {
		gen.mu.Lock()
		status := ConsumerStatus{
			Consumer:    consumer,
			CurrentFile: gen.currentFile,
			TotalFiles:  len(gen.cidrFiles),
			ActiveCIDRs: len(gen.activeGens),
			TotalCIDRs:  gen.totalCIDRsCount,
		}
		if gen.currentFile > 0 && gen.currentFile <= len(gen.cidrFiles) {
			status.CurrentFilePath = gen.cidrFiles[gen.currentFile-1]
		}
		gen.mu.Unlock()

		response.Consumers = append(response.Consumers, status)
	}

	w.Header().Set("Content-Type", "application/json")
	encoder := json.NewEncoder(w)
	encoder.SetIndent("", "  ")
	if err := encoder.Encode(response); err != nil {
		log.Printf("❌ Failed to encode status response: %v", err)
	}
}

type ExportResponse struct {
	ExportedAt time.Time                  `json:"exported_at"`
	States     map[string]GeneratorState  `json:"states"`
}

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

	// Force save all current states first
	s.saveAllStates()

	s.mu.RLock()
	defer s.mu.RUnlock()

	response := ExportResponse{
		ExportedAt: time.Now(),
		States:     make(map[string]GeneratorState),
	}

	for consumer, gen := range s.generators {
		response.States[consumer] = gen.getState()
	}

	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=state-export-%s.json", 
		time.Now().Format("2006-01-02-150405")))
	
	encoder := json.NewEncoder(w)
	encoder.SetIndent("", "  ")
	if err := encoder.Encode(response); err != nil {
		log.Printf("❌ Failed to encode export response: %v", err)
	}
}

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

	var exportData ExportResponse
	if err := json.NewDecoder(r.Body).Decode(&exportData); err != nil {
		http.Error(w, fmt.Sprintf("Failed to decode import data: %v", err), http.StatusBadRequest)
		return
	}

	// Ensure state directory exists
	if err := os.MkdirAll(stateDir, 0755); err != nil {
		http.Error(w, fmt.Sprintf("Failed to create state directory: %v", err), http.StatusInternalServerError)
		return
	}

	imported := 0
	failed := 0

	for consumer, state := range exportData.States {
		// Sanitize consumer name for filename
		filename := strings.ReplaceAll(consumer, ":", "_")
		filename = strings.ReplaceAll(filename, "/", "_")
		filepath := filepath.Join(stateDir, filename+".json")

		file, err := os.Create(filepath)
		if err != nil {
			log.Printf("⚠️  Failed to create state file for %s: %v", consumer, err)
			failed++
			continue
		}

		encoder := json.NewEncoder(file)
		encoder.SetIndent("", "  ")
		if err := encoder.Encode(state); err != nil {
			file.Close()
			log.Printf("⚠️  Failed to encode state for %s: %v", consumer, err)
			failed++
			continue
		}

		file.Close()
		imported++
	}

	response := map[string]interface{}{
		"imported": imported,
		"failed":   failed,
		"total":    len(exportData.States),
		"message":  fmt.Sprintf("Successfully imported %d consumer states", imported),
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(response)
	
	log.Printf("📥 Imported %d consumer states (%d failed)", imported, failed)
}

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)
	mux.HandleFunc("/status", server.handleStatus)
	mux.HandleFunc("/export", server.handleExport)
	mux.HandleFunc("/import", server.handleImport)

	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()

		// Stop periodic saver and save final state
		server.shutdown()

		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("   💾 Progress saved every %v to '%s' directory", saveInterval, stateDir)
	log.Printf("   📊 Status endpoint: http://localhost:%d/status", port)
	log.Printf("   📤 Export endpoint: http://localhost:%d/export", port)
	log.Printf("   📥 Import endpoint: http://localhost:%d/import (POST)", port)
	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")
}