DRMDTool/utils.go

package main

import (
	"encoding/json"
	"fmt"
	"io"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"sync"

	"github.com/beevik/etree"
)

type JobInfo struct {
	AbortChan  chan struct{}
	ResumeChan chan struct{}
	Cmd        *exec.Cmd
	Paused     bool
	TempDir    string
}

var (
	jobsMutex sync.Mutex
	jobs      = make(map[string]*JobInfo)
)

func sanitizeFilename(filename string) string {
	filename = regexp.MustCompile(`[<>:"/\\|?*]`).ReplaceAllString(filename, "_")

	filename = strings.Trim(filename, ".")

	return filename
}

func isValidURL(toTest string) bool {
	_, err := url.ParseRequestURI(toTest)
	return err == nil
}

func fixGoPlay(mpdContent string) (string, error) {
	doc := etree.NewDocument()
	if err := doc.ReadFromString(mpdContent); err != nil {
		return "", fmt.Errorf("error parsing MPD content: %v", err)
	}

	root := doc.Root()

	// Remove ad periods
	for _, period := range root.SelectElements("Period") {
		if strings.Contains(period.SelectAttrValue("id", ""), "-ad-") {
			root.RemoveChild(period)
		}
	}

	// Find highest bandwidth for video
	highestBandwidth := 0
	for _, adaptationSet := range root.FindElements("//AdaptationSet") {
		if strings.Contains(adaptationSet.SelectAttrValue("mimeType", ""), "video") {
			for _, representation := range adaptationSet.SelectElements("Representation") {
				bandwidth, _ := strconv.Atoi(representation.SelectAttrValue("bandwidth", "0"))
				if bandwidth > highestBandwidth {
					highestBandwidth = bandwidth
				}
			}
		}
	}

	// Remove lower bitrate representations
	for _, adaptationSet := range root.FindElements("//AdaptationSet") {
		if strings.Contains(adaptationSet.SelectAttrValue("mimeType", ""), "video") {
			for _, representation := range adaptationSet.SelectElements("Representation") {
				bandwidth, _ := strconv.Atoi(representation.SelectAttrValue("bandwidth", "0"))
				if bandwidth != highestBandwidth {
					adaptationSet.RemoveChild(representation)
				}
			}
		}
	}

	// Combine periods
	periods := root.SelectElements("Period")
	if len(periods) > 1 {
		firstPeriod := periods[0]
		var newVideoTimeline, newAudioTimeline *etree.Element

		// Find or create SegmentTimeline elements
		for _, adaptationSet := range firstPeriod.SelectElements("AdaptationSet") {
			mimeType := adaptationSet.SelectAttrValue("mimeType", "")
			if strings.Contains(mimeType, "video") && newVideoTimeline == nil {
				newVideoTimeline = findOrCreateSegmentTimeline(adaptationSet)
			} else if strings.Contains(mimeType, "audio") && newAudioTimeline == nil {
				newAudioTimeline = findOrCreateSegmentTimeline(adaptationSet)
			}
		}

		for _, period := range periods[1:] {
			for _, adaptationSet := range period.SelectElements("AdaptationSet") {
				mimeType := adaptationSet.SelectAttrValue("mimeType", "")
				var timeline *etree.Element
				if strings.Contains(mimeType, "video") {
					timeline = newVideoTimeline
				} else if strings.Contains(mimeType, "audio") {
					timeline = newAudioTimeline
				}

				if timeline != nil {
					segmentTimeline := findOrCreateSegmentTimeline(adaptationSet)
					for _, s := range segmentTimeline.SelectElements("S") {
						timeline.AddChild(s.Copy())
					}
				}
			}
			root.RemoveChild(period)
		}
	}

	return doc.WriteToString()
}

func findOrCreateSegmentTimeline(adaptationSet *etree.Element) *etree.Element {
	for _, representation := range adaptationSet.SelectElements("Representation") {
		for _, segmentTemplate := range representation.SelectElements("SegmentTemplate") {
			timeline := segmentTemplate.SelectElement("SegmentTimeline")
			if timeline != nil {
				return timeline
			}
		}
	}

	// If no SegmentTimeline found, create one
	representation := adaptationSet.CreateElement("Representation")
	segmentTemplate := representation.CreateElement("SegmentTemplate")
	return segmentTemplate.CreateElement("SegmentTimeline")
}

func parseInputFile(filename string) ([]Item, error) {
	fileInfo, err := os.Stat(filename)
	if err != nil {
		return nil, err
	}
	if fileInfo.IsDir() {
		return nil, fmt.Errorf("%s is a directory", filename)
	}

	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	byteValue, err := io.ReadAll(file)
	if err != nil {
		return nil, err
	}

	byteValue = removeBOM(byteValue)

	var items Items
	err = json.Unmarshal(byteValue, &items)
	if err != nil {
		return nil, err
	}

	return items.Items, nil
}

