Files
ai-media-hub/backend/handlers/api.go
T
AI Assistant 4dbb963256
build-push / docker (push) Successful in 4m5s
Proxy and cache preview media
2026-03-16 14:28:00 +09:00

700 lines
22 KiB
Go

package handlers
import (
"bufio"
"crypto/sha1"
"database/sql"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"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
PreviewCacheDir 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"`
PreviewStreamURL string `json:"previewStreamUrl"`
Duration string `json:"duration"`
DurationSeconds int `json:"durationSeconds"`
StartDefault string `json:"startDefault"`
EndDefault string `json:"endDefault"`
Qualities []map[string]any `json:"qualities"`
}
type searchDebugSummary struct {
Total int `json:"total"`
BySource map[string]int `json:"bySource"`
WithPreview int `json:"withPreview"`
WithThumbnail int `json:"withThumbnail"`
Top []map[string]any `json:"top"`
Warning string `json:"warning,omitempty"`
DurationMS int64 `json:"durationMs,omitempty"`
GeminiCandidateCap int `json:"geminiCandidateCap,omitempty"`
}
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.GET("/api/preview/stream", app.streamPreview)
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) debug(message string, data any) {
a.Hub.Broadcast("debug", gin.H{"message": message, "data": data})
}
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) streamPreview(c *gin.Context) {
target := strings.TrimSpace(c.Query("url"))
if target == "" {
c.JSON(http.StatusBadRequest, gin.H{"error": "url is required"})
return
}
req, err := http.NewRequest(http.MethodGet, target, nil)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
req.Header.Set("User-Agent", "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36")
req.Header.Set("Accept-Language", "en-US,en;q=0.9")
req.Header.Set("Referer", inferPreviewReferer(target))
resp, err := a.SearchService.Client.Do(req)
if err != nil {
c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
return
}
defer resp.Body.Close()
if resp.StatusCode >= 300 {
c.JSON(http.StatusBadGateway, gin.H{"error": fmt.Sprintf("preview source returned %d", resp.StatusCode)})
return
}
contentType := resp.Header.Get("Content-Type")
if strings.Contains(strings.ToLower(target), ".m3u8") || strings.Contains(strings.ToLower(contentType), "mpegurl") {
body, err := io.ReadAll(io.LimitReader(resp.Body, 2*1024*1024))
if err != nil {
c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
return
}
rewritten := rewriteM3U8Playlist(string(body), target)
c.Header("Content-Type", "application/vnd.apple.mpegurl")
c.String(http.StatusOK, rewritten)
return
}
if isCacheablePreview(target, contentType) {
if cachedPath, err := a.cachePreviewResponse(target, contentType, resp.Body); err == nil {
c.File(cachedPath)
return
}
}
if contentType != "" {
c.Header("Content-Type", contentType)
}
c.Status(http.StatusOK)
_, _ = io.Copy(c.Writer, resp.Body)
}
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
}
a.debug("download preview requested", gin.H{"url": req.URL})
cmd := exec.Command("python3", a.WorkerScript, "--mode", "probe", "--url", req.URL)
output, err := cmd.CombinedOutput()
if err != nil {
a.debug("download preview failed", gin.H{"url": req.URL, "output": string(output), "error": err.Error()})
c.JSON(http.StatusBadGateway, gin.H{"error": summarizeOutput("download preview probe failed", output, err)})
return
}
var preview PreviewResponse
if err := json.Unmarshal(output, &preview); err != nil {
a.debug("download preview invalid json", gin.H{"url": req.URL, "output": string(output)})
c.JSON(http.StatusBadGateway, gin.H{"error": summarizeOutput("download preview returned invalid JSON", output, err)})
return
}
a.debug("download preview succeeded", preview)
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})
a.debug("download command started", gin.H{"url": url, "start": start, "end": end, "quality": quality, "outputPath": outputPath})
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.debug("download worker event", msg)
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()})
a.debug("download command failed", gin.H{"url": url, "error": err.Error()})
}
a.debug("download command completed", gin.H{"url": url, "status": status, "outputPath": outputPath})
_ = models.MarkDownloadCompleted(a.DB, recordID, status)
}
func (a *App) searchMedia(c *gin.Context) {
started := time.Now()
var req struct {
Query string `json:"query"`
Platforms []string `json:"platforms"`
}
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
}
a.Hub.Broadcast("progress", gin.H{"type": "search", "status": "expanding query with Gemini", "progress": 10})
queryVariants, _ := a.GeminiService.ExpandQuery(req.Query)
if len(queryVariants) == 0 {
queryVariants = []string{req.Query}
}
a.debug("search query variants", gin.H{
"query": req.Query,
"platforms": req.Platforms,
"variants": queryVariants,
"variantCount": len(queryVariants),
"requestIdHint": time.Now().UnixNano(),
})
enabledPlatforms := normalizePlatforms(req.Platforms)
a.Hub.Broadcast("progress", gin.H{"type": "search", "status": "searching " + selectedPlatformLabel(enabledPlatforms), "progress": 35})
results, err := a.SearchService.SearchMedia(queryVariants, enabledPlatforms)
if err != nil {
a.debug("search backend failed", gin.H{"error": err.Error(), "variants": queryVariants, "durationMs": time.Since(started).Milliseconds()})
a.Hub.Broadcast("progress", gin.H{"type": "search", "status": "search failed", "progress": 100, "message": err.Error()})
c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
return
}
a.debug("search backend summary", summarizeSearchResults(results, time.Since(started), 0, ""))
if len(results) == 0 {
warning := "SearXNG returned no renderable results."
a.Hub.Broadcast("progress", gin.H{"type": "search", "status": "no renderable search results", "progress": 100, "message": warning})
c.JSON(http.StatusOK, gin.H{"results": []services.AIRecommendation{}, "warning": warning})
return
}
a.Hub.Broadcast("progress", gin.H{"type": "search", "status": "ranking thumbnail candidates", "progress": 55})
rankQuery := req.Query
if len(queryVariants) > 0 {
rankQuery = strings.Join(queryVariants[:min(len(queryVariants), 3)], " ")
}
scored := services.RankSearchResults(rankQuery, results)
a.debug("search ranked summary", summarizeSearchResults(scored, time.Since(started), services.GeminiCandidateLimit(len(scored)), ""))
a.Hub.Broadcast("progress", gin.H{"type": "search", "status": "analyzing all candidate visuals with Gemini Vision", "progress": 75})
recommended, geminiStats, geminiErr := services.EvaluateAllCandidatesWithGemini(a.GeminiService, req.Query, scored)
a.debug("search gemini evaluation", geminiStats)
if services.NeedsSupplementalExploration(recommended) {
a.Hub.Broadcast("progress", gin.H{"type": "search", "status": "Gemini 평가가 약해 추가 후보를 탐색하는 중", "progress": 82})
explorationQueries := buildSupplementalQueries(req.Query, queryVariants)
extraResults, extraErr := a.SearchService.SearchMedia(explorationQueries, enabledPlatforms)
if extraErr == nil && len(extraResults) > 0 {
results = mergeSearchResults(results, extraResults)
scored = services.RankSearchResults(strings.Join(explorationQueries[:min(len(explorationQueries), 3)], " "), results)
recommended, geminiStats, geminiErr = services.EvaluateAllCandidatesWithGemini(a.GeminiService, req.Query, scored)
a.debug("search supplemental query variants", gin.H{"variants": explorationQueries, "variantCount": len(explorationQueries)})
a.debug("search gemini evaluation after supplemental search", geminiStats)
}
}
if geminiErr != nil && len(recommended) == 0 {
warning := geminiErr.Error()
fallback := services.BuildFallbackRecommendations(scored, 20, "")
a.debug("search fallback summary", summarizeRecommendationResults(fallback, time.Since(started), warning))
a.Hub.Broadcast("progress", gin.H{"type": "search", "status": "Gemini Vision fallback to ranked results", "progress": 90, "message": warning})
c.JSON(http.StatusOK, gin.H{"results": fallback, "warning": warning, "queries": queryVariants})
return
}
merged := services.MergeRecommendations(recommended, scored, 20)
merged = services.RandomizeTopRecommendations(merged, 8)
warning := ""
if geminiErr != nil {
warning = geminiErr.Error()
}
a.debug("search complete summary", summarizeRecommendationResults(merged, time.Since(started), warning))
response := gin.H{"results": merged, "queries": queryVariants}
if warning != "" {
response["warning"] = warning
}
a.Hub.Broadcast("progress", gin.H{"type": "search", "status": "search complete", "progress": 100})
c.JSON(http.StatusOK, response)
}
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 normalizePlatforms(platforms []string) map[string]bool {
if len(platforms) == 0 {
return map[string]bool{
"envato": true,
"artgrid": true,
"google video": true,
}
}
normalized := map[string]bool{}
for _, item := range platforms {
switch strings.ToLower(strings.TrimSpace(item)) {
case "envato":
normalized["envato"] = true
case "artgrid":
normalized["artgrid"] = true
case "google video", "google_video", "google":
normalized["google video"] = true
}
}
return normalized
}
func selectedPlatformLabel(platforms map[string]bool) string {
labels := make([]string, 0, len(platforms))
if platforms["envato"] {
labels = append(labels, "Envato")
}
if platforms["artgrid"] {
labels = append(labels, "Artgrid")
}
if platforms["google video"] {
labels = append(labels, "Google Video")
}
if len(labels) == 0 {
return "selected platforms"
}
return strings.Join(labels, ", ")
}
func buildSupplementalQueries(query string, existing []string) []string {
candidates := append([]string{}, existing...)
candidates = append(candidates,
query+" cinematic stock footage",
query+" editorial b-roll",
query+" establishing shot",
query+" drone footage",
)
seen := map[string]bool{}
result := make([]string, 0, len(candidates))
for _, item := range candidates {
trimmed := strings.Join(strings.Fields(strings.TrimSpace(item)), " ")
if trimmed == "" {
continue
}
key := strings.ToLower(trimmed)
if seen[key] {
continue
}
seen[key] = true
result = append(result, trimmed)
}
return result
}
func mergeSearchResults(base, extra []services.SearchResult) []services.SearchResult {
merged := make([]services.SearchResult, 0, len(base)+len(extra))
seen := map[string]bool{}
for _, item := range append(base, extra...) {
if item.Link == "" || seen[item.Link] {
continue
}
seen[item.Link] = true
merged = append(merged, item)
}
return merged
}
func summarizeSearchResults(results []services.SearchResult, duration time.Duration, geminiCap int, warning string) searchDebugSummary {
bySource := map[string]int{}
withPreview := 0
withThumbnail := 0
top := make([]map[string]any, 0, min(6, len(results)))
for idx, item := range results {
bySource[item.Source]++
if strings.TrimSpace(item.PreviewVideoURL) != "" {
withPreview++
}
if strings.TrimSpace(item.ThumbnailURL) != "" {
withThumbnail++
}
if idx < 6 {
top = append(top, map[string]any{
"title": truncateText(item.Title, 120),
"source": item.Source,
"hasPreview": item.PreviewVideoURL != "",
"hasThumbnail": item.ThumbnailURL != "",
"displayLink": item.DisplayLink,
"snippetSample": truncateText(item.Snippet, 160),
})
}
}
return searchDebugSummary{
Total: len(results),
BySource: bySource,
WithPreview: withPreview,
WithThumbnail: withThumbnail,
Top: top,
Warning: warning,
DurationMS: duration.Milliseconds(),
GeminiCandidateCap: geminiCap,
}
}
func summarizeRecommendationResults(results []services.AIRecommendation, duration time.Duration, warning string) searchDebugSummary {
bySource := map[string]int{}
withPreview := 0
withThumbnail := 0
top := make([]map[string]any, 0, min(6, len(results)))
for idx, item := range results {
bySource[item.Source]++
if strings.TrimSpace(item.PreviewVideoURL) != "" {
withPreview++
}
if strings.TrimSpace(item.ThumbnailURL) != "" {
withThumbnail++
}
if idx < 6 {
top = append(top, map[string]any{
"title": truncateText(item.Title, 120),
"source": item.Source,
"hasPreview": item.PreviewVideoURL != "",
"hasThumbnail": item.ThumbnailURL != "",
"reasonSample": truncateText(item.Reason, 120),
"snippetSample": truncateText(item.Snippet, 160),
})
}
}
return searchDebugSummary{
Total: len(results),
BySource: bySource,
WithPreview: withPreview,
WithThumbnail: withThumbnail,
Top: top,
Warning: warning,
DurationMS: duration.Milliseconds(),
}
}
func truncateText(text string, limit int) string {
trimmed := strings.TrimSpace(text)
if len(trimmed) <= limit {
return trimmed
}
return trimmed[:limit] + "..."
}
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
}
func inferPreviewReferer(target string) string {
lower := strings.ToLower(target)
switch {
case strings.Contains(lower, "envatousercontent.com"), strings.Contains(lower, "elements.envato.com"):
return "https://elements.envato.com/"
case strings.Contains(lower, "artgrid"), strings.Contains(lower, "artlist"):
return "https://artgrid.io/"
default:
return ""
}
}
func rewriteM3U8Playlist(body, target string) string {
lines := strings.Split(body, "\n")
baseURL := target
if idx := strings.LastIndex(baseURL, "/"); idx >= 0 {
baseURL = baseURL[:idx+1]
}
for idx, line := range lines {
trimmed := strings.TrimSpace(line)
if trimmed == "" || strings.HasPrefix(trimmed, "#") {
continue
}
resolved := trimmed
if !strings.HasPrefix(trimmed, "http://") && !strings.HasPrefix(trimmed, "https://") {
resolved = baseURL + strings.TrimPrefix(trimmed, "/")
}
lines[idx] = "/api/preview/stream?url=" + url.QueryEscape(resolved)
}
return strings.Join(lines, "\n")
}
func isCacheablePreview(target, contentType string) bool {
lower := strings.ToLower(target + " " + contentType)
return strings.Contains(lower, ".mp4") || strings.Contains(lower, "video/mp4")
}
func (a *App) cachePreviewResponse(target, contentType string, body io.Reader) (string, error) {
if a.PreviewCacheDir == "" {
return "", fmt.Errorf("preview cache dir is not configured")
}
if err := os.MkdirAll(a.PreviewCacheDir, 0o755); err != nil {
return "", err
}
sum := sha1.Sum([]byte(target))
ext := ".bin"
if strings.Contains(strings.ToLower(target), ".mp4") || strings.Contains(strings.ToLower(contentType), "video/mp4") {
ext = ".mp4"
}
path := filepath.Join(a.PreviewCacheDir, fmt.Sprintf("%x%s", sum, ext))
file, err := os.Create(path)
if err != nil {
return "", err
}
defer file.Close()
if _, err := io.Copy(file, body); err != nil {
return "", err
}
return path, nil
}
func summarizeOutput(prefix string, output []byte, err error) string {
trimmed := strings.TrimSpace(string(output))
if trimmed == "" && err != nil {
return prefix + ": " + err.Error()
}
if trimmed == "" {
return prefix
}
if len(trimmed) > 1200 {
trimmed = trimmed[:1200] + "..."
}
return prefix + ": " + trimmed
}