yutu/pkg/caption/caption.go

// Copyright 2025 eat-pray-ai & OpenWaygate
// SPDX-License-Identifier: Apache-2.0

package caption

import (
	"errors"
	"fmt"
	"io"
	"os"

	"github.com/eat-pray-ai/yutu/pkg"
	"github.com/eat-pray-ai/yutu/pkg/common"
	"github.com/jedib0t/go-pretty/v6/table"
	"google.golang.org/api/youtube/v3"
)

var (
	errGetCaption      = errors.New("failed to get caption")
	errUpdateCaption   = errors.New("failed to update caption")
	errDeleteCaption   = errors.New("failed to delete caption")
	errInsertCaption   = errors.New("failed to insert caption")
	errDownloadCaption = errors.New("failed to download caption")
)

type Caption struct {
	common.Fields
	File           string `yaml:"file" json:"file,omitempty"`
	AudioTrackType string `yaml:"audio_track_type" json:"audio_track_type,omitempty"`
	IsAutoSynced   *bool  `yaml:"is_auto_synced" json:"is_auto_synced,omitzero"`
	IsCC           *bool  `yaml:"is_cc" json:"is_cc,omitzero"`
	IsDraft        *bool  `yaml:"is_draft" json:"is_draft,omitzero"`
	IsEasyReader   *bool  `yaml:"is_easy_reader" json:"is_easy_reader,omitzero"`
	IsLarge        *bool  `yaml:"is_large" json:"is_large,omitzero"`
	Language       string `yaml:"language" json:"language,omitempty"`
	Name           string `yaml:"name" json:"name,omitempty"`
	TrackKind      string `yaml:"track_kind" json:"track_kind,omitempty"`
	OnBehalfOf     string `yaml:"on_behalf_of" json:"on_behalf_of,omitempty"`
	VideoId        string `yaml:"video_id" json:"video_id,omitempty"`
	Tfmt           string `yaml:"tfmt" json:"tfmt,omitempty"`
	Tlang          string `yaml:"tlang" json:"tlang,omitempty"`
}

type ICaption[T youtube.Caption] interface {
	Get() ([]*T, error)
	List(io.Writer) error
	Insert(io.Writer) error
	Update(io.Writer) error
	Delete(io.Writer) error
	Download(io.Writer) error
}

type Option func(*Caption)

func NewCaption(opts ...Option) ICaption[youtube.Caption] {
	c := &Caption{Fields: common.Fields{}}
	for _, opt := range opts {
		opt(c)
	}
	return c
}

func (c *Caption) Get() ([]*youtube.Caption, error) {
	if err := c.EnsureService(); err != nil {
		return nil, err
	}
	call := c.Service.Captions.List(c.Parts, c.VideoId)
	if len(c.Ids) > 0 {
		call = call.Id(c.Ids...)
	}
	if c.OnBehalfOf != "" {
		call = call.OnBehalfOf(c.OnBehalfOf)
	}
	if c.OnBehalfOfContentOwner != "" {
		call = call.OnBehalfOfContentOwner(c.OnBehalfOfContentOwner)
	}

	res, err := call.Do()
	if err != nil {
		return nil, errors.Join(errGetCaption, err)
	}

	return res.Items, nil
}

func (c *Caption) List(writer io.Writer) error {
	captions, err := c.Get()
	if err != nil {
		return err
	}

	common.PrintList(
		c.Output, captions, writer, table.Row{"ID", "Video ID", "Name", "Language"},
		func(cap *youtube.Caption) table.Row {
			return table.Row{cap.Id, cap.Snippet.VideoId, cap.Snippet.Name, cap.Snippet.Language}
		},
	)
	return nil
}

