fetch_ml/cmd/tui/internal/services/services.go

package services

import (
	"context"
	"fmt"

	"github.com/jfraeys/fetch_ml/cmd/tui/internal/config"
	"github.com/jfraeys/fetch_ml/cmd/tui/internal/model"
	"github.com/jfraeys/fetch_ml/internal/experiment"
	"github.com/jfraeys/fetch_ml/internal/network"
	"github.com/jfraeys/fetch_ml/internal/queue"
)

// TaskQueue wraps the internal queue.TaskQueue for TUI compatibility
type TaskQueue struct {
	internal   *queue.TaskQueue
	expManager *experiment.Manager
	ctx        context.Context
}

func NewTaskQueue(cfg *config.Config) (*TaskQueue, error) {
	// Create internal queue config
	queueCfg := queue.Config{
		RedisAddr:     cfg.RedisAddr,
		RedisPassword: cfg.RedisPassword,
		RedisDB:       cfg.RedisDB,
	}

	internalQueue, err := queue.NewTaskQueue(queueCfg)
	if err != nil {
		return nil, fmt.Errorf("failed to create task queue: %w", err)
	}

	// Initialize experiment manager
	// TODO: Get base path from config
	expManager := experiment.NewManager("./experiments")

	return &TaskQueue{
		internal:   internalQueue,
		expManager: expManager,
		ctx:        context.Background(),
	}, nil
}

func (tq *TaskQueue) EnqueueTask(jobName, args string, priority int64) (*model.Task, error) {
	// Create internal task
	internalTask := &queue.Task{
		JobName:  jobName,
		Args:     args,
		Priority: priority,
	}

	// Use internal queue to enqueue
	err := tq.internal.AddTask(internalTask)
	if err != nil {
		return nil, err
	}

	// Convert to TUI model
	return &model.Task{
		ID:        internalTask.ID,
		JobName:   internalTask.JobName,
		Args:      internalTask.Args,
		Status:    "queued",
		Priority:  int64(internalTask.Priority),
		CreatedAt: internalTask.CreatedAt,
		Metadata:  internalTask.Metadata,
	}, nil
}

func (tq *TaskQueue) GetNextTask() (*model.Task, error) {
	internalTask, err := tq.internal.GetNextTask()
	if err != nil {
		return nil, err
	}
	if internalTask == nil {
		return nil, nil
	}

	// Convert to TUI model
	return &model.Task{
		ID:        internalTask.ID,
		JobName:   internalTask.JobName,
		Args:      internalTask.Args,
		Status:    internalTask.Status,
		Priority:  internalTask.Priority,
		CreatedAt: internalTask.CreatedAt,
		Metadata:  internalTask.Metadata,
	}, nil
}

func (tq *TaskQueue) GetTask(taskID string) (*model.Task, error) {
	internalTask, err := tq.internal.GetTask(taskID)
	if err != nil {
		return nil, err
	}

	// Convert to TUI model
	return &model.Task{
		ID:        internalTask.ID,
		JobName:   internalTask.JobName,
		Args:      internalTask.Args,
		Status:    internalTask.Status,
		Priority:  internalTask.Priority,
		CreatedAt: internalTask.CreatedAt,
		Metadata:  internalTask.Metadata,
	}, nil
}

func (tq *TaskQueue) UpdateTask(task *model.Task) error {
	// Convert to internal task
	internalTask := &queue.Task{
		ID:        task.ID,
		JobName:   task.JobName,
		Args:      task.Args,
		Status:    task.Status,
		Priority:  task.Priority,
		CreatedAt: task.CreatedAt,
		Metadata:  task.Metadata,
	}

	return tq.internal.UpdateTask(internalTask)
}

func (tq *TaskQueue) GetQueuedTasks() ([]*model.Task, error) {
	internalTasks, err := tq.internal.GetAllTasks()
	if err != nil {
		return nil, err
	}

	// Convert to TUI models
	tasks := make([]*model.Task, len(internalTasks))
	for i, task := range internalTasks {
		tasks[i] = &model.Task{
			ID:        task.ID,
			JobName:   task.JobName,
			Args:      task.Args,
			Status:    task.Status,
			Priority:  task.Priority,
			CreatedAt: task.CreatedAt,
			Metadata:  task.Metadata,
		}
	}

	return tasks, nil
}

func (tq *TaskQueue) GetJobStatus(jobName string) (map[string]string, error) {
	// This method doesn't exist in internal queue, implement basic version
	task, err := tq.internal.GetTaskByName(jobName)
	if err != nil {
		return nil, err
	}
	if task == nil {
		return map[string]string{"status": "not_found"}, nil
	}

	return map[string]string{
		"status":  task.Status,
		"task_id": task.ID,
	}, nil
}

func (tq *TaskQueue) RecordMetric(jobName, metric string, value float64) error {
	return tq.internal.RecordMetric(jobName, metric, value)
}

func (tq *TaskQueue) GetMetrics(jobName string) (map[string]string, error) {
	// This method doesn't exist in internal queue, return empty for now
	return map[string]string{}, nil
}

func (tq *TaskQueue) ListDatasets() ([]model.DatasetInfo, error) {
	// This method doesn't exist in internal queue, return empty for now
	return []model.DatasetInfo{}, nil
}

func (tq *TaskQueue) CancelTask(taskID string) error {
	return tq.internal.CancelTask(taskID)
}

func (tq *TaskQueue) ListExperiments() ([]string, error) {
	return tq.expManager.ListExperiments()
}

func (tq *TaskQueue) GetExperimentDetails(commitID string) (string, error) {
	meta, err := tq.expManager.ReadMetadata(commitID)
	if err != nil {
		return "", err
	}

	metrics, err := tq.expManager.GetMetrics(commitID)
	if err != nil {
		return "", err
	}

	output := fmt.Sprintf("Experiment: %s\n", meta.JobName)
	output += fmt.Sprintf("Commit ID: %s\n", meta.CommitID)
	output += fmt.Sprintf("User: %s\n", meta.User)
	output += fmt.Sprintf("Timestamp: %d\n\n", meta.Timestamp)
	output += "Metrics:\n"

	if len(metrics) == 0 {
		output += "  No metrics logged.\n"
	} else {
		for _, m := range metrics {
			output += fmt.Sprintf("  %s: %.4f (Step: %d)\n", m.Name, m.Value, m.Step)
		}
	}

	return output, nil
}

func (tq *TaskQueue) Close() error {
	return tq.internal.Close()
}

// MLServer wraps network.SSHClient for backward compatibility
type MLServer struct {
	*network.SSHClient
	addr string
}

func NewMLServer(cfg *config.Config) (*MLServer, error) {
	// Local mode: skip SSH entirely
	if cfg.Host == "" {
		client, _ := network.NewSSHClient("", "", "", 0, "")
		return &MLServer{SSHClient: client, addr: "localhost"}, nil
	}

	client, err := network.NewSSHClient(cfg.Host, cfg.User, cfg.SSHKey, cfg.Port, cfg.KnownHosts)
	if err != nil {
		return nil, err
	}
	addr := fmt.Sprintf("%s:%d", cfg.Host, cfg.Port)
	return &MLServer{SSHClient: client, addr: addr}, nil
}