vivaplusdl/downloader.go

package main

import (
	"database/sql"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"time"
)

func DownloadAllVideos(db *sql.DB) error {
	// Ensure no partial videos exist
	tempDir := os.Getenv("VIVAPLUS_DOWNLOAD")
	if tempDir == "" {
		tempDir = "./temp" // Replace with your desired default path
	}
	err := os.RemoveAll(tempDir)
	if err != nil {
		return fmt.Errorf("error cleaning temporary directory: %w", err)
	}
	err = os.MkdirAll(tempDir, 0750)
	if err != nil {
		return fmt.Errorf("error creating temporary directory: %w", err)
	}

	// Find the target directory
	downloadDir := os.Getenv("VIVAPLUS_DESTINATION")
	if downloadDir == "" {
		downloadDir = "./downloads"
	}
	err = os.MkdirAll(downloadDir, 0750)
	if err != nil {
		return fmt.Errorf("error creating destination directory: %w", err)
	}

	log.Printf("Starting downloads...")
	for {
		// Fetch the next record from the database
		row := db.QueryRow("select id, url, `cast`, title, description, upload_date, thumbnail, episode from videos where state = 'pending' and episode is not null order by year, episode limit 1")
		var id, episode int
		var href, cast, title, description, uploadDateStr, thumbnailUrl string
		err = row.Scan(&id, &href, &cast, &title, &description, &uploadDateStr, &thumbnailUrl, &episode)
		if errors.Is(err, sql.ErrNoRows) {
			log.Printf("No videos found for downloading")
			return nil
		}
		if err != nil {
			return fmt.Errorf("error fetching record: %w", err)
		}
		uploadDate, err := time.Parse(time.DateOnly, uploadDateStr)
		if err != nil {
			return fmt.Errorf("error parsing upload date '%s': %w", uploadDateStr, err)
		}

		// Download the actual video
		baseFileName := fmt.Sprintf("/S%dE%03d", uploadDate.Year(), episode)
		log.Printf("Downloading %s", href)
		cmd := exec.Command("yt-dlp", cast, "-o", filepath.Join(tempDir, baseFileName+".%(ext)s"))
		cmd.Stderr = os.Stderr
		cmd.Stdout = os.Stdout
		err = cmd.Run()
		if err != nil {
			return fmt.Errorf("error downloading video: %w", err)
		}

		// Move it into the destination directory
		destinationDir := filepath.Join(downloadDir, fmt.Sprintf("Season %d", uploadDate.Year()))
		err = os.MkdirAll(destinationDir, 0750)
		if err != nil {
			return fmt.Errorf("error creating target directory: %w", err)
		}
		files, err := os.ReadDir(tempDir)
		if err != nil {
			return fmt.Errorf("error retrieving downloaded files: %w", err)
		}

		for _, file := range files {
			srcPath := filepath.Join(tempDir, file.Name())
			destPath := filepath.Join(destinationDir, file.Name())
			fmt.Printf("Moving %s to %s", srcPath, destPath)
			err = os.Rename(srcPath, destPath)
			if err != nil {
				return fmt.Errorf("error moving file: %w", err)
			}
		}

		// Write XML sidecar
		nfo := NFO{}
		nfo.Plot = description
		nfo.Title = title
		nfo.Year = uploadDate.Year()
		nfo.Aired = uploadDate.Format(time.DateOnly)
		nfoData, err := xml.MarshalIndent(nfo, "", "    ")
		if err != nil {
			return fmt.Errorf("error marshalling NFO data: %w", err)
		}
		nfoString := xml.Header + string(nfoData)
		nfoFile, err := os.Create(filepath.Join(destinationDir, baseFileName+".nfo"))
		defer nfoFile.Close()
		_, err = nfoFile.WriteString(nfoString)
		if err != nil {
			return fmt.Errorf("error writing NFO file: %w", err)
		}

		// Write thumbnail
		parsedThumbnailUrl, err := url.Parse(thumbnailUrl)
		if err != nil {
			return fmt.Errorf("error parsing thumbnail url: %w", err)
		}
		thumbnailExt := filepath.Ext(parsedThumbnailUrl.Path)
		thumbnailFile, err := os.Create(filepath.Join(destinationDir, baseFileName+"-thumb"+thumbnailExt))
		if err != nil {
			return fmt.Errorf("error creating thumbnail: %w", err)
		}
		defer thumbnailFile.Close()
		thumbnailResp, err := http.Get(thumbnailUrl)
		if err != nil {
			return fmt.Errorf("error fetching thumbnail: %w", err)
		}
		defer thumbnailResp.Body.Close()
		_, err = io.Copy(thumbnailFile, thumbnailResp.Body)
		if err != nil {
			return fmt.Errorf("error writing thumbnail: %w", err)
		}

		// Set the database state to done
		tx, err := db.Begin()
		if err != nil {
			return fmt.Errorf("error starting transaction: %w", err)
		}
		defer tx.Rollback()
		result, err := tx.Exec("update videos set state = 'done' where id = ?", id)
		if err != nil {
			return fmt.Errorf("error updating database: %w", err)
		}
		count, err := result.RowsAffected()
		if err != nil {
			return fmt.Errorf("error getting number of rows affected: %d", err)
		}
		if count != 1 {
			return fmt.Errorf("unexpected number of rows changed (expected exactly 1): %d", count)
		}
		err = tx.Commit()
		if err != nil {
			return fmt.Errorf("error commiting transaction: %w", err)
		}

		fmt.Printf("Finished downloading file")
	}
}