nsfwapp/backend/ml_setup.go

package main

import (
	"bufio"
	"context"
	"errors"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

type mlPythonSetupStatus struct {
	Required   bool   `json:"required"`
	Running    bool   `json:"running"`
	Done       bool   `json:"done"`
	OK         bool   `json:"ok"`
	Progress   int    `json:"progress"`
	Phase      string `json:"phase"`
	Message    string `json:"message"`
	Error      string `json:"error,omitempty"`
	LogFile    string `json:"logFile,omitempty"`
	StartedAt  string `json:"startedAt,omitempty"`
	UpdatedAt  string `json:"updatedAt,omitempty"`
	FinishedAt string `json:"finishedAt,omitempty"`
}

type mlPythonSetupConfig struct {
	enabled          bool
	requirementsPath string
	venvDir          string
	venvPython       string
	requirementsMark string
	logPath          string
}

var (
	mlSetupMu     sync.RWMutex
	mlSetupRunMu  sync.Mutex
	mlSetupStatus mlPythonSetupStatus
)

func initMLPythonSetupStatus() {
	cfg := mlPythonSetupConfigFromEnv()
	now := time.Now().UTC().Format(time.RFC3339Nano)

	mlSetupMu.Lock()
	defer mlSetupMu.Unlock()

	if cfg.enabled {
		mlSetupStatus = mlPythonSetupStatus{
			Required:  true,
			Running:   false,
			Done:      false,
			OK:        false,
			Progress:  0,
			Phase:     "pending",
			Message:   "ML-Python-Umgebung wartet auf Start.",
			LogFile:   filepath.Base(cfg.logPath),
			UpdatedAt: now,
		}
		return
	}

	mlSetupStatus = mlPythonSetupStatus{
		Required:  false,
		Running:   false,
		Done:      true,
		OK:        true,
		Progress:  100,
		Phase:     "skipped",
		Message:   "ML-Python-Setup nicht erforderlich.",
		UpdatedAt: now,
	}
}

func mlPythonSetupConfigFromEnv() mlPythonSetupConfig {
	if strings.TrimSpace(os.Getenv("NSFWAPP_SKIP_ML_SETUP")) == "1" {
		return mlPythonSetupConfig{}
	}
	if strings.TrimSpace(os.Getenv("AI_SERVER_PYTHON")) != "" {
		return mlPythonSetupConfig{}
	}
	if strings.TrimSpace(os.Getenv("NSFWAPP_ML_SETUP")) != "1" {
		return mlPythonSetupConfig{}
	}

	appDir := strings.TrimSpace(os.Getenv("NSFWAPP_HOME"))
	if appDir == "" {
		if wd, err := os.Getwd(); err == nil {
			appDir = wd
		}
	}
	if appDir == "" {
		appDir = "."
	}

	requirementsPath := strings.TrimSpace(os.Getenv("NSFWAPP_ML_REQUIREMENTS"))
	if requirementsPath == "" {
		requirementsPath = filepath.Join(appDir, "requirements-ml.txt")
	}

	venvDir := strings.TrimSpace(os.Getenv("NSFWAPP_ML_VENV"))
	if venvDir == "" {
		venvDir = filepath.Join(appDir, "ml-venv")
	}

	logPath := strings.TrimSpace(os.Getenv("NSFWAPP_ML_SETUP_LOG"))
	if logPath == "" {
		logPath = filepath.Join(appDir, "ml-setup.log")
	}

	return mlPythonSetupConfig{
		enabled:          true,
		requirementsPath: filepath.Clean(requirementsPath),
		venvDir:          filepath.Clean(venvDir),
		venvPython:       filepath.Join(filepath.Clean(venvDir), "bin", "python"),
		requirementsMark: filepath.Join(filepath.Clean(venvDir), ".requirements-ml.txt"),
		logPath:          filepath.Clean(logPath),
	}
}

func mlPythonSetupStatusSnapshot() mlPythonSetupStatus {
	mlSetupMu.RLock()
	defer mlSetupMu.RUnlock()
	return mlSetupStatus
}

func updateMLPythonSetupStatus(update func(*mlPythonSetupStatus)) {
	mlSetupMu.Lock()
	defer mlSetupMu.Unlock()

	update(&mlSetupStatus)
	mlSetupStatus.UpdatedAt = time.Now().UTC().Format(time.RFC3339Nano)
}

func mlPythonSetupStatusHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "Nur GET erlaubt", http.StatusMethodNotAllowed)
		return
	}
	writeJSON(w, http.StatusOK, mlPythonSetupStatusSnapshot())
}

