stream-poc2/import.go

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path"
	"strconv"
	"stream-poc2/ffmpeg"
	"stream-poc2/model"
	"strings"
)

func isFile(path string) bool {
	stat, err := os.Stat(path)
	if err != nil {
		return false
	}
	return !stat.IsDir()
}

func isDir(path string) bool {
	stat, err := os.Stat(path)
	if err != nil {
		return false
	}
	return stat.IsDir()
}

func createRawStructure(root string, audioStreams []int) error {
	dirs := []string{
		"",
		"video",
		"audio",
		"cache",
		"cache/video",
	}

	for _, dir := range dirs {
		log.Println("mkdir:", dir)
		err := os.MkdirAll(path.Join(root, dir), 0o755)
		if err != nil {
			return err
		}
	}

	for _, streamIdx := range audioStreams {
		dir := path.Join(root, "audio", strconv.Itoa(streamIdx))
		log.Println("mkdir:", dir)
		err := os.MkdirAll(dir, 0o755)
		if err != nil {
			return err
		}
	}

	return nil
}

func writeJSON(fpath string, data any) error {
	log.Println("writing", fpath)
	f, err := os.OpenFile(fpath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o644)
	if err != nil {
		return err
	}
	defer func(f *os.File) {
		err := f.Close()
		if err != nil {
			log.Println("Failed to close", fpath, err)
		}
	}(f)

	err = json.NewEncoder(f).Encode(data)
	if err != nil {
		return err
	}
	return nil
}

