gontrol/main.go

package main

import (
	"context"
	"os"
	"strings"
	"time"

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

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

	"os/signal"

	_ "github.com/go-sql-driver/mysql"
	"github.com/gorilla/websocket"
)

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

const (
	keyServerAddr = "serverAddr"
	dbHost = "172.17.0.2"
	dbUser = "root"
	dbPort = 3306
	dbPassword = "root"
	dbName = "gomatic"
)

type Agent struct {
	agentId int
	agentName string
	initialContact string
	lastContact string
}

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

func initDB () {
	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 fetchAgents(w http.ResponseWriter, r *http.Request) {
	agents, _ := getAgents(db)

	log.Printf("%s",agents)
	tmpl.ExecuteTemplate(w, "agentList", agents)
}

func getAgents(dbPointer *sql.DB) ([]Agent, error) {
	query := "Select agentId, agentName, initialContact, lastContact from agents"
	rows, err := dbPointer.Query(query)

	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var agents []Agent
	for rows.Next() {
		var agent Agent

		rowErr := rows.Scan(&agent.agentId, &agent.agentName, &agent.initialContact, &agent.lastContact)
		if rowErr != nil {
			return nil, err
		}
		agents = append(agents, agent)
	}

	if err = rows.Err(); err != nil {
		return nil, err
	}

	return agents, nil
	// return agents, rows.Err()
}

func getRoot(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	hasFirst := r.URL.Query().Has("first")
	first := r.URL.Query().Get("first")
	hasSecond := r.URL.Query().Has("second")
	second := r.URL.Query().Get("second")

	body, err := ioutil.ReadAll(r.Body)

	if err != nil {
		fmt.Printf("Could not read body: %s\n", err)
	}

	fmt.Printf("%s: got / request. first(%t)=%s, second(%t)=%s body:\n%s\n",
		ctx.Value(keyServerAddr),
		hasFirst, first,
		hasSecond, second,
		body,
	)

	// fmt.Printf("%s: got / request\n", ctx.Value(keyServerAddr))
	io.WriteString(w, "This is my website!\n")
}

func getHello(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	fmt.Printf("%s: go /hello request\n", ctx.Value(keyServerAddr))

	myName := r.PostFormValue("myName")
	if myName == "" {
		w.Header().Set("x-missing-field", "myName")
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	io.WriteString(w, fmt.Sprintf("Hello, %s!\n", myName))
}

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

func agentHandler(w http.ResponseWriter, r *http.Request) {
	parts := strings.Split(strings.TrimPrefix(r.URL.Path, "/agents/"), "/")
	if len (parts) < 1 || parts[0] == "" {
		http.Error(w, "Agent ID required", http.StatusBadRequest)
		return
	}

	agentId := parts[0]

	switch r.Method {
		case http.MethodGet:
			getAgent(w, r, agentId)
		// case http.MethodPost:
		// 	createAgent(w, r, agentId)
		// case http.MethodPut:
		// 	updateAgent(w, r, agentId)
		// case http.MethodDelete:
		// 	deleteAgent(w, r, agentId)
		default:
			http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		}
}

type webSocketHandler struct {
	upgrader websocket.Upgrader
}

func getAgent(w http.ResponseWriter, r *http.Request, agentId string) {
	query := "Select agentId, agentName, initialContact, lastContact from agents where agentId = ?"
	var agent Agent
	err := db.QueryRow(query, agentId).Scan(&agentId, &agent.agentName, &agent.initialContact, &agent.lastContact)
	if err == sql.ErrNoRows {
		http.Error(w, "Agent not found", http.StatusNotFound)
		return
	} else if err != nil {
		http.Error(w, "Failed to fetch agent", http.StatusInternalServerError)
		return
	}

	// return agent, nil

}

func (wsh webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	c, err := wsh.upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("Error %s when upgrading connection to websocket", err)
		return
	}
	defer c.Close()

	for {
		mt, message, err := c.ReadMessage()
		if err != nil {
			log.Printf("Error reading: message: %s", err)
		}

		log.Printf("Received message: %s", message)
		if err = c.WriteMessage(mt, message); err !=nil {
			log.Printf("Error writing the message: %s", err)
			break
		}
	}
}

func main() {

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

	var wg sync.WaitGroup

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

	webSocketMux := http.NewServeMux()
	webSocketMux.Handle("/data", webSocketHandler)
	websocketServer := &http.Server{
		Addr: ":5555",
		Handler: webSocketMux,
	}

	webMux := http.NewServeMux()
	webMux.HandleFunc("/", getHomepage)
	webMux.HandleFunc("/index", getRoot)
	webMux.HandleFunc("/hello", getHello)
	webMux.HandleFunc("/agents", fetchAgents)
	// webMux.HandleFunc("/newagentform", getAgentForm)
	// webMux.HandleFunc("/getagentupdateform/{agentId}", getAgentUpdateForm)
	webMux.HandleFunc("/agents/{agentId}", agentHandler)
	// webMux.HandleFunc


	initDB()
	defer db.Close()

	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")
}