Sync from /srv/compose/unified-media-manager
This commit is contained in:
Christopher Mayor
378
internal/service/subtitle.go
Normal file
378
internal/service/subtitle.go
Normal file
@@ -0,0 +1,378 @@
|
||||
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, ".")
|
||||
}
|
||||
Reference in New Issue
Block a user