gooseek/backend/cmd/api-gateway/main.go

package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/gofiber/fiber/v2"
	"github.com/gofiber/fiber/v2/middleware/cors"
	"github.com/gofiber/fiber/v2/middleware/logger"
	"github.com/gooseek/backend/pkg/config"
)

var svcURLs map[string]string

func main() {
	cfg, err := config.Load()
	if err != nil {
		log.Fatal("Failed to load config:", err)
	}

	svcURLs = map[string]string{
		"chat":     cfg.ChatSvcURL,
		"agents":   cfg.AgentSvcURL,
		"search":   cfg.SearchSvcURL,
		"llm":      cfg.LLMSvcURL,
		"scraper":  cfg.ScraperSvcURL,
		"memory":   cfg.MemorySvcURL,
		"library":  cfg.LibrarySvcURL,
		"thread":   cfg.ThreadSvcURL,
		"discover": cfg.DiscoverSvcURL,
		"finance":  cfg.FinanceHeatmapURL,
		"learning": cfg.LearningSvcURL,
		"computer": cfg.ComputerSvcURL,
	}

	app := fiber.New(fiber.Config{
		StreamRequestBody: true,
		BodyLimit:         50 * 1024 * 1024,
		ReadTimeout:       time.Duration(cfg.HTTPTimeout),
		WriteTimeout:      5 * time.Minute,
		IdleTimeout:       2 * time.Minute,
	})

	app.Use(logger.New())
	app.Use(cors.New(cors.Config{
		AllowOrigins: strings.Join(cfg.AllowedOrigins, ","),
		AllowHeaders: "Origin, Content-Type, Accept, Authorization",
		AllowMethods: "GET, POST, PUT, PATCH, DELETE, OPTIONS",
	}))

	app.Get("/health", func(c *fiber.Ctx) error {
		return c.JSON(fiber.Map{"status": "ok"})
	})

	app.Get("/ready", func(c *fiber.Ctx) error {
		return c.JSON(fiber.Map{"status": "ready"})
	})

	app.Post("/api/chat", handleChat)
	app.All("/api/*", handleProxy)

	port := cfg.APIGatewayPort
	log.Printf("api-gateway listening on :%d", port)
	log.Fatal(app.Listen(fmt.Sprintf(":%d", port)))
}

func getTarget(path string) (base, rewrite string) {
	switch {
	case path == "/api/chat" || strings.HasPrefix(path, "/api/chat?"):
		return svcURLs["chat"], "/api/v1/chat"
	case strings.HasPrefix(path, "/api/v1/agents"):
		return svcURLs["agents"], path
	case strings.HasPrefix(path, "/api/v1/search"):
		return svcURLs["search"], path
	case strings.HasPrefix(path, "/api/v1/llm"), strings.HasPrefix(path, "/api/v1/providers"):
		return svcURLs["llm"], path
	case strings.HasPrefix(path, "/api/v1/memory"):
		return svcURLs["memory"], path
	case strings.HasPrefix(path, "/api/v1/library"):
		return svcURLs["library"], path
	case strings.HasPrefix(path, "/api/v1/threads"):
		return svcURLs["thread"], path
	case strings.HasPrefix(path, "/api/v1/spaces"):
		return svcURLs["thread"], path
	case strings.HasPrefix(path, "/api/v1/pages"):
		return svcURLs["thread"], path
	case strings.HasPrefix(path, "/api/v1/share"):
		return svcURLs["thread"], path
	case strings.HasPrefix(path, "/api/v1/discover"):
		return svcURLs["discover"], path
	case strings.HasPrefix(path, "/api/v1/heatmap"):
		return svcURLs["finance"], path
	case strings.HasPrefix(path, "/api/v1/movers"):
		return svcURLs["finance"], path
	case strings.HasPrefix(path, "/api/v1/markets"):
		return svcURLs["finance"], path
	case strings.HasPrefix(path, "/api/v1/learning"):
		return svcURLs["learning"], path
	case strings.HasPrefix(path, "/api/v1/computer"):
		return svcURLs["computer"], path
	default:
		return "", ""
	}
}

