opencatd-open/backend/internal/proxy/gateway.go

package proxy

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"opencatd-open/internal/channel"
	"opencatd-open/internal/dao"
	"opencatd-open/internal/proxy/convert"
	"opencatd-open/internal/store"
	"opencatd-open/pkg/config"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"gorm.io/gorm"
)

type Gateway struct {
	ctx        context.Context
	cfg        *config.Config
	db         *gorm.DB
	wg         *sync.WaitGroup
	httpClient *http.Client

	userDAO    *dao.UserDAO
	apiKeyDAO  *dao.ApiKeyDAO
	usageDAO   *dao.UsageDAO
	dailyDAO   *dao.DailyUsageDAO
	channelSvc *channel.Service
}

func NewGateway(ctx context.Context, cfg *config.Config, db *gorm.DB, wg *sync.WaitGroup, userDAO *dao.UserDAO, apiKeyDAO *dao.ApiKeyDAO, usageDAO *dao.UsageDAO, dailyDAO *dao.DailyUsageDAO) *Gateway {
	client := &http.Client{Timeout: 120 * time.Second}
	if os.Getenv("LOCAL_PROXY") != "" {
		proxyUrl, err := url.Parse(os.Getenv("LOCAL_PROXY"))
		if err == nil {
			tr := &http.Transport{
				Proxy: http.ProxyURL(proxyUrl),
			}
			client.Transport = tr
		}
	}

	return &Gateway{
		ctx:        ctx,
		cfg:        cfg,
		db:         db,
		wg:         wg,
		httpClient: client,
		userDAO:    userDAO,
		apiKeyDAO:  apiKeyDAO,
		usageDAO:   usageDAO,
		dailyDAO:   dailyDAO,
		channelSvc: nil,
	}
}

func (g *Gateway) SetChannelService(svc *channel.Service) {
	g.channelSvc = svc
}

// Request represents a parsed incoming request
type Request struct {
	Model     string
	Stream    bool
	Protocol  string // "chat", "messages", "responses"
	Body      []byte
	APIKey    *store.APIKey
	UserID    uint64
}

// ParseRequest parses the incoming request and extracts key fields
func (g *Gateway) ParseRequest(c *gin.Context, protocol string) (*Request, error) {
	body, err := io.ReadAll(c.Request.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read body: %w", err)
	}

	apiKey, _ := c.Get("api_key")
	userID, _ := c.Get("user_id")

	req := &Request{
		Protocol: protocol,
		Body:     body,
		UserID:   userID.(uint64),
	}

	if ak, ok := apiKey.(*store.APIKey); ok {
		req.APIKey = ak
	}

	// Parse model and stream based on protocol
	switch protocol {
	case "chat":
		var parsed convert.ChatCompletionRequest
		if err := json.Unmarshal(body, &parsed); err != nil {
			return nil, fmt.Errorf("invalid chat request: %w", err)
		}
		req.Model = parsed.Model
		req.Stream = parsed.Stream
	case "messages":
		var parsed convert.MessagesRequest
		if err := json.Unmarshal(body, &parsed); err != nil {
			return nil, fmt.Errorf("invalid messages request: %w", err)
		}
		req.Model = parsed.Model
		req.Stream = parsed.Stream
	case "responses":
		var parsed convert.ResponsesRequest
		if err := json.Unmarshal(body, &parsed); err != nil {
			return nil, fmt.Errorf("invalid responses request: %w", err)
		}
		req.Model = parsed.Model
		req.Stream = parsed.Stream
	}

	return req, nil
}

// Dispatch routes the request to the appropriate upstream
func (g *Gateway) Dispatch(c *gin.Context, req *Request) {
	if g.channelSvc == nil {
		g.writeError(c, http.StatusBadGateway, "channel service not available")
		return
	}

	ch, err := g.channelSvc.SelectChannel(g.ctx, req.Model)
	if err != nil {
		g.writeError(c, http.StatusBadGateway, err.Error())
		return
	}

	apiKey, err := g.channelSvc.GetAPIKey(ch)
	if err != nil {
		g.writeError(c, http.StatusBadGateway, "failed to decrypt API key")
		return
	}

	// Determine target format and convert if needed
	targetFormat := req.Protocol
	if len(ch.FormatsEffective()) > 0 {
		// Prefer the channel's native format
		for _, f := range ch.FormatsEffective() {
			if f == req.Protocol {
				targetFormat = f
				break
			}
		}
	}

	// Build upstream URL
	upstreamPath := g.getUpstreamPath(req.Protocol)
	upstreamURL := ch.UpstreamURL(req.Protocol, upstreamPath)

	// Convert request if needed
	var requestBody []byte
	if targetFormat != req.Protocol {
		requestBody, err = g.convertRequest(req.Body, req.Protocol, targetFormat)
		if err != nil {
			g.writeError(c, http.StatusBadRequest, "conversion failed: "+err.Error())
			return
		}
	} else {
		requestBody = req.Body
	}

	// Create upstream request
	httpReq, err := http.NewRequestWithContext(g.ctx, "POST", upstreamURL, bytes.NewReader(requestBody))
	if err != nil {
		g.writeError(c, http.StatusBadGateway, "failed to create request")
		return
	}

	// Set headers
	g.setHeaders(httpReq, ch, apiKey, targetFormat)

	// Execute request
	start := time.Now()
	resp, err := g.httpClient.Do(httpReq)
	latency := time.Since(start)
	if err != nil {
		g.channelSvc.RecordFailure(ch.ID)
		g.writeError(c, http.StatusBadGateway, fmt.Sprintf("upstream error: %v (latency: %v)", err, latency))
		return
	}
	defer resp.Body.Close()

	// Record success
	g.channelSvc.RecordSuccess(ch.ID)

	// Handle response
	if resp.StatusCode >= 400 {
		body, _ := io.ReadAll(resp.Body)
		log.Printf("Upstream error: status=%d body=%s", resp.StatusCode, string(body))
		c.Data(resp.StatusCode, "application/json", body)
		return
	}

	// Stream or buffer response
	if req.Stream {
		g.streamResponse(c, resp, req.Protocol, ch)
	} else {
		g.bufferResponse(c, resp, req.Protocol, ch)
	}
}

