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

package proxy

import (
	"bufio"
	"bytes"
	"context"
	"crypto/rand"
	"encoding/hex"
	"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/internal/usage"
	"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
	modelDAO   *dao.ModelDAO
	channelSvc *channel.Service
	usageRec   *usage.Recorder
}

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,
		modelDAO:   dao.NewModelDAO(db),
		channelSvc: nil,
	}
}

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

// SetUsageRecorder 注入异步用量记录器；nil 时网关跳过用量上报。
func (g *Gateway) SetUsageRecorder(r *usage.Recorder) {
	g.usageRec = r
}

// generateRequestID 生成请求级唯一 ID，用于用量明细关联与排障。
func generateRequestID() string {
	b := make([]byte, 12)
	if _, err := rand.Read(b); err != nil {
		return fmt.Sprintf("req-%d", time.Now().UnixNano())
	}
	return "req-" + hex.EncodeToString(b)
}

// 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
	KeyID     uint64
	RequestID string
}

// 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),
		RequestID: c.GetHeader("X-Request-Id"),
	}
	if req.RequestID == "" {
		req.RequestID = generateRequestID()
	}

	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
	}

	cands := g.channelSvc.Candidates(req.Model)
	// 内存健康过滤：连续失败进入 cooldown 的渠道不再尝试（渠道级健康自愈靠冷却过期）。
	cands = g.channelSvc.FilterHealthy(cands)
	if len(cands) == 0 {
		g.writeError(c, http.StatusServiceUnavailable, "no enabled channels for model: "+req.Model)
		g.recordUsage(req, nil, nil, usage.Event{
			IsError: true, ErrorCode: "no_channel",
		}, convert.TokenUsage{})
		return
	}

	var lastCh *store.Channel
	_ = lastCh // 保留变量名便于断点排查；失败渠道已在循环内各自 RecordFailure
	lastErrStatus := http.StatusBadGateway
	lastErrBody := "all upstream channels failed"

	for i := range cands {
		cand := &cands[i]
		ch := cand.Channel
		lastCh = ch

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

		// Determine target format: channel declares support for the client protocol
		// then passthrough, otherwise convert to its first supported protocol
		// (chat > messages > responses).
		targetFormat := g.conversionTarget(ch, req.Protocol)
		if targetFormat == "" {
			continue // 渠道不支持该协议，换下一个
		}

		// Build upstream URL
		upstreamURL := ch.UpstreamURL(targetFormat, g.getUpstreamPath(targetFormat))

		// Convert request if needed
		var requestBody []byte
		if targetFormat != req.Protocol {
			var err error
			requestBody, err = convert.ConvertRequest(req.Body, req.Protocol, targetFormat)
			if err != nil {
				lastErrStatus, lastErrBody = http.StatusBadRequest, "conversion failed: "+err.Error()
				continue
			}
		} else {
			requestBody = req.Body
		}

		// 绑定了 upstream_model 时把请求体里的 model 重写为上游模型名（别名映射）。
		if cand.Binding != nil && cand.Binding.UpstreamModel != "" &&
			cand.Binding.UpstreamModel != req.Model {
			requestBody = rewriteModel(requestBody, cand.Binding.UpstreamModel)
		}

		// Create upstream request
		httpReq, err := http.NewRequestWithContext(g.ctx, "POST", upstreamURL, bytes.NewReader(requestBody))
		if err != nil {
			lastErrStatus, lastErrBody = http.StatusBadGateway, "failed to create request"
			continue
		}
		g.setHeaders(httpReq, ch, apiKey, targetFormat)

		// Execute request
		start := time.Now()
		resp, err := g.httpClient.Do(httpReq)
		if err != nil {
			g.channelSvc.RecordFailure(ch.ID)
			lastErrStatus = http.StatusBadGateway
			lastErrBody = fmt.Sprintf("upstream error: %v", err)
			g.recordUsage(req, cand, ch, usage.Event{
				IsError:   true,
				ErrorCode: "upstream_error",
				LatencyMS: int(time.Since(start).Milliseconds()),
			}, convert.TokenUsage{})
			continue // 可重试：换下一个渠道
		}

		// Handle upstream error responses
		if resp.StatusCode >= 400 {
			body, _ := io.ReadAll(resp.Body)
			resp.Body.Close()
			log.Printf("Upstream error: status=%d body=%s", resp.StatusCode, string(body))
			g.recordUsage(req, cand, ch, usage.Event{
				IsError:   true,
				ErrorCode: fmt.Sprintf("upstream_%d", resp.StatusCode),
				LatencyMS: int(time.Since(start).Milliseconds()),
			}, convert.TokenUsage{})
			// 429/5xx 可换渠道重试；4xx 直接透传
			if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500 {
				lastErrStatus, lastErrBody = resp.StatusCode, string(body)
				continue
			}
			c.Data(resp.StatusCode, "application/json", body)
			return
		}

		g.channelSvc.RecordSuccess(ch.ID)

		// Stream or buffer response；tok 从上游响应（SSE usage 块或非流式 JSON）提取。
		var tok convert.TokenUsage
		if req.Stream {
			tok = g.streamResponse(c, resp, req.Protocol, targetFormat)
		} else {
			tok = g.bufferResponse(c, resp, req.Protocol, targetFormat)
		}
		resp.Body.Close()
		// 成功记录：用量 + 定价计费。
		g.recordUsage(req, cand, ch, usage.Event{
			LatencyMS: int(time.Since(start).Milliseconds()),
		}, tok)
		return
	}

	// 全部候选失败（每个候选失败时已各自 RecordFailure，不再重复计数）
	g.writeError(c, lastErrStatus, lastErrBody)
}

