// main.go

package main

import (
	"log"
	"net/http"
	"net/http/httputil"
)

func main() {
	// Initialize all servers as healthy for start-up
	for i := range HealthStatus {
		HealthStatus[i] = true
	}

	// Start health check as a go routine in the background
	go RunHealthCheck()

	// Set up the single reverse proxy with dynamic routing
	proxy := httputil.NewSingleHostReverseProxy(nil)
	// Reverse Proxy Struct requires a director function so 
    // I made my own custom Director called dynamicDirector
	proxy.Director = dynamicDirector

	// Runs the load balancer server
	http.HandleFunc("/", func(res http.ResponseWriter, req *http.Request) {
		proxy.ServeHTTP(res, req)
	})
	log.Fatal(http.ListenAndServe(":8000", nil))
}

// The Director function is a request transformer. 
// When a request comes to the reverse proxy, 
// Director modifies it so that the request looks as if it 
// was originally intended for the backend server. 
// After the Director function makes its modifications, 
// the proxy forwards the request to the appropriate target server.
func dynamicDirector(req *http.Request) {
	// targetURL := getNextServer()
	targetURL := GetNextHealthyServer()
	req.URL.Scheme = targetURL.Scheme
	req.URL.Host = targetURL.Host
	req.URL.Path = targetURL.Path
}

// healthChecker.go

package main

import (
	"github.com/go-co-op/gocron"
	"log"
	"net/http"
	"net/url"
	"sync"
	"time"
)

var (
	ServerURLs = []*url.URL{
		ParseURL("http://127.0.0.1:5000/"),
		ParseURL("http://127.0.0.1:5001/"),
		ParseURL("http://127.0.0.1:5002/"),
		ParseURL("http://127.0.0.1:5003/"),
		ParseURL("http://127.0.0.1:5004/"),
	}
	HealthStatus = make([]bool, len(ServerURLs))
	Mu           sync.RWMutex
)

// Functions to perform health checks on servers

// Parses a string URL and returns a pointer to a url object which is required by rev proxy's Director func
func ParseURL(urlStr string) *url.URL {
	u, _ := url.Parse(urlStr)
	return u
}

// Go routine to run health checks on a separate thread
func RunHealthCheck() {
	// Utilizes cron jobs to run checkServerHealth every 2 seconds
	// Here we create a Cron Scheduler Object from the gocron library
	s := gocron.NewScheduler(time.Local)

	// We loop through every server URL and check each server's health
	for idx, server := range ServerURLs {
		// We want the scheduler to perform a given function, every 2 seconds
		// The Do method takes a function as an argument which is where we will put checkServerHealth
		s.Every(2).Second().Do(func() {
			// We run the function and then log the server's health
			CheckServerHealth(idx, server)
			if HealthStatus[idx] {
				log.Printf("'%s' is healthy!", server)
			} else {
				log.Printf("'%s' is not healthy!", server)
			}
		})
	}
	// cmd to start the func, use "<-" syntax if the function has a return value
	s.StartAsync()
}

// Pings each server to assess health status, updating shared state accordingly
func CheckServerHealth(index int, server *url.URL) {
	// We check server health by pinging the server and checking the response
	resp, err := http.Get(server.String())

	// Acquire a write lock to access HealthStatus atomically
	Mu.Lock()
	// Deferring delays the unlock until the function returns
	defer Mu.Unlock()

	// If we don't get a 200, it means the server is offline, set the health status to false
	if err != nil || resp.StatusCode != http.StatusOK {
		HealthStatus[index] = false
		// If we do get a 200, it means the server is online, set the health status to true
	} else {
		HealthStatus[index] = true
	}
	if resp != nil {
		resp.Body.Close()
	}
}

func GetNextHealthyServer() *url.URL {
	// Acquire a read lock to access HealthStatus atomically
	Mu.RLock()
	// Deferring delays the unlock until the function returns
	defer Mu.RUnlock()

	// Loop to find the next healthy server
	for i := 0; i < len(ServerURLs); i++ {
		if HealthStatus[i] {
			return ServerURLs[i]
		}
	}
	// Fallback: if all servers are unhealthy, return a default server or error
	log.Printf("Error")
	return nil
}

# server.py

import sys
from flask import Flask
app = Flask(__name__)

serverName = sys.argv[1]

@app.route('/')
def hello():
    return "Hello from " + serverName + " !"

if __name__ == '__main__':
    app.run(port=int(sys.argv[2]))