StreamRecorder/pkg/processing/service.go

package processing

import (
	"bufio"
	"context"
	"fmt"
	"log"
	"m3u8-downloader/pkg/constants"
	"m3u8-downloader/pkg/nas"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"sync"
)

type ProcessingService struct {
	config *ProcessConfig
	nas    *nas.NASService
}

func NewProcessingService(eventName string) (*ProcessingService, error) {
	config := &ProcessConfig{
		WorkerCount:     constants.ProcessWorkerCount,
		SourcePath:      constants.NASOutputPath,
		DestinationPath: constants.ProcessOutputPath,
		Enabled:         constants.ProcessingEnabled,
		EventName:       eventName,
	}

	nasConfig := &nas.NASConfig{
		Path:       constants.NASOutputPath,
		Username:   constants.NASUsername,
		Password:   constants.NASPassword,
		Timeout:    constants.TransferTimeout,
		RetryLimit: constants.TransferRetryLimit,
		VerifySize: true,
	}

	nasService := nas.NewNASService(*nasConfig)

	if err := nasService.TestConnection(); err != nil {
		return nil, fmt.Errorf("Failed to connect to NAS: %w", err)
	}

	return &ProcessingService{
		config: config,
		nas:    nasService,
	}, nil
}

func (ps *ProcessingService) GetSegmentInfo() (map[int]string, error) {
	return nil, nil
}

func (ps *ProcessingService) GetEventDirs() ([]string, error) {
	if ps.config.EventName == "" {
		dirs, err := os.ReadDir(ps.config.SourcePath)
		if err != nil {
			return nil, fmt.Errorf("Failed to read directory: %w", err)
		}
		var eventDirs []string
		for _, dir := range dirs {
			if dir.IsDir() {
				eventDirs = append(eventDirs, dir.Name())
			}
		}
		return eventDirs, nil
	} else {
		return []string{ps.config.EventName}, nil
	}
}

func (ps *ProcessingService) Start(ctx context.Context) error {
	if !ps.config.Enabled {
		log.Println("Processing service disabled")
		return nil
	}

	if ps.config.EventName == "" {
		events, err := ps.GetEventDirs()
		if err != nil {
			return fmt.Errorf("Failed to get event directories: %w", err)
		}
		if len(events) == 0 {
			return fmt.Errorf("No events found")
		}
		if len(events) > 1 {
			fmt.Println("Multiple events found, please select one:")
			for i, event := range events {
				fmt.Printf("%d. %s\n", i+1, event)
			}
			reader := bufio.NewReader(os.Stdin)
			input, _ := reader.ReadString('\n')
			input = strings.TrimSpace(input)
			index, err := strconv.Atoi(input)
			if err != nil {
				return fmt.Errorf("Failed to parse input: %w", err)
			}
			if index < 1 || index > len(events) {
				return fmt.Errorf("Invalid input")
			}
			ps.config.EventName = events[index-1]
		} else {
			ps.config.EventName = events[0]
		}
	}

	//Get all present resolutions
	dirs, err := ps.GetResolutions()
	if err != nil {
		return fmt.Errorf("Failed to get resolutions: %w", err)
	}

	//Spawn a worker per resolution
	ch := make(chan SegmentInfo, 100)
	var wg sync.WaitGroup

	for _, resolution := range dirs {
		wg.Add(1)
		go ps.ParseResolutionDirectory(resolution, ch, &wg)
	}
	go func() {
		wg.Wait()
		close(ch)
	}()

	segments, err := ps.AggregateSegmentInfo(ch)
	if err != nil {
		return fmt.Errorf("Failed to aggregate segment info: %w", err)
	}

	aggFile, err := ps.WriteConcatFile(segments)
	if err != nil {
		return fmt.Errorf("Failed to write concat file: %w", err)
	}

	//Feed info to ffmpeg to stitch files together
	outPath := filepath.Join(constants.NASOutputPath, ps.config.DestinationPath, ps.config.EventName)
	concatErr := ps.RunFFmpeg(aggFile, outPath)
	if concatErr != nil {
		return concatErr
	}

	return nil
}

func (ps *ProcessingService) GetResolutions() ([]string, error) {
	dirs, err := os.ReadDir(ps.config.SourcePath + "/" + ps.config.EventName)
	if err != nil {
		return nil, fmt.Errorf("Failed to read source directory: %w", err)
	}

	re := regexp.MustCompile(`^\d+p$`)

	var resolutions []string
	for _, dir := range dirs {
		if dir.IsDir() && re.MatchString(dir.Name()) {
			resolutions = append(resolutions, dir.Name())
		}
	}

	return resolutions, nil
}