// rewriteModel 把 JSON 请求体顶层的 model 字段替换为 upstreamModel。
func rewriteModel(body []byte, upstreamModel string) []byte {
	var m map[string]json.RawMessage
	if json.Unmarshal(body, &m) != nil {
		return body
	}
	if _, ok := m["model"]; !ok {
		return body
	}
	m["model"], _ = json.Marshal(upstreamModel)
	out, err := json.Marshal(m)
	if err != nil {
		return body
	}
	return out
}

// recordUsage 汇总一次请求的用量事件并异步落库。tok 为从上游响应提取的用量。
// cand/ch 可为 nil（无可用渠道的失败场景）。
func (g *Gateway) recordUsage(req *Request, cand *channel.Candidate, ch *store.Channel, ev usage.Event, tok convert.TokenUsage) {
	if g.usageRec == nil {
		return
	}
	ev.UserID = req.UserID
	ev.ModelName = req.Model
	ev.Protocol = req.Protocol
	ev.RequestID = req.RequestID
	if req.APIKey != nil {
		ev.KeyID = req.APIKey.ID
	}
	if ch != nil {
		ev.ChannelID = ch.ID
	}
	if cand != nil && cand.Binding != nil {
		ev.ModelID = cand.Binding.ModelID
	}
	ev.PromptTokens = tok.InputTokens
	ev.CompletionTokens = tok.OutputTokens
	ev.CacheReadTokens = tok.CacheReadTokens
	ev.CacheCreationTokens = tok.CacheCreationTokens
	// 定价与成本（价格按每百万 token 的 USD 单价）。
	// 成本口径：非缓存输入 × 输入价 + 缓存读 × 缓存价 + 缓存写与输出 × 输出价。
	if ev.ModelID != 0 {
		if m, err := g.modelDAO.GetByID(ev.ModelID); err == nil {
			ev.InputPrice = m.InputPrice
			ev.OutputPrice = m.OutputPrice
			ev.CacheReadPrice = m.CacheReadPrice
		}
	}
	if !ev.IsError {
		ev.Cost = (float64(ev.PromptTokens-ev.CacheReadTokens)*ev.InputPrice +
			float64(ev.CacheReadTokens)*ev.CacheReadPrice +
			float64(ev.CacheCreationTokens)*ev.OutputPrice +
			float64(ev.CompletionTokens)*ev.OutputPrice) / 1e6
	}
	g.usageRec.Record(ev)
}

// conversionTarget 决定客户端协议在渠道上的处理方式：
// 渠道声明支持该协议则直通；否则转为其首选支持协议（chat > messages > responses）。
func (g *Gateway) conversionTarget(ch *store.Channel, clientProto string) string {
	formats := ch.FormatsEffective()
	for _, f := range formats {
		if f == clientProto {
			return clientProto
		}
	}
	for _, p := range []string{convert.ProtoChat, convert.ProtoMessages, convert.ProtoResponses} {
		for _, f := range formats {
			if f == p {
				return p
			}
		}
	}
	return ""
}

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")
	}
}


