unified-media-manager/internal/service/subtitle.go

package service

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/TopherMayor/unified-media-manager/internal/db"
)

type SubtitleSearchResult struct {
	ID            string `json:"id"`
	FileName      string `json:"file_name"`
	Language      string `json:"language"`
	LanguageCode  string `json:"language_code"`
	HI            bool   `json:"hi"`
	Forced        bool   `json:"forced"`
	DownloadCount int    `json:"download_count"`
	ReleaseName   string `json:"release_name"`
	Provider      string `json:"provider"`
}

type SubtitleFile struct {
	Path         string `json:"path"`
	Language     string `json:"language"`
	LanguageCode string `json:"language_code"`
	HI           bool   `json:"hi"`
	Forced       bool   `json:"forced"`
	Source       string `json:"source"`
}

type SubtitleSearchOptions struct {
	LanguageCodes []string
	HI            bool
	Forced        bool
}

type SubtitleService struct {
	db          *db.DB
	apiKey      string
	baseURL     string
	httpClient  *http.Client
	ffmpegPath  string
	ffprobePath string
}

func NewSubtitleService(database *db.DB, apiKey string) *SubtitleService {
	return &SubtitleService{
		db:      database,
		apiKey:  apiKey,
		baseURL: "https://api.opensubtitles.com/api/v1",
		httpClient: &http.Client{
			Timeout: 15 * time.Second,
		},
		ffmpegPath:  "ffmpeg",
		ffprobePath: "ffprobe",
	}
}

type osSearchResponse struct {
	TotalPages int          `json:"total_pages"`
	Data       []osSubtitle `json:"data"`
}

type osSubtitle struct {
	ID                string `json:"id"`
	FileName          string `json:"file_name"`
	Language          string `json:"language"`
	MovieFileNameMatch string `json:"movie_file_name_match"`
	DownloadCount     int    `json:"download_count"`
	HearingImpaired   bool   `json:"hearing_impaired"`
	ForeignPartsOnly  bool   `json:"foreign_parts_only"`
	ReleaseName       string `json:"release_name"`
	MovieHash         string `json:"movie_hash"`
}

type osDownloadResponse struct {
	Link     string `json:"link"`
	FileName string `json:"file_name"`
}

type osLoginResponse struct {
	Token string `json:"token"`
}

type ffprobeStream struct {
	Streams []ffprobeStreamInfo `json:"streams"`
}

type ffprobeStreamInfo struct {
	Index       int               `json:"index"`
	CodecName   string            `json:"codec_name"`
	CodecType   string            `json:"codec_type"`
	Disposition ffprobeDisposition `json:"disposition"`
	Tags        ffprobeTags       `json:"tags"`
}

type ffprobeDisposition struct {
	Forced int `json:"forced"`
}

type ffprobeTags struct {
	Language string `json:"language"`
	Title    string `json:"title"`
}

func (s *SubtitleService) osLogin(ctx context.Context) (string, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, s.baseURL+"/login", nil)
	if err != nil {
		return "", fmt.Errorf("create login request: %w", err)
	}
	req.Header.Set("Api-Key", s.apiKey)
	req.Header.Set("Content-Type", "application/json")

	resp, err := s.httpClient.Do(req)
	if err != nil {
		return "", fmt.Errorf("opensubtitles login: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("opensubtitles login failed: status %d", resp.StatusCode)
	}

	var loginResp osLoginResponse
	if err := json.NewDecoder(resp.Body).Decode(&loginResp); err != nil {
		return "", fmt.Errorf("decode login response: %w", err)
	}

	return loginResp.Token, nil
}

func (s *SubtitleService) Search(ctx context.Context, query string, opts SubtitleSearchOptions) ([]SubtitleSearchResult, error) {
	token, err := s.osLogin(ctx)
	if err != nil {
		return nil, fmt.Errorf("opensubtitles login: %w", err)
	}

	langs := strings.Join(opts.LanguageCodes, ",")
	url := fmt.Sprintf("%s/subtitles?query=%s&languages=%s", s.baseURL, query, langs)
	if opts.HI {
		url += "&hearing_impaired=true"
	}
	if opts.Forced {
		url += "&foreign_parts_only=true"
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, fmt.Errorf("create search request: %w", err)
	}
	req.Header.Set("Api-Key", s.apiKey)
	req.Header.Set("Authorization", "Bearer "+token)

	resp, err := s.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("search opensubtitles: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("opensubtitles search failed: status %d", resp.StatusCode)
	}

	var searchResp osSearchResponse
	if err := json.NewDecoder(resp.Body).Decode(&searchResp); err != nil {
		return nil, fmt.Errorf("decode search response: %w", err)
	}

	var results []SubtitleSearchResult
	for _, sub := range searchResp.Data {
		langCode := ""
		if parts := strings.Split(sub.Language, "-"); len(parts) > 0 {
			langCode = strings.ToLower(parts[0])
		}

		results = append(results, SubtitleSearchResult{
			ID:            sub.ID,
			FileName:      sub.FileName,
			Language:      sub.Language,
			LanguageCode:  langCode,
			HI:            sub.HearingImpaired,
			Forced:        sub.ForeignPartsOnly,
			DownloadCount: sub.DownloadCount,
			ReleaseName:   sub.ReleaseName,
			Provider:      "opensubtitles",
		})
	}

	return results, nil
}

