fetch_ml/cmd/tui/main.go

// Package main implements the ML TUI
package main

import (
	"log"
	"os"
	"os/signal"
	"syscall"

	tea "github.com/charmbracelet/bubbletea"
	"github.com/jfraeys/fetch_ml/cmd/tui/internal/config"
	"github.com/jfraeys/fetch_ml/cmd/tui/internal/controller"
	"github.com/jfraeys/fetch_ml/cmd/tui/internal/model"
	"github.com/jfraeys/fetch_ml/cmd/tui/internal/services"
	"github.com/jfraeys/fetch_ml/cmd/tui/internal/view"
	"github.com/jfraeys/fetch_ml/internal/auth"
	"github.com/jfraeys/fetch_ml/internal/logging"
)

type AppModel struct {
	state      model.State
	controller *controller.Controller
}

func (m AppModel) Init() tea.Cmd {
	return m.controller.Init()
}

func (m AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	newState, cmd := m.controller.Update(msg, m.state)
	m.state = newState
	return m, cmd
}

func (m AppModel) View() string {
	return view.Render(m.state)
}

func main() {
	// Parse authentication flags
	authFlags := auth.ParseAuthFlags()
	if err := auth.ValidateAuthFlags(authFlags); err != nil {
		log.Fatalf("Authentication flag error: %v", err)
	}

	// Get API key from various sources
	apiKey := auth.GetAPIKeyFromSources(authFlags)

	var (
		cfg         *config.Config
		cliConfig   *config.CLIConfig
		cliConfPath string
	)

	configFlag := authFlags.ConfigFile

	// Only support TOML configuration
	var err error
	cliConfig, cliConfPath, err = config.LoadCLIConfig(configFlag)
	if err != nil {
		if configFlag != "" {
			log.Fatalf("Failed to load TOML config %s: %v", configFlag, err)
		} else {
			// Provide helpful error message for data scientists
			log.Printf("=== Fetch ML TUI - Configuration Required ===")
			log.Printf("")
			log.Printf("Error: %v", err)
			log.Printf("")
			log.Printf("To get started with the TUI, you need to initialize your configuration:")
			log.Printf("")
			log.Printf("Option 1: Using the Zig CLI (Recommended)")
			log.Printf("  1. Build the CLI: cd cli && make build")
			log.Printf("  2. Initialize config: ./cli/zig-out/bin/ml init")
			log.Printf("  3. Edit ~/.ml/config.toml with your settings")
			log.Printf("  4. Run TUI: ./bin/tui")
			log.Printf("")
			log.Printf("Option 2: Manual Configuration")
			log.Printf("  1. Create directory: mkdir -p ~/.ml")
			log.Printf("  2. Create config: touch ~/.ml/config.toml")
			log.Printf("  3. Add your settings to the file")
			log.Printf("  4. Run TUI: ./bin/tui")
			log.Printf("")
			log.Printf("Example ~/.ml/config.toml:")
			log.Printf("  worker_host = \"localhost\"")
			log.Printf("  worker_user = \"your_username\"")
			log.Printf("  worker_base = \"~/ml_jobs\"")
			log.Printf("  worker_port = 22")
			log.Printf("  api_key = \"your_api_key_here\"")
			log.Printf("")
			log.Printf("For more help, see: https://github.com/jfraeys/fetch_ml/docs")
			os.Exit(1)
		}
	}

	cfg = cliConfig.ToTUIConfig()
	log.Printf("Loaded TOML configuration from %s", cliConfPath)

	// Validate authentication configuration
	if err := cfg.Auth.ValidateAuthConfig(); err != nil {
		log.Fatalf("Invalid authentication configuration: %v", err)
	}

	if err := cfg.Validate(); err != nil {
		log.Fatalf("Invalid configuration: %v", err)
	}

	// Test authentication if enabled
	if cfg.Auth.Enabled {
		// Use API key from CLI config if available, otherwise use from flags
		var effectiveAPIKey string
		if cliConfig != nil && cliConfig.APIKey != "" {
			effectiveAPIKey = cliConfig.APIKey
		} else if apiKey != "" {
			effectiveAPIKey = apiKey
		} else {
			log.Fatal("Authentication required but no API key provided")
		}

		if _, err := cfg.Auth.ValidateAPIKey(effectiveAPIKey); err != nil {
			log.Fatalf("Authentication failed: %v", err)
		}
	}

	srv, err := services.NewMLServer(cfg)
	if err != nil {
		log.Fatalf("Failed to connect to server: %v", err)
	}
	defer func() {
		if err := srv.Close(); err != nil {
			log.Printf("server close error: %v", err)
		}
	}()

	tq, err := services.NewTaskQueue(cfg)
	if err != nil {
		log.Fatalf("Failed to connect to Redis: %v", err)
	}
	defer func() {
		if err := tq.Close(); err != nil {
			log.Printf("task queue close error: %v", err)
		}
	}()

	// Initialize logger
	// Note: In original code, logger was created inside initialModel.
	// Here we create it and pass it to controller.
	// We use slog.LevelError as default from original code.
	// But original code imported "log/slog".
	// We use internal/logging package.
	// Check logging package signature.
	// Original: logger := logging.NewLogger(slog.LevelError, false)
	// We need to import "log/slog" in main if we use slog constants.
	// Or use logging package constants if available.
	// Let's check logging package.
	// Assuming logging.NewLogger takes (slog.Level, bool).
	// I'll import "log/slog".

	// Wait, I need to import "log/slog"
	logger := logging.NewLogger(-4, false) // -4 is slog.LevelError value. Or I can import log/slog.

	// Initialize State and Controller
	var effectiveAPIKey string
	if cliConfig != nil && cliConfig.APIKey != "" {
		effectiveAPIKey = cliConfig.APIKey
	} else {
		effectiveAPIKey = apiKey
	}

	initialState := model.InitialState(effectiveAPIKey)
	ctrl := controller.New(cfg, srv, tq, logger)

	appModel := AppModel{
		state:      initialState,
		controller: ctrl,
	}

	// Run TUI app
	p := tea.NewProgram(appModel, tea.WithAltScreen(), tea.WithMouseAllMotion())

	// Ensure we restore the terminal even if panic or error occurs
	// Note: p.Run() usually handles this, but explicit cleanup is safer
	// if we want to ensure the alt screen is exited.
	// We can't defer p.ReleaseTerminal() here because p is created here.
	// But we can defer a function that calls it.

	// Set up signal handling for graceful shutdown
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	// Run program and handle signals
	go func() {
		<-sigChan
		p.Quit()
	}()

	if _, err := p.Run(); err != nil {
		// Attempt to restore terminal before logging fatal error
		p.ReleaseTerminal()
		log.Fatalf("Error running TUI: %v", err)
	}

	// Explicitly restore terminal after program exits
	p.ReleaseTerminal()
}