func ensureMLPythonSetup(ctx context.Context) error {
	cfg := mlPythonSetupConfigFromEnv()
	if !cfg.enabled {
		initMLPythonSetupStatus()
		return nil
	}

	mlSetupRunMu.Lock()
	defer mlSetupRunMu.Unlock()

	if strings.TrimSpace(os.Getenv("AI_SERVER_PYTHON")) != "" {
		initMLPythonSetupStatus()
		return nil
	}

	startedAt := time.Now().UTC().Format(time.RFC3339Nano)
	updateMLPythonSetupStatus(func(s *mlPythonSetupStatus) {
		s.Required = true
		s.Running = true
		s.Done = false
		s.OK = false
		s.Progress = 1
		s.Phase = "prepare"
		s.Message = "ML-Python-Umgebung wird vorbereitet."
		s.Error = ""
		s.LogFile = filepath.Base(cfg.logPath)
		s.StartedAt = startedAt
		s.FinishedAt = ""
	})

	reqBytes, err := os.ReadFile(cfg.requirementsPath)
	if err != nil {
		return finishMLPythonSetupError("Requirements konnten nicht gelesen werden.", err)
	}

	if err := os.MkdirAll(cfg.venvDir, 0755); err != nil {
		return finishMLPythonSetupError("ML-Venv-Ordner konnte nicht erstellt werden.", err)
	}
	if err := os.MkdirAll(filepath.Dir(cfg.logPath), 0755); err != nil {
		return finishMLPythonSetupError("ML-Setup-Logordner konnte nicht erstellt werden.", err)
	}

	if current, err := os.ReadFile(cfg.requirementsMark); err == nil && string(current) == string(reqBytes) {
		if fi, statErr := os.Stat(cfg.venvPython); statErr == nil && fi != nil && !fi.IsDir() {
			_ = os.Setenv("AI_SERVER_PYTHON", cfg.venvPython)
			finishMLPythonSetupSuccess("ML-Python-Pakete sind bereits installiert.")
			return nil
		}
	}

	if fi, err := os.Stat(cfg.venvPython); err != nil || fi == nil || fi.IsDir() {
		python3, lookErr := exec.LookPath("python3")
		if lookErr != nil {
			return finishMLPythonSetupError("python3 wurde nicht gefunden.", lookErr)
		}
		if err := runMLPythonSetupCommand(ctx, cfg, "venv", "Python-Umgebung wird erstellt.", 5, 20, python3, "-m", "venv", cfg.venvDir); err != nil {
			return finishMLPythonSetupError("Python-Umgebung konnte nicht erstellt werden.", err)
		}
	}

	if err := runMLPythonSetupCommand(ctx, cfg, "pip", "pip wird aktualisiert.", 22, 38, cfg.venvPython, "-m", "pip", "install", "--upgrade", "pip"); err != nil {
		return finishMLPythonSetupError("pip konnte nicht aktualisiert werden.", err)
	}

	if err := runMLPythonSetupCommand(ctx, cfg, "requirements", "ML-Python-Pakete werden installiert.", 40, 94, cfg.venvPython, "-m", "pip", "install", "-r", cfg.requirementsPath); err != nil {
		return finishMLPythonSetupError("ML-Python-Pakete konnten nicht installiert werden.", err)
	}

	if err := atomicWriteFile(cfg.requirementsMark, reqBytes); err != nil {
		return finishMLPythonSetupError("ML-Requirements-Marker konnte nicht geschrieben werden.", err)
	}

	_ = os.Setenv("AI_SERVER_PYTHON", cfg.venvPython)
	finishMLPythonSetupSuccess("ML-Python-Pakete wurden installiert.")
	return nil
}

func runMLPythonSetupCommand(ctx context.Context, cfg mlPythonSetupConfig, phase, message string, startProgress, endProgress int, name string, args ...string) error {
	updateMLPythonSetupStatus(func(s *mlPythonSetupStatus) {
		s.Running = true
		s.Progress = startProgress
		s.Phase = phase
		s.Message = message
	})

	logFile, err := os.OpenFile(cfg.logPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
	if err != nil {
		return err
	}
	defer logFile.Close()

	_, _ = logFile.WriteString("\n[" + time.Now().Format(time.RFC3339) + "] " + name + " " + strings.Join(args, " ") + "\n")

	cmd := exec.CommandContext(ctx, name, args...)
	cmd.Dir = filepath.Dir(cfg.logPath)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return err
	}

	if err := cmd.Start(); err != nil {
		return err
	}

	lines := make(chan string, 64)
	var scanWG sync.WaitGroup
	scan := func(r io.Reader) {
		defer scanWG.Done()
		scanner := bufio.NewScanner(r)
		scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
		for scanner.Scan() {
			line := scanner.Text()
			_, _ = logFile.WriteString(line + "\n")
			select {
			case lines <- line:
			case <-ctx.Done():
				return
			}
		}
	}

	scanWG.Add(2)
	go scan(stdout)
	go scan(stderr)

	done := make(chan error, 1)
	go func() {
		scanWG.Wait()
		done <- cmd.Wait()
	}()

	lineCount := 0
	progressRange := endProgress - startProgress
	if progressRange < 1 {
		progressRange = 1
	}

	for {
		select {
		case line := <-lines:
			line = strings.TrimSpace(line)
			if line == "" {
				continue
			}
			lineCount++
			progress := startProgress + minInt(progressRange-1, lineCount)
			updateMLPythonSetupStatus(func(s *mlPythonSetupStatus) {
				s.Progress = progress
				s.Phase = phase
				s.Message = line
			})

		case err := <-done:
			if err != nil {
				return err
			}
			updateMLPythonSetupStatus(func(s *mlPythonSetupStatus) {
				s.Progress = endProgress
				s.Phase = phase
				s.Message = message
			})
			return nil

		case <-ctx.Done():
			_ = cmd.Process.Kill()
			return ctx.Err()
		}
	}
}

func finishMLPythonSetupSuccess(message string) {
	updateMLPythonSetupStatus(func(s *mlPythonSetupStatus) {
		s.Required = true
		s.Running = false
		s.Done = true
		s.OK = true
		s.Progress = 100
		s.Phase = "done"
		s.Message = message
		s.Error = ""
		s.FinishedAt = time.Now().UTC().Format(time.RFC3339Nano)
	})
}

func finishMLPythonSetupError(message string, err error) error {
	if err == nil {
		err = errors.New(message)
	}
	updateMLPythonSetupStatus(func(s *mlPythonSetupStatus) {
		s.Required = true
		s.Running = false
		s.Done = true
		s.OK = false
		s.Phase = "error"
		s.Message = message
		s.Error = err.Error()
		s.FinishedAt = time.Now().UTC().Format(time.RFC3339Nano)
	})
	return appErrorf("%s: %w", message, err)
}

func minInt(a, b int) int {
	if a < b {
		return a
	}
	return b
}