// streamResponse 流式响应：按 \n\n 分块零缓冲转发；跨协议时逐行转换。
// 返回从上游 SSE usage 块累计的 token 用量（按上游协议解析）。
func (g *Gateway) streamResponse(c *gin.Context, resp *http.Response, clientProto, upstreamProto string) convert.TokenUsage {
	w := c.Writer
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Status(http.StatusOK)

	flusher, _ := w.(http.Flusher)

	// 跨协议时按行转换；同协议直通（lineConv 为 nil）。
	var lineConv func([]byte) []byte
	if upstreamProto != clientProto {
		lineConv = convert.NewStreamTransformer(upstreamProto, clientProto)
	}

	// 上游原始行按 \n\n 分块，避免把 data 行内的转义换行当成事件边界。
	// 同时喂入用量累计器（usage 块可能出现在任一事件）。
	r := bufio.NewReaderSize(resp.Body, 32*1024)
	accum := convert.NewStreamUsageAccum()
	for {
		buf := []byte{}
		for {
			line, err := r.ReadSlice('\n')
			if err == bufio.ErrBufferFull {
				buf = append(buf, line...)
				continue
			}
			buf = append(buf, line...)
			if err == io.EOF {
				if len(buf) == 0 {
					return accum.Usage()
				}
				if !bytes.HasSuffix(buf, []byte("\n")) {
					buf = append(buf, '\n')
				}
			} else if err != nil {
				log.Printf("stream read error: %v", err)
				return accum.Usage()
			}
			if len(buf) >= 2 && bytes.HasSuffix(buf, []byte("\n\n")) {
				break
			}
		}

		// 先解析用量（data: {...} 行），再决定转发内容。
		for _, data := range sseDataPayloads(buf) {
			accum.Feed(data, upstreamProto)
		}

		out := buf
		if lineConv != nil {
			out = lineConv(buf)
		}
		if len(out) == 0 {
			continue
		}
		if _, err := w.Write(out); err != nil {
			return accum.Usage() // 客户端已断开
		}
		if flusher != nil {
			flusher.Flush()
		}

		// 流结束标记：chat/messages 上游以 data: [DONE] 收尾。部分上游（keep-alive）
		// 发完 [DONE] 后不关连接，继续读会阻塞到超时；据此主动收尾。
		// responses 协议没有 [DONE]，以 response.completed 事件收尾。
		if streamTerminated(buf, upstreamProto) {
			return accum.Usage()
		}
	}
}

// streamTerminated 判断一块 SSE 是否为上游流的结束事件。
func streamTerminated(chunk []byte, proto string) bool {
	switch proto {
	case convert.ProtoChat:
		// chat 上游以 data: [DONE] 收尾；keep-alive 上游发完不关连接。
		return bytes.Contains(chunk, []byte("data: [DONE]"))
	case convert.ProtoMessages:
		// messages 上游以 message_stop 事件结束（无 [DONE]）。
		return bytes.Contains(chunk, []byte(`"type":"message_stop"`)) ||
			bytes.Contains(chunk, []byte(`"type": "message_stop"`)) ||
			bytes.Contains(chunk, []byte("data: [DONE]"))
	case convert.ProtoResponses:
		return bytes.Contains(chunk, []byte(`"response.completed"`)) ||
			bytes.Contains(chunk, []byte(`"type":"response.completed"`))
	}
	return false
}

// sseDataPayloads 从一块 SSE（一个完整事件，\n\n 结尾）中取出所有 data 行的原始载荷。
func sseDataPayloads(chunk []byte) [][]byte {
	var out [][]byte
	for _, line := range bytes.Split(chunk, []byte("\n")) {
		line = bytes.TrimSuffix(line, []byte("\r"))
		if !bytes.HasPrefix(line, []byte("data:")) {
			continue
		}
		payload := bytes.TrimPrefix(line, []byte("data:"))
		payload = bytes.TrimPrefix(payload, []byte(" "))
		if len(payload) == 0 || bytes.Equal(payload, []byte("[DONE]")) {
			continue
		}
		out = append(out, payload)
	}
	return out
}

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

	// 用量从上游原始响应体提取（先于转换，转换会改字段名）。
	tok, _ := convert.ExtractUsageJSON(body, upstreamProto)

	out := body
	if upstreamProto != clientProto {
		if converted, cerr := convert.ConvertResponse(body, upstreamProto, clientProto); cerr == nil {
			out = converted
		} else {
			// 转换失败时至少剥掉非 JSON 前缀，让客户端能解析出正文
			out = convert.CleanJSON(body)
		}
	} else {
		// 直通：部分上游（如 OpenRouter）的 non-stream 响应在 JSON 前夹带空白/注释
		out = convert.CleanJSON(body)
	}
	c.Data(resp.StatusCode, "application/json", out)
	return tok
}

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",
			},
		})
	}
}