func (s *SubtitleService) Download(ctx context.Context, subtitleID string, targetDir string, baseName string, langCode string, hi bool, forced bool) (*SubtitleFile, error) {
	token, err := s.osLogin(ctx)
	if err != nil {
		return nil, fmt.Errorf("opensubtitles login: %w", err)
	}

	body, _ := json.Marshal(map[string]string{"file_id": subtitleID})
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, s.baseURL+"/download", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("create download request: %w", err)
	}
	req.Header.Set("Api-Key", s.apiKey)
	req.Header.Set("Authorization", "Bearer "+token)
	req.Header.Set("Content-Type", "application/json")

	resp, err := s.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("request subtitle download: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("subtitle download request failed: status %d", resp.StatusCode)
	}

	var downloadResp osDownloadResponse
	if err := json.NewDecoder(resp.Body).Decode(&downloadResp); err != nil {
		return nil, fmt.Errorf("decode download response: %w", err)
	}

	fileReq, err := http.NewRequestWithContext(ctx, http.MethodGet, downloadResp.Link, nil)
	if err != nil {
		return nil, fmt.Errorf("create file download request: %w", err)
	}

	fileResp, err := s.httpClient.Do(fileReq)
	if err != nil {
		return nil, fmt.Errorf("download subtitle file: %w", err)
	}
	defer fileResp.Body.Close()

	if fileResp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("subtitle file download failed: status %d", fileResp.StatusCode)
	}

	filename := baseName + "." + langCode + ".srt"
	if hi {
		filename = baseName + "." + langCode + ".sdh.srt"
	}
	if forced {
		filename = baseName + "." + langCode + ".forced.srt"
	}

	if err := os.MkdirAll(targetDir, 0755); err != nil {
		return nil, fmt.Errorf("create target directory: %w", err)
	}

	destPath := filepath.Join(targetDir, filename)
	f, err := os.Create(destPath)
	if err != nil {
		return nil, fmt.Errorf("create subtitle file: %w", err)
	}
	defer f.Close()

	if _, err := io.Copy(f, fileResp.Body); err != nil {
		return nil, fmt.Errorf("write subtitle file: %w", err)
	}

	return &SubtitleFile{
		Path:         destPath,
		Language:     langCode,
		LanguageCode: langCode,
		HI:           hi,
		Forced:       forced,
		Source:       "downloaded",
	}, nil
}

func (s *SubtitleService) ExtractSubtitles(ctx context.Context, mediaFilePath string, targetDir string, baseName string) ([]SubtitleFile, error) {
	probeCtx, probeCancel := context.WithTimeout(ctx, 30*time.Second)
	defer probeCancel()

	cmd := exec.CommandContext(probeCtx, s.ffprobePath,
		"-v", "quiet",
		"-print_format", "json",
		"-show_streams",
		"-select_streams", "s",
		mediaFilePath,
	)

	var stdout bytes.Buffer
	cmd.Stdout = &stdout
	if err := cmd.Run(); err != nil {
		return nil, fmt.Errorf("ffprobe subtitle streams: %w", err)
	}

	var probeResult ffprobeStream
	if err := json.Unmarshal(stdout.Bytes(), &probeResult); err != nil {
		return nil, fmt.Errorf("parse ffprobe output: %w", err)
	}

	if len(probeResult.Streams) == 0 {
		return nil, nil
	}

	var results []SubtitleFile
	for i, stream := range probeResult.Streams {
		langCode := strings.ToLower(stream.Tags.Language)
		if langCode == "" {
			langCode = "und"
		}

		hi := false
		if strings.Contains(strings.ToLower(stream.Tags.Title), "sdh") {
			hi = true
		}

		forced := stream.Disposition.Forced == 1

		filename := baseName + "." + langCode + ".srt"
		if hi {
			filename = baseName + "." + langCode + ".sdh.srt"
		}
		if forced {
			filename = baseName + "." + langCode + ".forced.srt"
		}

		outputPath := filepath.Join(targetDir, filename)

		extractCtx, extractCancel := context.WithTimeout(ctx, 30*time.Second)
		extractCmd := exec.CommandContext(extractCtx, s.ffmpegPath,
			"-i", mediaFilePath,
			"-map", fmt.Sprintf("0:s:%d", i),
			"-f", "srt",
			outputPath,
		)
		if err := extractCmd.Run(); err != nil {
			slog.Error("failed to extract subtitle stream", "error", err, "stream_index", i)
			extractCancel()
			continue
		}
		extractCancel()

		results = append(results, SubtitleFile{
			Path:         outputPath,
			Language:     langCode,
			LanguageCode: langCode,
			HI:           hi,
			Forced:       forced,
			Source:       "extracted",
		})
	}

	return results, nil
}

func BuildSubtitleBaseName(title string, year *int, season, episode int) string {
	parts := []string{sanitizeSubtitleName(title)}
	if year != nil {
		parts = append(parts, fmt.Sprintf("%d", *year))
	}
	if season > 0 && episode > 0 {
		parts = append(parts, fmt.Sprintf("S%02dE%02d", season, episode))
	}
	return strings.Join(parts, ".")
}

var nonAlphaNumRe = regexp.MustCompile(`[^a-zA-Z0-9]+`)

func sanitizeSubtitleName(s string) string {
	s = nonAlphaNumRe.ReplaceAllString(s, ".")
	for strings.Contains(s, "..") {
		s = strings.ReplaceAll(s, "..", ".")
	}
	return strings.Trim(s, ".")
}