func groupItemsBySeason(items []Item) map[string][]Item {
	grouped := make(map[string][]Item)
	for _, item := range items {
		metadata := parseMetadata(item.Metadata)
		if metadata.Type == "serie" {
			key := fmt.Sprintf("%s - %s", metadata.Title, metadata.Season)
			grouped[key] = append(grouped[key], item)
		} else {
			grouped["Movies"] = append(grouped["Movies"], item)
		}
	}
	return grouped
}

func filterSelectedItems(items []Item, selectedItems []string) []Item {
	var filtered []Item
	for _, item := range items {
		for _, selected := range selectedItems {
			if item.Filename == selected {
				filtered = append(filtered, item)
				break
			}
		}
	}
	return filtered
}

func sortItems(items []Item) {
	sort.Slice(items, func(i, j int) bool {
		iMeta := parseMetadata(items[i].Metadata)
		jMeta := parseMetadata(items[j].Metadata)

		if iMeta.Title != jMeta.Title {
			return iMeta.Title < jMeta.Title
		}

		iSeason := extractNumber(iMeta.Season)
		jSeason := extractNumber(jMeta.Season)

		if iSeason != jSeason {
			return iSeason < jSeason
		}

		iEpisode := extractEpisodeNumber(items[i].Filename)
		jEpisode := extractEpisodeNumber(items[j].Filename)

		return iEpisode < jEpisode
	})
}

func extractNumber(s string) int {
	num, _ := strconv.Atoi(strings.TrimLeft(s, "S"))
	return num
}

func extractEpisodeNumber(filename string) int {
	parts := strings.Split(filename, "E")
	if len(parts) > 1 {
		num, _ := strconv.Atoi(parts[1])
		return num
	}
	return 0
}

func processItems(filename string, items []Item) error {
	jobsMutex.Lock()
	jobInfo := &JobInfo{
		AbortChan:  make(chan struct{}),
		ResumeChan: make(chan struct{}),
	}
	jobs[filename] = jobInfo
	jobsMutex.Unlock()

	defer func() {
		jobsMutex.Lock()
		delete(jobs, filename)
		jobsMutex.Unlock()

		if jobInfo.TempDir != "" {
			os.RemoveAll(jobInfo.TempDir)
		}
	}()

	sortItems(items)

	for i := 0; i < len(items); i++ {
		select {
		case <-jobInfo.AbortChan:
			updateProgress(filename, 100, "Aborted")
			return fmt.Errorf("download aborted")
		default:
			if jobInfo.Paused {
				select {
				case <-jobInfo.ResumeChan:
					jobInfo.Paused = false
				case <-jobInfo.AbortChan:
					updateProgress(filename, 100, "Aborted")
					return fmt.Errorf("download aborted")
				}
			}
			updateProgress(filename, float64(i)/float64(len(items))*100, items[i].Filename)
			err := downloadFile(items[i], jobInfo)
			if err != nil {
				if err.Error() == "download paused" {
					removeCompletedEpisodes(filename, items[:i])
					i--
					continue
				}
			}
		}
	}
	updateProgress(filename, 100, "")
	return nil
}

func removeCompletedEpisodes(filename string, completedItems []Item) error {
	inputFile := filepath.Join(uploadDir, filename)
	items, err := parseInputFile(inputFile)
	if err != nil {
		return fmt.Errorf("error parsing input file: %v", err)
	}

	remainingItems := make([]Item, 0)
	for _, item := range items {
		if !isItemCompleted(item, completedItems) || isLastCompletedItem(item, completedItems) {
			remainingItems = append(remainingItems, item)
		}
	}

	updatedItems := Items{Items: remainingItems}
	jsonData, err := json.MarshalIndent(updatedItems, "", "  ")
	if err != nil {
		return fmt.Errorf("error marshaling updated items: %v", err)
	}

	err = os.WriteFile(inputFile, jsonData, 0644)
	if err != nil {
		return fmt.Errorf("error writing updated DRMD file: %v", err)
	}

	return nil
}

func isItemCompleted(item Item, completedItems []Item) bool {
	for _, completedItem := range completedItems {
		if item.Filename == completedItem.Filename {
			return true
		}
	}
	return false
}

func isLastCompletedItem(item Item, completedItems []Item) bool {
	if len(completedItems) == 0 {
		return false
	}
	return item.Filename == completedItems[len(completedItems)-1].Filename
}