package proxy

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"sync"
	"sync/atomic"
	"time"
)

// Backend represents a target server with its health state
type Backend struct {
	URL   *url.URL
	Alive bool
	mu    sync.RWMutex
}

func (b *Backend) SetAlive(alive bool) {
	b.mu.Lock()
	b.Alive = alive
	b.mu.Unlock()
}

func (b *Backend) IsAlive() bool {
	b.mu.RLock()
	alive := b.Alive
	b.mu.RUnlock()
	return alive
}

// LoadBalancer holds a slice of backends and a counter for Round-Robin
type LoadBalancer struct {
	backends []*Backend
	current  uint64
}

// Next returns the next healthy backend URL using Round-Robin
func (lb *LoadBalancer) Next() *url.URL {
	// Loop at most len(backends) times to find a healthy one
	for i := 0; i < len(lb.backends); i++ {
		next := atomic.AddUint64(&lb.current, 1)
		idx := next % uint64(len(lb.backends))
		b := lb.backends[idx]

		if b.IsAlive() {
			return b.URL
		}
	}
	return nil // No healthy backends available
}

// StartHealthCheck begins a background routine to ping backends
func (lb *LoadBalancer) StartHealthCheck() {
	go func() {
		for {
			for _, b := range lb.backends {
				go func(backend *Backend) {
					// TCP Dial is a safe, universal health check for HTTP servers
					host := backend.URL.Host
					if backend.URL.Port() == "" {
						if backend.URL.Scheme == "https" {
							host += ":443"
						} else {
							host += ":80"
						}
					}
					
					conn, err := net.DialTimeout("tcp", host, 2*time.Second)
					if err != nil {
						if backend.IsAlive() {
							log.Printf("[HealthCheck] Backend %s went DOWN: %v", backend.URL.String(), err)
							backend.SetAlive(false)
						}
					} else {
						conn.Close()
						if !backend.IsAlive() {
							log.Printf("[HealthCheck] Backend %s came UP", backend.URL.String())
							backend.SetAlive(true)
						}
					}
				}(b)
			}
			time.Sleep(5 * time.Second) // Check every 5 seconds
		}
	}()
}

// NewLoadBalancedHTTPProxy creates a ReverseProxy that load balances across multiple backends
func NewLoadBalancedHTTPProxy(backends []string) (*httputil.ReverseProxy, error) {
	if len(backends) == 0 {
		return nil, fmt.Errorf("no backends provided")
	}

	lb := &LoadBalancer{}
	for _, b := range backends {
		u, err := url.Parse(b)
		if err != nil {
			return nil, fmt.Errorf("invalid backend url %s: %w", b, err)
		}
		// Initially mark all as alive
		lb.backends = append(lb.backends, &Backend{URL: u, Alive: true})
	}

	// Start health checking in the background
	lb.StartHealthCheck()

	proxy := &httputil.ReverseProxy{
		Director: func(req *http.Request) {
			target := lb.Next()
			if target == nil {
				// We can't return an error directly from Director, but setting Scheme to empty
				// will cause the RoundTripper to fail, returning 502 Bad Gateway by default.
				return
			}
			req.URL.Scheme = target.Scheme
			req.URL.Host = target.Host
			// Note: We don't overwrite req.URL.Path here, as it's modified by the router
			req.Header.Set("X-Forwarded-Host", req.Header.Get("Host"))
		},
		// Optional custom error handler when backends are unreachable
		ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
			log.Printf("Proxy error: %v", err)
			http.Error(w, "502 Bad Gateway - No healthy backends available", http.StatusBadGateway)
		},
	}

	return proxy, nil
}