Initial AI media hub implementation

backend/handlers/api.go

@@ -0,0 +1,282 @@
+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()
+}
+
+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/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"`
+		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
+	}
+
+	go a.runDownload(recordID, req.URL, req.Start, req.End, outputPath)
+	c.JSON(http.StatusAccepted, gin.H{"message": "download started", "recordId": recordID})
+}
+
+func (a *App) runDownload(recordID int64, url, start, end, 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, "--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)
+		}
+	}
+
+	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
+}