func importMain(filepath, target string) {
	targetFullPath := path.Join("db", target)

	if isDir(targetFullPath) {
		fmt.Println("target already exists")
		return
	}

	if !isFile(filepath) {
		fmt.Println("Source does not exists")
		return
	}

	log.Println("Running FFProbe...")

	ffProbed, err := ffmpeg.FFProbe(filepath)
	if err != nil {
		panic(err)
	}

	var numVideoStreams int

	// we will fill those out from ffprobe's output
	nativeVideoStream := model.VideoStream{}
	var audioStreams []model.AudioStream

	for _, stream := range ffProbed.Streams {
		log.Println("Discovered stream:", stream.Index, stream.CodecType)
		if stream.CodecType == "audio" {
			lang, ok := stream.Tags["language"]
			if !ok {
				lang = "eng" // use as default
			}
			title, _ := stream.Tags["title"]
			if !ok {
				title = fmt.Sprintf("Unknown (%d)", stream.Index) // Title is used as name which must be different
			}

			sampleRate, err := strconv.Atoi(stream.SampleRate)
			if err != nil {
				log.Println("!!! Couldn't get audio sample rate !!!")
			}

			bitRate, err := strconv.Atoi(stream.BitRate)
			if err != nil {
				log.Println("!!! Couldn't get audio bit rate !!!")
			}

			audioStreams = append(audioStreams, model.AudioStream{
				Index:       stream.Index,
				Lang:        lang,
				Title:       title,
				Default:     stream.Disposition.Default == 1,
				Codec:       stream.CodecName,
				SampleRate:  sampleRate,
				Bitrate:     bitRate,
				Channels:    stream.Channels,
				AudioChunks: nil,
			})
		}
		if stream.CodecType == "video" {
			// funnily mkv can encode any sort of files, including "cover images", those represented as video frames.
			frameRateParts := strings.Split(stream.AvgFrameRate, "/")
			a, err := strconv.Atoi(frameRateParts[0])
			if err != nil {
				panic(err)
			}
			b, err := strconv.Atoi(frameRateParts[1])
			if err != nil {
				panic(err)
			}

			if a > 0 && b > 0 {
				nativeVideoStream.FrameRate = float64(a) / float64(b)
			}

			if nativeVideoStream.FrameRate == 0 || strings.HasPrefix(stream.Tags["mimetype"], "image") {
				log.Println("This is a still image, ignoring...")
				continue
			}

			numVideoStreams++
			nativeVideoStream.Resolution = model.Resolution{
				Width:  stream.Width,
				Height: stream.Height,
			}
			if stream.BitRate != "" {
				nativeVideoStream.PeakBitrate, err = strconv.Atoi(stream.BitRate)
				if err != nil {
					panic(err)
				}

				nativeVideoStream.AvgBitrate, err = strconv.Atoi(stream.BitRate)
				if err != nil {
					panic(err)
				}
			} else {
				log.Println("Warning: couldn't extract video bitrate")
			}

			nativeVideoStream.HLSCodecsString, _ = ffmpeg.FigureOutHLSCodecsParam(stream) // ignore if it's unrecognizable..
		}

	}

	if numVideoStreams != 1 {
		log.Printf("!!! The input media should have exactly one video stream, this has %d !!!\n", numVideoStreams)
	}

	metadata := model.MediaMetadata{
		FriendlyName:   path.Base(filepath),
		TargetDuration: ffmpeg.SlicerTargetDurationSec,
		VideoStream:    nativeVideoStream,
		AudioStreams:   audioStreams,
	}

	audioStreamIndexes := make([]int, len(audioStreams))
	for i, audioStream := range audioStreams {
		audioStreamIndexes[i] = audioStream.Index
	}

	// create structure
	err = createRawStructure(targetFullPath, audioStreamIndexes)
	if err != nil {
		panic(err)
	}

	// write ffprobeOutput
	err = writeJSON(path.Join(targetFullPath, "ffprobe.json"), ffProbed)
	if err != nil {
		panic(err)
	}

	log.Println("Now segmenting...")
	segmentsVideoCSV, audioSegmentsCSVs, err := ffmpeg.MakeSegments(filepath, targetFullPath, audioStreams)
	if err != nil {
		panic(err)
	}
	log.Println("Segment data: ", segmentsVideoCSV, audioSegmentsCSVs)

	log.Println("Now processing segment data...")

	// read slice stuff
	segments, err := ffmpeg.ReadSegmentsCSV(segmentsVideoCSV)
	if err != nil {
		panic(err)
	}

	// load videoStream chunks
	videoChunks := make([]model.Chunk, len(segments))
	for i, segment := range segments {
		videoChunks[i] = model.Chunk{
			FName:  segment.Name,
			Length: segment.Len(),
		}
	}

	metadata.VideoStream.VideoChunks = videoChunks

	// load audioStream chunks
	for i, slicesFile := range audioSegmentsCSVs {
		segments, err = ffmpeg.ReadSegmentsCSV(slicesFile)
		if err != nil {
			panic(err)
		}
		audioChunks := make([]model.Chunk, len(segments))
		for j, segment := range segments {
			audioChunks[j] = model.Chunk{
				FName:  segment.Name,
				Length: segment.Len(),
			}
		}
		metadata.AudioStreams[i].AudioChunks = audioChunks
	}

	log.Println("Writing metadata...")

	// write metadata
	err = writeJSON(path.Join(targetFullPath, "meta.json"), metadata)
	if err != nil {
		panic(err)
	}

	// pre-encode
	log.Println("Pre-encoding some slices...")

	for i, chunk := range videoChunks {
		// encode 7 slices to begin with...
		if i > 6 {
			break
		}

		origSegmentPath := path.Join(targetFullPath, "video", chunk.FName)
		for profileName, profileSettings := range model.VideoProfiles {
			preEncodedSegmentDir := path.Join(targetFullPath, "cache", "video", profileName)
			err = os.MkdirAll(preEncodedSegmentDir, 0o755)
			if err != nil {
				panic(err)
			}

			preEncodedSegmentPath := path.Join(preEncodedSegmentDir, chunk.FName)
			log.Println(origSegmentPath, " -> ", preEncodedSegmentPath)

			var f *os.File
			f, err = os.OpenFile(preEncodedSegmentPath, os.O_WRONLY|os.O_CREATE, 0o644)
			if err != nil {
				log.Println("failed to open chunk", err)
				panic(err)
			}

			err = ffmpeg.EncodeChunk(profileSettings, nativeVideoStream.Resolution.Height, "", origSegmentPath, f)
			_ = f.Close()
			if err != nil {
				panic(err)
			}
		}
	}

}