gontrol/main.go

package main

import (
	"context"
	// "encoding/json"
	"os"
	"strings"
	"time"

	// "errors"
	"fmt"
	// "io"
	// "io/ioutil"

	// "net"
	"database/sql"
	"html/template"
	"log"
	"net/http"
	"sync"
	"syscall"

	"gontrol/src/randomname"
	api "gontrol/src/server/api"
	websocketserver "gontrol/src/server/websocket"

	"os/signal"

	_ "github.com/go-sql-driver/mysql"
	"github.com/kelseyhightower/envconfig"
)

var tmpl *template.Template
var db 	 *sql.DB

type Config struct {
	Database struct {
		Username string `envconfig:"DB_USERNAME"`
		Password string `envconfig:"DB_PASSWORD"`
		Port     int16  `envconfig:"DB_PORT"`
		Name     string `envconfig:"DB_NAME"`
		Host     string `envconfig:"DB_HOST"`
	}
}

func readEnv(cfg *Config) {
	err := envconfig.Process("", cfg)
	if err != nil {
		processError(err)
	}
}

func processError(err error) {
	fmt.Println(err)
	os.Exit(2)
}

type Agent struct {
	AgentID        int    `json:"agentId"`
	AgentName      string `json:"agentName"`
	InitialContact string `json:"initialContact"`
	LastContact    string `json:"lastContact"`
	IPv4Address    string `json:"IPv4Address"`
	Status         string `json:"status"`
}

func init() {
	tmpl, _ = template.ParseGlob("templates/*.html")
}

func initDB (dbUser string, dbPassword string, dbHost string, dbPort int16, dbName string ) {
	var err  error
	dbOpen := fmt.Sprintf("%s:%s@(%s:%d)/%s?parseTime=true", dbUser,dbPassword, dbHost, dbPort, dbName)
	db, err = sql.Open("mysql", dbOpen)
	if err != nil {
		log.Fatal(err)
	}

	if err = db.Ping(); err != nil {
		log.Fatal(err)
	}
	log.Println("Connected to database")
}

func renderTemplate(w http.ResponseWriter, tmpl string, data interface{}) {
	t, err := template.ParseFiles(tmpl)
		if err != nil {
		log.Printf("Failed to load template %s: %v", tmpl, err)
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}
	if err := t.Execute(w, data); err != nil {
		log.Printf("Failed to render template %s: %v", tmpl, err)
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}
}

func agentsHandler(w http.ResponseWriter, r *http.Request) {
	parts := strings.Split(strings.TrimPrefix(r.URL.Path, "/agents/"), "/")
	agentId := ""
	if len(parts) > 0 && parts[0] != "" {
		agentId = parts[0]
	}

	switch r.Method {
	case http.MethodDelete:
		// deleteAgent(w, r, agentId)
		api.DeleteAgent(db, w, r, agentId)
		listAgents(w,r)
		// renderTemplate(w, "templates/partials/agent_list.html", agents)
	case http.MethodGet:
		if agentId == "" {
			listAgents(w, r)
		} else {
			// getAgent(w, r, agentId)
			agent, _ := api.GetAgent(db, w, r, agentId)
			renderTemplate(w, "templates/partials/agent_detail.html", agent)
		}
	case http.MethodPost:
		// createAgent(w, r)
		api.CreateAgent(db, w, r)
		listAgents(w, r)
	case http.MethodPut:
		// updateAgent(w, r, agentId)
		api.UpdateAgent(db, w, r, agentId)
		listAgents(w, r)
	default:
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
	}
}

func getHomepage(w http.ResponseWriter, r *http.Request) {
	tmpl.ExecuteTemplate(w, "index.html", nil)
}

func listAgents(w http.ResponseWriter, r *http.Request) {
	// agents, err := getAgents()
	agents, err := api.GetAgents(db)
	if err != nil {
		http.Error(w, "Failed to fetch agents", http.StatusInternalServerError)
		return
	}

	renderTemplate(w, "templates/partials/agent_list.html", agents)
}

func getAgentNames(w http.ResponseWriter, r *http.Request) {
	api.GetAgentNames(db, w, r)
	return
}


func main() {

	var cfg Config
	readEnv(&cfg)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	var wg sync.WaitGroup

	// webSocketHandler := webSocketHandler {
	// 	upgrader: websocket.Upgrader{},
	// }

	websocketServer := websocketserver.Server()


	webMux := http.NewServeMux()
	webMux.HandleFunc("/", getHomepage)
	// webMux.HandleFunc("/index", getRoot)
	// webMux.HandleFunc("/hello", getHello)
	// webMux.HandleFunc("/agents", fetchAgents)
	webMux.HandleFunc("/agents", agentsHandler)
	webMux.HandleFunc("/agentNames", getAgentNames)
	webMux.HandleFunc("/agents/{agentId}", agentsHandler)

	initDB (cfg.Database.Username, cfg.Database.Password, cfg.Database.Host, cfg.Database.Port, cfg.Database.Name)
	defer db.Close()
	name := randomname.GenerateRandomName()
	fmt.Println(name)

	webServer := &http.Server {
		Addr: ":3333",
		Handler: webMux,
		// BaseContext: func(l net.Listener) context.Context {
		// 	ctx = context.WithValue(ctx, keyServerAddr, l.Addr().String())
		// 	return ctx
		// },
	}

	wg.Add(1)
	go func() {
		defer wg.Done()
		log.Println("Websocket server is running on port 5555")
		if err := websocketServer.ListenAndServe(); err != http.ErrServerClosed {
			log.Fatalf("Websocket server failed: %s", err)
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		log.Println("Web server is running on port 3333")

		if err := webServer.ListenAndServe(); err != http.ErrServerClosed {
			log.Fatalf("Web server failed: %s", err)
		}
	}()

	shutdownCh := make(chan os.Signal, 1)
	signal.Notify(shutdownCh, os.Interrupt, syscall.SIGTERM)

	<-shutdownCh
	log.Println("Shutdown signal received")

	shutdownCtx, shutdownCancel := context.WithTimeout(ctx, 10*time.Second)
	defer shutdownCancel()

	if err := websocketServer.Shutdown(shutdownCtx); err != nil {
		log.Printf("error shutting down websocket server: %s", err)
	}

	if err := webServer.Shutdown(shutdownCtx); err != nil {
		log.Printf("Error shutting down web server: %s", err)
	}

	wg.Wait()
	log.Println("All servers stopped")
}