yt-dlp-telegram-bot/upload.go
2026-03-02 13:11:38 +08:00

128 lines
4.0 KiB
Go

package main
import (
"bytes"
"context"
"fmt"
"io"
"math/big"
"os"
"github.com/dustin/go-humanize"
"github.com/flytam/filenamify"
"github.com/gotd/td/telegram/message"
"github.com/gotd/td/telegram/uploader"
"github.com/gotd/td/tg"
)
type Uploader struct{}
var dlUploader Uploader
func (p Uploader) Chunk(ctx context.Context, state uploader.ProgressState) error {
dlQueue.HandleProgressPercentUpdate(uploadStr, int(state.Uploaded*100/state.Total))
return nil
}
func (p *Uploader) UploadFile(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage, f io.ReadCloser, format, title string, width, height int) error {
// Get file size by seeking if it's a file
var fileSize int64
if file, ok := f.(*os.File); ok {
stat, err := file.Stat()
if err != nil {
return fmt.Errorf("getting file stat error: %w", err)
}
fileSize = stat.Size()
if params.MaxSize > 0 && fileSize > params.MaxSize {
return fmt.Errorf("file is too big, max. allowed size is %s", humanize.BigBytes(big.NewInt(int64(params.MaxSize))))
}
} else {
// Fallback: read to buffer for non-file readers
return p.uploadFromBuffer(ctx, entities, u, f, format, title, width, height)
}
fmt.Println(" got", fileSize, "bytes, uploading...")
dlQueue.currentlyDownloadedEntry.progressInfo = fmt.Sprint(" (", humanize.BigBytes(big.NewInt(fileSize)), ")")
// Reset file pointer to beginning
if _, err := f.(*os.File).Seek(0, 0); err != nil {
return fmt.Errorf("seeking file error: %w", err)
}
// Use uploader.NewUpload with progress callback
upload, err := telegramUploader.Upload(ctx, uploader.NewUpload("yt-dlp", f, fileSize))
if err != nil {
return fmt.Errorf("uploading %w", err)
}
// Now we have uploaded file handle, sending it as styled message. First, preparing message.
var document message.MediaOption
filename, _ := filenamify.Filenamify(title+"."+format, filenamify.Options{Replacement: " "})
if format == "mp3" {
document = message.UploadedDocument(upload).Filename(filename).Audio().Title(title)
} else {
doc := message.UploadedDocument(upload).Filename(filename).Video()
// Set resolution to help Telegram display correct aspect ratio
if width > 0 && height > 0 {
doc = doc.Resolution(width, height)
}
document = doc
}
// Sending message with media.
if _, err := telegramSender.Answer(entities, u).Media(ctx, document); err != nil {
return fmt.Errorf("send: %w", err)
}
return nil
}
func (p *Uploader) uploadFromBuffer(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage, f io.ReadCloser, format, title string, width, height int) error {
// Fallback for non-file io.ReadCloser - read all to buffer
buf := make([]byte, 0)
tempBuf := make([]byte, 8192)
for {
n, err := f.Read(tempBuf)
if err != nil && err != io.EOF {
return fmt.Errorf("reading to buffer error: %w", err)
}
if n == 0 {
break
}
buf = append(buf, tempBuf[:n]...)
if params.MaxSize > 0 && len(buf) > int(params.MaxSize) {
return fmt.Errorf("file is too big, max. allowed size is %s", humanize.BigBytes(big.NewInt(int64(params.MaxSize))))
}
}
fmt.Println(" got", len(buf), "bytes, uploading...")
dlQueue.currentlyDownloadedEntry.progressInfo = fmt.Sprint(" (", humanize.BigBytes(big.NewInt(int64(len(buf)))), ")")
// Use Upload with progress for buffer too
upload, err := telegramUploader.Upload(ctx, uploader.NewUpload("yt-dlp", bytes.NewReader(buf), int64(len(buf))))
if err != nil {
return fmt.Errorf("uploading %w", err)
}
var document message.MediaOption
filename, _ := filenamify.Filenamify(title+"."+format, filenamify.Options{Replacement: " "})
if format == "mp3" {
document = message.UploadedDocument(upload).Filename(filename).Audio().Title(title)
} else {
doc := message.UploadedDocument(upload).Filename(filename).Video()
// Set resolution to help Telegram display correct aspect ratio
if width > 0 && height > 0 {
doc = doc.Resolution(width, height)
}
document = doc
}
if _, err := telegramSender.Answer(entities, u).Media(ctx, document); err != nil {
return fmt.Errorf("send: %w", err)
}
return nil
}