331 lines
8.9 KiB
Go
331 lines
8.9 KiB
Go
package handlers
|
|
|
|
import (
|
|
"bufio"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"regexp"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"ai-media-hub/backend/models"
|
|
"ai-media-hub/backend/services"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/google/uuid"
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
type App struct {
|
|
DB *sql.DB
|
|
DownloadsDir string
|
|
WorkerScript string
|
|
SearchService *services.SearchService
|
|
GeminiService *services.GeminiService
|
|
Hub *Hub
|
|
}
|
|
|
|
type Hub struct {
|
|
clients map[*websocket.Conn]bool
|
|
mu sync.Mutex
|
|
}
|
|
|
|
func NewHub() *Hub {
|
|
return &Hub{clients: map[*websocket.Conn]bool{}}
|
|
}
|
|
|
|
func (h *Hub) Broadcast(event string, data any) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
|
|
payload, _ := json.Marshal(gin.H{"event": event, "data": data})
|
|
for conn := range h.clients {
|
|
_ = conn.WriteMessage(websocket.TextMessage, payload)
|
|
}
|
|
}
|
|
|
|
func (h *Hub) Add(conn *websocket.Conn) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
h.clients[conn] = true
|
|
}
|
|
|
|
func (h *Hub) Remove(conn *websocket.Conn) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
delete(h.clients, conn)
|
|
_ = conn.Close()
|
|
}
|
|
|
|
type PreviewResponse struct {
|
|
Title string `json:"title"`
|
|
Thumbnail string `json:"thumbnail"`
|
|
Duration string `json:"duration"`
|
|
DurationSeconds int `json:"durationSeconds"`
|
|
StartDefault string `json:"startDefault"`
|
|
EndDefault string `json:"endDefault"`
|
|
Qualities []map[string]any `json:"qualities"`
|
|
}
|
|
|
|
func RegisterRoutes(router *gin.Engine, app *App) {
|
|
router.GET("/healthz", func(c *gin.Context) {
|
|
c.JSON(http.StatusOK, gin.H{"status": "ok"})
|
|
})
|
|
router.GET("/ws", app.handleWS)
|
|
router.GET("/api/history/check", app.checkDuplicate)
|
|
router.POST("/api/download/preview", app.previewDownload)
|
|
router.POST("/api/upload", app.uploadFile)
|
|
router.POST("/api/download", app.startDownload)
|
|
router.POST("/api/search", app.searchMedia)
|
|
}
|
|
|
|
func (a *App) handleWS(c *gin.Context) {
|
|
upgrader := websocket.Upgrader{
|
|
CheckOrigin: func(r *http.Request) bool { return true },
|
|
}
|
|
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
|
if err != nil {
|
|
return
|
|
}
|
|
a.Hub.Add(conn)
|
|
defer a.Hub.Remove(conn)
|
|
|
|
for {
|
|
if _, _, err := conn.ReadMessage(); err != nil {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (a *App) checkDuplicate(c *gin.Context) {
|
|
url := strings.TrimSpace(c.Query("url"))
|
|
if url == "" {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "url is required"})
|
|
return
|
|
}
|
|
record, err := models.FindByURL(a.DB, url)
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
c.JSON(http.StatusOK, gin.H{"exists": record != nil, "record": record})
|
|
}
|
|
|
|
func (a *App) uploadFile(c *gin.Context) {
|
|
file, err := c.FormFile("file")
|
|
if err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "file is required"})
|
|
return
|
|
}
|
|
|
|
a.Hub.Broadcast("progress", gin.H{"type": "upload", "status": "started", "progress": 5, "filename": file.Filename})
|
|
|
|
safeName := normalizeFilename(file.Filename)
|
|
targetPath := filepath.Join(a.DownloadsDir, safeName)
|
|
if err := c.SaveUploadedFile(file, targetPath); err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
|
|
a.Hub.Broadcast("progress", gin.H{"type": "upload", "status": "completed", "progress": 100, "filename": safeName})
|
|
c.JSON(http.StatusOK, gin.H{"message": "uploaded", "path": targetPath, "filename": safeName})
|
|
}
|
|
|
|
func (a *App) startDownload(c *gin.Context) {
|
|
var req struct {
|
|
URL string `json:"url"`
|
|
Start string `json:"start"`
|
|
End string `json:"end"`
|
|
Quality string `json:"quality"`
|
|
Force bool `json:"force"`
|
|
}
|
|
if err := c.ShouldBindJSON(&req); err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
|
|
rec, err := models.FindByURL(a.DB, req.URL)
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
if rec != nil && !req.Force {
|
|
c.JSON(http.StatusConflict, gin.H{"error": "duplicate url", "record": rec})
|
|
return
|
|
}
|
|
|
|
outputBase := uuid.NewString()
|
|
outputPath := filepath.Join(a.DownloadsDir, outputBase+".mp4")
|
|
recordID, err := models.InsertDownload(a.DB, req.URL, detectSource(req.URL), outputPath, "queued")
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
|
|
quality := strings.TrimSpace(req.Quality)
|
|
if quality == "" {
|
|
quality = "best"
|
|
}
|
|
|
|
go a.runDownload(recordID, req.URL, req.Start, req.End, quality, outputPath)
|
|
c.JSON(http.StatusAccepted, gin.H{"message": "download started", "recordId": recordID})
|
|
}
|
|
|
|
func (a *App) previewDownload(c *gin.Context) {
|
|
var req struct {
|
|
URL string `json:"url"`
|
|
}
|
|
if err := c.ShouldBindJSON(&req); err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
if strings.TrimSpace(req.URL) == "" {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "url is required"})
|
|
return
|
|
}
|
|
|
|
cmd := exec.Command("python3", a.WorkerScript, "--mode", "probe", "--url", req.URL, "--output", filepath.Join(a.DownloadsDir, "probe.tmp"))
|
|
output, err := cmd.CombinedOutput()
|
|
if err != nil {
|
|
c.JSON(http.StatusBadGateway, gin.H{"error": strings.TrimSpace(string(output))})
|
|
return
|
|
}
|
|
|
|
var preview PreviewResponse
|
|
if err := json.Unmarshal(output, &preview); err != nil {
|
|
c.JSON(http.StatusBadGateway, gin.H{"error": "invalid probe response"})
|
|
return
|
|
}
|
|
c.JSON(http.StatusOK, preview)
|
|
}
|
|
|
|
func (a *App) runDownload(recordID int64, url, start, end, quality, outputPath string) {
|
|
a.Hub.Broadcast("progress", gin.H{"type": "download", "status": "queued", "progress": 0, "url": url})
|
|
cmd := exec.Command("python3", a.WorkerScript, "--url", url, "--start", start, "--end", end, "--quality", quality, "--output", outputPath)
|
|
stdout, err := cmd.StdoutPipe()
|
|
if err != nil {
|
|
a.Hub.Broadcast("progress", gin.H{"type": "download", "status": "error", "progress": 0, "message": err.Error()})
|
|
_ = models.MarkDownloadCompleted(a.DB, recordID, "failed")
|
|
return
|
|
}
|
|
cmd.Stderr = cmd.Stdout
|
|
|
|
if err := cmd.Start(); err != nil {
|
|
a.Hub.Broadcast("progress", gin.H{"type": "download", "status": "error", "progress": 0, "message": err.Error()})
|
|
_ = models.MarkDownloadCompleted(a.DB, recordID, "failed")
|
|
return
|
|
}
|
|
|
|
scanner := bufio.NewScanner(stdout)
|
|
for scanner.Scan() {
|
|
line := scanner.Bytes()
|
|
var msg map[string]any
|
|
if err := json.Unmarshal(line, &msg); err == nil {
|
|
msg["type"] = "download"
|
|
a.Hub.Broadcast("progress", msg)
|
|
}
|
|
}
|
|
if err := scanner.Err(); err != nil {
|
|
a.Hub.Broadcast("progress", gin.H{"type": "download", "status": "error", "progress": 100, "message": err.Error()})
|
|
}
|
|
|
|
status := "completed"
|
|
if err := cmd.Wait(); err != nil {
|
|
status = "failed"
|
|
a.Hub.Broadcast("progress", gin.H{"type": "download", "status": "error", "progress": 100, "message": err.Error()})
|
|
}
|
|
_ = models.MarkDownloadCompleted(a.DB, recordID, status)
|
|
}
|
|
|
|
func (a *App) searchMedia(c *gin.Context) {
|
|
var req struct {
|
|
Query string `json:"query"`
|
|
}
|
|
if err := c.ShouldBindJSON(&req); err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
if strings.TrimSpace(req.Query) == "" {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "query is required"})
|
|
return
|
|
}
|
|
|
|
results, err := a.SearchService.SearchMedia(req.Query)
|
|
if err != nil {
|
|
c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
|
|
recommended, err := a.GeminiService.Recommend(req.Query, results)
|
|
if err != nil {
|
|
fallback := make([]services.AIRecommendation, 0, min(4, len(results)))
|
|
for _, result := range results[:min(4, len(results))] {
|
|
fallback = append(fallback, services.AIRecommendation{
|
|
Title: result.Title,
|
|
Link: result.Link,
|
|
ThumbnailURL: result.ThumbnailURL,
|
|
Source: result.Source,
|
|
Reason: "Gemini recommendation failed, showing raw search result.",
|
|
Recommended: true,
|
|
})
|
|
}
|
|
c.JSON(http.StatusOK, gin.H{"results": fallback, "warning": err.Error()})
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, gin.H{"results": recommended})
|
|
}
|
|
|
|
func normalizeFilename(name string) string {
|
|
base := strings.ToLower(strings.TrimSpace(name))
|
|
ext := filepath.Ext(base)
|
|
base = strings.TrimSuffix(base, ext)
|
|
re := regexp.MustCompile(`[^a-z0-9]+`)
|
|
base = strings.Trim(re.ReplaceAllString(base, "-"), "-")
|
|
if base == "" {
|
|
base = fmt.Sprintf("upload-%d", time.Now().Unix())
|
|
}
|
|
if ext == "" {
|
|
ext = ".bin"
|
|
}
|
|
return base + ext
|
|
}
|
|
|
|
func detectSource(url string) string {
|
|
switch {
|
|
case strings.Contains(url, "youtube"):
|
|
return "YouTube"
|
|
case strings.Contains(url, "tiktok"):
|
|
return "TikTok"
|
|
default:
|
|
return "direct"
|
|
}
|
|
}
|
|
|
|
func min(a, b int) int {
|
|
if a < b {
|
|
return a
|
|
}
|
|
return b
|
|
}
|
|
|
|
func EnsurePaths(downloadsDir, workerScript string) error {
|
|
if err := os.MkdirAll(downloadsDir, 0o755); err != nil {
|
|
return err
|
|
}
|
|
if _, err := os.Stat(workerScript); err != nil {
|
|
if errors.Is(err, os.ErrNotExist) {
|
|
return fmt.Errorf("worker script not found: %s", workerScript)
|
|
}
|
|
return err
|
|
}
|
|
return nil
|
|
}
|