func (g *Gateway) getUpstreamPath(protocol string) string {
	switch protocol {
	case "chat":
		return "/chat/completions"
	case "messages":
		return "/messages"
	case "responses":
		return "/responses"
	default:
		return "/chat/completions"
	}
}

func (g *Gateway) setHeaders(req *http.Request, ch *store.Channel, apiKey string, format string) {
	req.Header.Set("Content-Type", "application/json")

	switch ch.Provider {
	case store.ChannelProviderOpenAI, store.ChannelProviderCompatible:
		req.Header.Set("Authorization", "Bearer "+apiKey)
	case store.ChannelProviderAnthropic:
		req.Header.Set("x-api-key", apiKey)
		req.Header.Set("anthropic-version", "2023-06-01")
	}
}

func (g *Gateway) convertRequest(body []byte, from, to string) ([]byte, error) {
	switch {
	case from == "chat" && to == "messages":
		var req convert.ChatCompletionRequest
		if err := json.Unmarshal(body, &req); err != nil {
			return nil, err
		}
		msgReq, err := convert.ChatToMessages(&req)
		if err != nil {
			return nil, err
		}
		return json.Marshal(msgReq)

	case from == "chat" && to == "responses":
		var req convert.ChatCompletionRequest
		if err := json.Unmarshal(body, &req); err != nil {
			return nil, err
		}
		respReq, err := convert.ChatToResponses(&req)
		if err != nil {
			return nil, err
		}
		return json.Marshal(respReq)

	case from == "messages" && to == "chat":
		var req convert.MessagesRequest
		if err := json.Unmarshal(body, &req); err != nil {
			return nil, err
		}
		// Messages -> Chat: we need to construct a ChatCompletionRequest
		chatReq := &convert.ChatCompletionRequest{
			Model: req.Model,
		}
		for _, m := range req.Messages {
			chatReq.Messages = append(chatReq.Messages, m)
		}
		if req.Temperature != nil {
			chatReq.Temperature = req.Temperature
		}
		if req.TopP != nil {
			chatReq.TopP = req.TopP
		}
		chatReq.Tools = req.Tools
		chatReq.Stream = req.Stream
		return json.Marshal(chatReq)

	case from == "responses" && to == "chat":
		var req convert.ResponsesRequest
		if err := json.Unmarshal(body, &req); err != nil {
			return nil, err
		}
		chatReq := &convert.ChatCompletionRequest{
			Model: req.Model,
		}
		for _, item := range req.Input {
			chatReq.Messages = append(chatReq.Messages, convert.Message{
				Role:    item.Role,
				Content: item.Content,
			})
		}
		chatReq.Tools = req.Tools
		chatReq.Stream = req.Stream
		return json.Marshal(chatReq)

	default:
		return body, nil
	}
}

func (g *Gateway) streamResponse(c *gin.Context, resp *http.Response, protocol string, ch *store.Channel) {
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Status(http.StatusOK)

	writer := convert.NewSSEWriter(c.Writer)
	parser := convert.NewSSEParser(resp.Body)

	for {
		event, err := parser.ReadEvent()
		if err != nil {
			if err == io.EOF {
				break
			}
			log.Printf("Stream parse error: %v", err)
			break
		}

		if event.Event == "error" {
			log.Printf("Upstream stream error: %s", event.Data)
			break
		}

		// Write raw SSE event based on protocol
		if err := writer.WriteEvent("chat CompletionChunk", event.Data); err != nil {
			break
		}
	}

	writer.WriteDone()
}

func (g *Gateway) bufferResponse(c *gin.Context, resp *http.Response, protocol string, ch *store.Channel) {
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		g.writeError(c, http.StatusBadGateway, "failed to read response")
		return
	}

	c.Data(resp.StatusCode, "application/json", body)
}

func (g *Gateway) writeError(c *gin.Context, status int, message string) {
	protocol := c.GetHeader("X-Protocol")
	if protocol == "" {
		protocol = "chat"
	}

	switch {
	case strings.Contains(c.GetHeader("Accept"), "text/event-stream"):
		c.Header("Content-Type", "text/event-stream")
		c.Status(status)
		fmt.Fprintf(c.Writer, "data: {\"error\":{\"message\":\"%s\"}}\n\n", message)
		fmt.Fprintf(c.Writer, "data: [DONE]\n\n")
	case protocol == "messages":
		c.JSON(status, gin.H{
			"type": "error",
			"error": gin.H{
				"type":    "api_error",
				"message": message,
			},
		})
	default:
		c.JSON(status, gin.H{
			"error": gin.H{
				"message": message,
				"type":    "invalid_request_error",
			},
		})
	}
}