func (ps *ProcessingService) ParseResolutionDirectory(resolution string, ch chan<- SegmentInfo, wg *sync.WaitGroup) {
	defer wg.Done()

	files, err := os.ReadDir(ps.config.SourcePath + "/" + ps.config.EventName + "/" + resolution)
	if err != nil {
		log.Printf("Failed to read resolution directory: %v", err)
		return
	}

	for _, file := range files {
		if !file.IsDir() {
			if !strings.HasSuffix(strings.ToLower(file.Name()), ".ts") {
				continue
			}
			no, err := strconv.Atoi(file.Name()[6:10])
			if err != nil {
				log.Printf("Failed to parse segment number: %v", err)
				continue
			}
			ch <- SegmentInfo{
				Name:       file.Name(),
				SeqNo:      no,
				Resolution: resolution,
			}
		}
	}
}

func (ps *ProcessingService) AggregateSegmentInfo(ch <-chan SegmentInfo) (map[int]SegmentInfo, error) {
	segmentMap := make(map[int]SegmentInfo)

	rank := map[string]int{
		"1080p": 1,
		"720p":  2,
		"540p":  3,
		"480p":  4,
		"450p":  5,
		"360p":  6,
		"270p":  7,
		"240p":  8,
	}

	for segment := range ch {
		fmt.Printf("Received segment %s in resolution %s \n", segment.Name, segment.Resolution)
		current, exists := segmentMap[segment.SeqNo]
		if !exists || rank[segment.Resolution] > rank[current.Resolution] {
			segmentMap[segment.SeqNo] = segment
		}
	}

	return segmentMap, nil
}

func (ps *ProcessingService) WriteConcatFile(segmentMap map[int]SegmentInfo) (string, error) {
	concatPath := filepath.Join(constants.NASOutputPath, constants.ProcessOutputPath, ps.config.EventName)
	// Ensure the directory exists
	if err := os.MkdirAll(concatPath, 0755); err != nil {
		return "", fmt.Errorf("failed to create directories for concat path: %w", err)
	}

	concatFilePath := filepath.Join(concatPath, ps.config.EventName+".txt")
	f, err := os.Create(concatFilePath)
	if err != nil {
		return "", fmt.Errorf("failed to create concat file: %w", err)
	}
	defer f.Close()

	// Sort keys to preserve order
	keys := make([]int, 0, len(segmentMap))
	for k := range segmentMap {
		keys = append(keys, k)
	}
	sort.Ints(keys)

	for _, seq := range keys {
		segment := segmentMap[seq]
		filePath := filepath.Join(ps.config.SourcePath, ps.config.EventName, segment.Resolution, segment.Name)
		line := fmt.Sprintf("file '%s'\n", filePath)
		if _, err := f.WriteString(line); err != nil {
			return "", fmt.Errorf("failed to write to concat file: %w", err)
		}
	}

	return concatFilePath, nil
}

func ffmpegPath() (string, error) {
	var baseDir string

	exePath, err := os.Executable()
	if err == nil {
		baseDir = filepath.Dir(exePath)
	} else {
		// fallback to current working directory
		baseDir, err = os.Getwd()
		if err != nil {
			return "", err
		}
	}

	// When running in GoLand, exePath might point to the IDE launcher,
	// so optionally check if ffmpeg exists here, else fallback to CWD
	ffmpeg := filepath.Join(baseDir, "bin", "ffmpeg")
	if runtime.GOOS == "windows" {
		ffmpeg += ".exe"
	}

	if _, err := os.Stat(ffmpeg); os.IsNotExist(err) {
		// try cwd instead
		cwd, err := os.Getwd()
		if err != nil {
			return "", err
		}
		ffmpeg = filepath.Join(cwd, "bin", "ffmpeg")
		if runtime.GOOS == "windows" {
			ffmpeg += ".exe"
		}
	}

	return ffmpeg, nil
}

func (ps *ProcessingService) RunFFmpeg(inputPath, outputPath string) error {
	fmt.Println("Running ffmpeg...")

	fileOutPath := filepath.Join(outputPath, ps.config.EventName+".mp4")
	fmt.Println("Input path:", inputPath)
	fmt.Println("Attempting to write to: ", outputPath+"/"+ps.config.EventName+".mp4")

	path, err := ffmpegPath()
	if err != nil {
		return err
	}
	cmd := exec.Command(path, "-f", "concat", "-safe", "0", "-i", inputPath, "-c", "copy", filepath.Join(outputPath, ps.config.EventName+".mp4"))
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	ffmpegErr := cmd.Run()
	if ffmpegErr != nil {
		return fmt.Errorf("Failed to run ffmpeg: %w", ffmpegErr)
	}
	fmt.Println("Output path:", fileOutPath)
	return nil
}