opantoantro/handlers/openai.go

package handlers

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/gofiber/fiber/v3"
	"gorm.io/gorm"

	"optoant/config"
	"optoant/internal/logger"
	"optoant/internal/proxy"
	"optoant/models"
)

// OpenAIHandler handles all /v1/* requests and forwards them to OPENAI_BACKEND.
// @Summary      OpenAI-compatible proxy
// @Description  Forwards any /v1/* request to the configured OpenAI-compatible backend (e.g. DeepSeek).
// @Tags         openai
// @Accept       json
// @Produce      json
// @Param        path  path  string  true  "OpenAI API path (e.g. chat/completions)"
// @Success      200   {object}  map[string]interface{}
// @Failure      502   {object}  map[string]string
// @Router       /v1/{path} [post]
func OpenAIHandler(cfg *config.Config, db *gorm.DB) fiber.Handler {
	return func(c fiber.Ctx) error {
		start := time.Now()

		// Build target URL: OPENAI_BACKEND + full original path
		targetURL := strings.TrimRight(cfg.OpenAIBackend, "/") + c.OriginalURL()

		// Read body
		bodyBytes := c.Body()

		// Inject default model from config if request body has no model field
		if cfg.OpenAIModel != "" {
			bodyBytes = injectDefaultModel(bodyBytes, cfg.OpenAIModel)
		}

		// Convert fiber headers to net/http.Header
		reqHeaders := fiberToHTTPHeaders(c)

		// Auto-inject API key from .env if client didn't provide one
		if cfg.OpenAIApiKey != "" && reqHeaders.Get("Authorization") == "" {
			reqHeaders.Set("Authorization", "Bearer "+cfg.OpenAIApiKey)
			logger.Debug("│  🔑 Auto-injected API key from OPENAI_KEY env")
		}

		// ── LOG: Incoming OpenAI Request ──
		logger.Info("┌─ [OPENAI] >>> %s %s | IP: %s", c.Method(), c.OriginalURL(), c.IP())
		logger.Debug("│  Headers: %v", proxy.MaskSensitiveHeaders(reqHeaders))
		logger.Debug("│  Body: %s", string(bodyBytes))
		logger.Debug("│  🚀 Forwarding to: %s", targetURL)

		result, err := proxy.Forward(
			c.Context(),
			c.Method(),
			targetURL,
			reqHeaders,
			bytes.NewReader(bodyBytes),
			cfg.RequestTimeoutSeconds,
		)

		latency := time.Since(start).Milliseconds()
		status := http.StatusBadGateway
		if err == nil {
			status = result.StatusCode
		}

		// Log to DB (fire-and-forget style; don't fail the request if logging fails)
		if db != nil {
			logEntry := &models.RequestLog{
				Endpoint:       c.OriginalURL(),
				Method:         c.Method(),
				ClientIP:       c.IP(),
				RequestBody:    models.TruncateBody(string(bodyBytes)),
				ResponseStatus: status,
				LatencyMs:      latency,
			}
			go db.Create(logEntry) //nolint:errcheck
		}

		if err != nil {
			logger.Warn("│  ❌ UPSTREAM ERROR [IP: %s]: %v", c.IP(), err)
			logger.Warn("└─ [OPENAI] <<< 502 (%dms) | IP: %s", latency, c.IP())
			return c.Status(fiber.StatusBadGateway).JSON(fiber.Map{
				"error": fmt.Sprintf("upstream error: %v", err),
			})
		}

		// ── LOG: Upstream Response ──
		logger.Debug("│  <--- UPSTREAM RESPONSE (%d bytes, status %d):", len(result.Body), result.StatusCode)
		logger.Debug("│  %s", string(result.Body))
		logger.Info("└─ [OPENAI] <<< %d (%dms)", result.StatusCode, latency)

		// Copy response headers
		copyResponseHeaders(c, result.Headers)

		return c.Status(result.StatusCode).Send(result.Body)
	}
}

// fiberToHTTPHeaders converts Fiber request headers to stdlib net/http.Header.
func fiberToHTTPHeaders(c fiber.Ctx) http.Header {
	h := make(http.Header)
	c.Request().Header.VisitAll(func(k, v []byte) {
		h.Add(string(k), string(v))
	})
	return h
}

// copyResponseHeaders writes upstream response headers back to the Fiber response.
func copyResponseHeaders(c fiber.Ctx, headers http.Header) {
	skipHeaders := map[string]bool{
		"Transfer-Encoding": true,
		"Content-Length":    true,
		"Connection":        true,
	}
	for key, vals := range headers {
		if skipHeaders[key] {
			continue
		}
		for _, v := range vals {
			c.Set(key, v)
		}
	}
}

// injectDefaultModel forces the model field in a JSON body to the configured
// default model. This ensures Claude Code's model selection is overridden by
// the upstream model from .env (OPENAI_MODEL).
func injectDefaultModel(body []byte, defaultModel string) []byte {
	if defaultModel == "" {
		return body
	}
	var req map[string]interface{}
	if err := json.Unmarshal(body, &req); err != nil {
		return body
	}
	if req["model"] == defaultModel {
		return body
	}
	oldModel := req["model"]
	req["model"] = defaultModel
	modified, err := json.Marshal(req)
	if err != nil {
		return body
	}
	logger.Info("│  🔧 Model override: %v → %s", oldModel, defaultModel)
	return modified
}