ai-media-hub/backend/main.go

package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/exec"
	"sync"

	"github.com/gofiber/fiber/v2"
	"github.com/gofiber/fiber/v2/middleware/cors"
	"github.com/gofiber/fiber/v2/middleware/logger"
	"github.com/gofiber/websocket/v2"
)

func main() {
	app := fiber.New()

	app.Use(logger.New())
	app.Use(cors.New())

	// Sub-routes for views/assets
	app.Static("/", "./frontend")

	// Global Hub for Websocket
	type Client struct {
		Conn *websocket.Conn
	}
	var clients = make(map[*Client]bool)
	var clientsMu sync.Mutex

	broadcast := func(msg string) {
		clientsMu.Lock()
		defer clientsMu.Unlock()
		for client := range clients {
			if err := client.Conn.WriteMessage(websocket.TextMessage, []byte(msg)); err != nil {
				client.Conn.Close()
				delete(clients, client)
			}
		}
	}

	// API Routes
	api := app.Group("/api")
	
	api.Get("/search", func(c *fiber.Ctx) error {
		query := c.Query("q")
		apiKey := os.Getenv("GCP_API_KEY")
		cx := os.Getenv("GCP_CX")

		if apiKey == "" || cx == "" {
			return c.Status(500).JSON(fiber.Map{"error": "Search API keys not configured"})
		}

		searchURL := fmt.Sprintf("https://www.googleapis.com/customsearch/v1?q=%s&key=%s&cx=%s&searchType=image", query, apiKey, cx)
		resp, err := http.Get(searchURL)
		if err != nil {
			return c.Status(500).JSON(fiber.Map{"error": err.Error()})
		}
		defer resp.Body.Close()

		var result map[string]interface{}
		if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
			return c.Status(500).JSON(fiber.Map{"error": "Failed to parse search results"})
		}

		return c.JSON(fiber.Map{
			"status": "success",
			"data":   result["items"],
			"query":  query,
		})
	})

	api.Post("/download", func(c *fiber.Ctx) error {
		type Req struct {
			URL   string `json:"url"`
			Start string `json:"start"`
			End   string `json:"end"`
		}
		var req Req
		if err := c.BodyParser(&req); err != nil {
			return c.Status(400).JSON(fiber.Map{"error": "Invalid request body"})
		}

		// Run python worker in background
		go func() {
			args := []string{"./worker/downloader.py", req.URL}
			if req.Start != "" {
				args = append(args, "--start", req.Start)
			}
			if req.End != "" {
				args = append(args, "--end", req.End)
			}
			
			cmd := exec.Command("python3", args...)
			cmd.Env = append(os.Environ(), "PYTHONUNBUFFERED=1")

			stdout, _ := cmd.StdoutPipe()
			cmd.Start()

			scanner := bufio.NewScanner(stdout)
			for scanner.Scan() {
				line := scanner.Text()
				broadcast(line)
				log.Println("Worker:", line)
			}
			cmd.Wait()
			broadcast("PROGRESS: 100%")
		}()

		return c.JSON(fiber.Map{
			"status": "success",
			"message": "Download task queued",
		})
	})

	// WebSocket for progress
	app.Use("/ws", func(c *fiber.Ctx) error {
		if websocket.IsWebSocketUpgrade(c) {
			c.Locals("allowed", true)
			return c.Next()
		}
		return fiber.ErrUpgradeRequired
	})

	app.Get("/ws", websocket.New(func(c *websocket.Conn) {
		client := &Client{Conn: c}
		clientsMu.Lock()
		clients[client] = true
		clientsMu.Unlock()

		defer func() {
			clientsMu.Lock()
			delete(clients, client)
			clientsMu.Unlock()
			c.Close()
		}()

		for {
			if _, _, err := c.ReadMessage(); err != nil {
				break
			}
		}
	}))

	port := os.Getenv("PORT")
	if port == "" {
		port = "8000"
	}

	log.Printf("Server listening on port %s", port)
	log.Fatal(app.Listen(":" + port))
}