func (c *Caption) Insert(writer io.Writer) error {
	if err := c.EnsureService(); err != nil {
		return err
	}
	file, err := pkg.Root.Open(c.File)
	if err != nil {
		return errors.Join(errInsertCaption, err)
	}
	defer func(file *os.File) {
		_ = file.Close()
	}(file)

	caption := &youtube.Caption{
		Snippet: &youtube.CaptionSnippet{
			AudioTrackType: c.AudioTrackType,
			IsAutoSynced:   c.IsAutoSynced != nil && *c.IsAutoSynced,
			IsCC:           c.IsCC != nil && *c.IsCC,
			IsDraft:        c.IsDraft != nil && *c.IsDraft,
			IsEasyReader:   c.IsEasyReader != nil && *c.IsEasyReader,
			IsLarge:        c.IsLarge != nil && *c.IsLarge,
			Language:       c.Language,
			Name:           c.Name,
			TrackKind:      c.TrackKind,
			VideoId:        c.VideoId,
		},
	}

	call := c.Service.Captions.Insert([]string{"snippet"}, caption).Media(file)
	if c.OnBehalfOf != "" {
		call = call.OnBehalfOf(c.OnBehalfOf)
	}
	if c.OnBehalfOfContentOwner != "" {
		call = call.OnBehalfOfContentOwner(c.OnBehalfOfContentOwner)
	}

	res, err := call.Do()
	if err != nil {
		return errors.Join(errInsertCaption, err)
	}

	common.PrintResult(c.Output, res, writer, "Caption inserted: %s\n", res.Id)
	return nil
}

func (c *Caption) Update(writer io.Writer) error {
	c.Parts = []string{"snippet"}
	captions, err := c.Get()
	if err != nil {
		return errors.Join(errUpdateCaption, err)
	}
	if len(captions) == 0 {
		return errGetCaption
	}

	caption := captions[0]
	if c.AudioTrackType != "" {
		caption.Snippet.AudioTrackType = c.AudioTrackType
	}
	if c.IsAutoSynced != nil {
		caption.Snippet.IsAutoSynced = *c.IsAutoSynced
		caption.Snippet.ForceSendFields = append(
			caption.Snippet.ForceSendFields, "IsAutoSynced",
		)
	}
	if c.IsCC != nil {
		caption.Snippet.IsCC = *c.IsCC
		caption.Snippet.ForceSendFields = append(
			caption.Snippet.ForceSendFields, "IsCC",
		)
	}
	if c.IsDraft != nil {
		caption.Snippet.IsDraft = *c.IsDraft
		caption.Snippet.ForceSendFields = append(
			caption.Snippet.ForceSendFields, "IsDraft",
		)
	}
	if c.IsEasyReader != nil {
		caption.Snippet.IsEasyReader = *c.IsEasyReader
		caption.Snippet.ForceSendFields = append(
			caption.Snippet.ForceSendFields, "IsEasyReader",
		)
	}
	if c.IsLarge != nil {
		caption.Snippet.IsLarge = *c.IsLarge
		caption.Snippet.ForceSendFields = append(
			caption.Snippet.ForceSendFields, "IsLarge",
		)
	}
	if c.Language != "" {
		caption.Snippet.Language = c.Language
	}
	if c.Name != "" {
		caption.Snippet.Name = c.Name
	}
	if c.TrackKind != "" {
		caption.Snippet.TrackKind = c.TrackKind
	}
	if c.VideoId != "" {
		caption.Snippet.VideoId = c.VideoId
	}

	call := c.Service.Captions.Update([]string{"snippet"}, caption)
	if c.File != "" {
		file, err := pkg.Root.Open(c.File)
		if err != nil {
			return errors.Join(errUpdateCaption, err)
		}
		defer func(file *os.File) {
			_ = file.Close()
		}(file)
		call = call.Media(file)
	}
	if c.OnBehalfOf != "" {
		call = call.OnBehalfOf(c.OnBehalfOf)
	}
	if c.OnBehalfOfContentOwner != "" {
		call = call.OnBehalfOfContentOwner(c.OnBehalfOfContentOwner)
	}

	res, err := call.Do()
	if err != nil {
		return errors.Join(errUpdateCaption, err)
	}

	common.PrintResult(c.Output, res, writer, "Caption updated: %s\n", res.Id)
	return nil
}

