package websocketserver

import (
	"net/http"
	"log"
	"sync"
	"github.com/gorilla/websocket"

)

type webSocketHandler struct {
	upgrader websocket.Upgrader
}

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

var agentSockets = make(map[string]*websocket.Conn)
var agentSocketsMutex sync.Mutex

var executeCommand http.HandlerFunc = func(w http.ResponseWriter, r *http.Request){
	err := r.ParseForm()
	if err != nil {
		http.Error(w, "Invalid form data", http.StatusBadRequest)
	}

	agentName := r.FormValue("agentName")
	command := r.FormValue("command")

	agentSocketsMutex.Lock()
	conn, ok := agentSockets[agentName]
	agentSocketsMutex.Unlock()

	if !ok {
		http.Error(w, "Agent not connected", http.StatusNotFound)
		return
	}

	err = conn.WriteMessage(websocket.TextMessage, []byte(command))
	if err != nil {
		http.Error(w, "Failed to send command to the agent", http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusOK)
	w.Write([]byte("Command sent successfully"))
}

func Server() (*http.Server) {


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

	return websocketServer
}