func handleChat(c *fiber.Ctx) error {
	base := svcURLs["chat"]
	if base == "" {
		return c.Status(503).JSON(fiber.Map{"error": "Chat service not configured"})
	}

	targetURL := strings.TrimSuffix(base, "/") + "/api/v1/chat"

	req, err := http.NewRequest("POST", targetURL, strings.NewReader(string(c.Body())))
	if err != nil {
		return c.Status(500).JSON(fiber.Map{"error": err.Error()})
	}

	req.Header.Set("Content-Type", "application/json")
	if auth := c.Get("Authorization"); auth != "" {
		req.Header.Set("Authorization", auth)
	}

	client := &http.Client{Timeout: 5 * time.Minute}
	resp, err := client.Do(req)
	if err != nil {
		return c.Status(503).JSON(fiber.Map{"error": "Service unavailable"})
	}

	if resp.StatusCode != http.StatusOK {
		defer resp.Body.Close()
		body, _ := io.ReadAll(resp.Body)
		return c.Status(resp.StatusCode).Send(body)
	}

	c.Set("Content-Type", "application/x-ndjson")
	c.Set("Cache-Control", "no-cache")
	c.Set("Transfer-Encoding", "chunked")

	c.Context().SetBodyStreamWriter(func(w *bufio.Writer) {
		defer resp.Body.Close()
		buf := make([]byte, 4096)
		for {
			n, err := resp.Body.Read(buf)
			if n > 0 {
				w.Write(buf[:n])
				w.Flush()
			}
			if err == io.EOF {
				break
			}
			if err != nil {
				break
			}
		}
	})

	return nil
}

func handleProxy(c *fiber.Ctx) error {
	path := c.Path()
	base, rewrite := getTarget(path)

	if base == "" {
		return c.Status(404).JSON(fiber.Map{"error": "Not found"})
	}

	targetURL := strings.TrimSuffix(base, "/") + rewrite
	if c.Context().QueryArgs().Len() > 0 {
		targetURL += "?" + string(c.Context().QueryArgs().QueryString())
	}

	method := c.Method()
	var body io.Reader
	if method != "GET" && method != "HEAD" {
		body = strings.NewReader(string(c.Body()))
	}

	req, err := http.NewRequest(method, targetURL, body)
	if err != nil {
		return c.Status(500).JSON(fiber.Map{"error": err.Error()})
	}

	passHeaders := []string{"Authorization", "Content-Type", "Accept", "User-Agent", "Accept-Language"}
	for _, h := range passHeaders {
		if v := c.Get(h); v != "" {
			req.Header.Set(h, v)
		}
	}

	client := &http.Client{Timeout: time.Minute}
	resp, err := client.Do(req)
	if err != nil {
		return handleFallback(c, path)
	}
	defer resp.Body.Close()

	for _, h := range []string{"Content-Type", "Cache-Control", "Set-Cookie"} {
		if v := resp.Header.Get(h); v != "" {
			c.Set(h, v)
		}
	}

	data, _ := io.ReadAll(resp.Body)
	return c.Status(resp.StatusCode).Send(data)
}

func handleFallback(c *fiber.Ctx, path string) error {
	switch {
	case strings.HasPrefix(path, "/api/v1/discover"):
		return c.JSON(fiber.Map{"items": []interface{}{}})
	case strings.HasPrefix(path, "/api/geo-context"):
		return c.JSON(fiber.Map{"country": nil, "city": nil})
	case strings.HasPrefix(path, "/api/translations"):
		return c.JSON(fiber.Map{})
	default:
		return c.Status(503).JSON(fiber.Map{"error": "Service unavailable"})
	}
}

func init() {
	if os.Getenv("PORT") == "" {
		os.Setenv("PORT", "3015")
	}
}