func (c *Caption) Delete(writer io.Writer) error {
	if err := c.EnsureService(); err != nil {
		return err
	}
	for _, id := range c.Ids {
		call := c.Service.Captions.Delete(id)
		if c.OnBehalfOf != "" {
			call = call.OnBehalfOf(c.OnBehalfOf)
		}
		if c.OnBehalfOfContentOwner != "" {
			call = call.OnBehalfOfContentOwner(c.OnBehalfOfContentOwner)
		}

		err := call.Do()
		if err != nil {
			return errors.Join(errDeleteCaption, err)
		}

		_, _ = fmt.Fprintf(writer, "Caption %s deleted\n", id)
	}
	return nil
}

func (c *Caption) Download(writer io.Writer) error {
	if err := c.EnsureService(); err != nil {
		return err
	}
	call := c.Service.Captions.Download(c.Ids[0])
	if c.Tfmt != "" {
		call = call.Tfmt(c.Tfmt)
	}
	if c.Tlang != "" {
		call = call.Tlang(c.Tlang)
	}
	if c.OnBehalfOf != "" {
		call = call.OnBehalfOf(c.OnBehalfOf)
	}
	if c.OnBehalfOfContentOwner != "" {
		call = call.OnBehalfOfContentOwner(c.OnBehalfOfContentOwner)
	}

	res, err := call.Download()
	if err != nil {
		return errors.Join(errDownloadCaption, err)
	}
	defer func(Body io.ReadCloser) {
		_ = Body.Close()
	}(res.Body)

	body, err := io.ReadAll(res.Body)
	if err != nil {
		return errors.Join(errDownloadCaption, err)
	}

	file, err := pkg.Root.OpenFile(
		c.File, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600,
	)
	if err != nil {
		return errors.Join(errDownloadCaption, err)
	}
	defer func(file *os.File) {
		_ = file.Close()
	}(file)

	_, err = file.Write(body)
	if err != nil {
		return errors.Join(errDownloadCaption, err)
	}

	_, _ = fmt.Fprintf(writer, "Caption %s downloaded to %s\n", c.Ids[0], c.File)
	return nil
}

func WithFile(file string) Option {
	return func(c *Caption) {
		c.File = file
	}
}

func WithAudioTrackType(audioTrackType string) Option {
	return func(c *Caption) {
		c.AudioTrackType = audioTrackType
	}
}

func WithIsAutoSynced(isAutoSynced *bool) Option {
	return func(c *Caption) {
		if isAutoSynced != nil {
			c.IsAutoSynced = isAutoSynced
		}
	}
}

func WithIsCC(isCC *bool) Option {
	return func(c *Caption) {
		if isCC != nil {
			c.IsCC = isCC
		}
	}
}

func WithIsDraft(isDraft *bool) Option {
	return func(c *Caption) {
		if isDraft != nil {
			c.IsDraft = isDraft
		}
	}
}

func WithIsEasyReader(isEasyReader *bool) Option {
	return func(c *Caption) {
		if isEasyReader != nil {
			c.IsEasyReader = isEasyReader
		}
	}
}

func WithIsLarge(isLarge *bool) Option {
	return func(c *Caption) {
		if isLarge != nil {
			c.IsLarge = isLarge
		}
	}
}

func WithLanguage(language string) Option {
	return func(c *Caption) {
		c.Language = language
	}
}

func WithName(name string) Option {
	return func(c *Caption) {
		c.Name = name
	}
}

func WithTrackKind(trackKind string) Option {
	return func(c *Caption) {
		c.TrackKind = trackKind
	}
}

func WithOnBehalfOf(onBehalfOf string) Option {
	return func(c *Caption) {
		c.OnBehalfOf = onBehalfOf
	}
}

func WithVideoId(videoId string) Option {
	return func(c *Caption) {
		c.VideoId = videoId
	}
}

func WithTfmt(tfmt string) Option {
	return func(c *Caption) {
		c.Tfmt = tfmt
	}
}

func WithTlang(tlang string) Option {
	return func(c *Caption) {
		c.Tlang = tlang
	}
}

var (
	WithIds     = common.WithIds[*Caption]
	WithParts   = common.WithParts[*Caption]
	WithOutput  = common.WithOutput[*Caption]
	WithService = common.WithService[*Caption]

	WithOnBehalfOfContentOwner = common.WithOnBehalfOfContentOwner[*Caption]
)