gochat/backend/internal/llm/openai_provider.go

package llm

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

	applogger "github.com/gochat/gochat/pkg/logger"
)

// OpenAIProvider implements the Provider interface using OpenAI-compatible APIs.
// Supports custom base_url for domestic providers like Volcengine/Doubao.
type OpenAIProvider struct {
	apiKey     string
	baseURL    string
	model      string
	embedModel string
	httpClient *http.Client
	maxRetries int
}

// OpenAIProviderConfig holds configuration for creating an OpenAIProvider.
type OpenAIProviderConfig struct {
	APIKey        string
	BaseURL       string // defaults to "https://api.openai.com/v1"
	Model         string // defaults to "gpt-4"
	EmbedModel    string // defaults to "text-embedding-3-small"
	MaxRetries    int    // defaults to 3
	MaxRetriesSet bool   // preserves an explicit zero-retry setting
	Timeout       int    // HTTP timeout in seconds, defaults to 60
}

// NewOpenAIProvider creates a new OpenAIProvider with the given configuration.
func NewOpenAIProvider(cfg OpenAIProviderConfig) *OpenAIProvider {
	if cfg.BaseURL == "" {
		cfg.BaseURL = "https://api.openai.com/v1"
	}
	// Ensure baseURL ends without trailing slash
	cfg.BaseURL = strings.TrimRight(cfg.BaseURL, "/")

	if cfg.Model == "" {
		cfg.Model = "gpt-4"
	}
	if cfg.EmbedModel == "" {
		cfg.EmbedModel = "text-embedding-3-small"
	}
	if !cfg.MaxRetriesSet && cfg.MaxRetries == 0 {
		cfg.MaxRetries = 3
	}
	if cfg.Timeout == 0 {
		cfg.Timeout = 60
	}

	return &OpenAIProvider{
		apiKey:     cfg.APIKey,
		baseURL:    cfg.BaseURL,
		model:      cfg.Model,
		embedModel: cfg.EmbedModel,
		httpClient: &http.Client{
			Timeout: time.Duration(cfg.Timeout) * time.Second,
		},
		maxRetries: cfg.MaxRetries,
	}
}

// ChatCompletion sends a chat completion request to the OpenAI-compatible API.
func (p *OpenAIProvider) ChatCompletion(ctx context.Context, req ChatRequest) (*ChatResponse, error) {
	// Set default model if not specified
	if req.Model == "" {
		req.Model = p.model
	}

	body, err := json.Marshal(req)
	if err != nil {
		applogger.L().Errorf("ChatCompletion: failed to marshal request: %v", err)
		return nil, fmt.Errorf("marshal chat request: %w", err)
	}

	respBody, err := p.doRequestWithRetry(ctx, "/chat/completions", body)
	if err != nil {
		return nil, fmt.Errorf("chat completion request: %w", err)
	}

	var resp ChatResponse
	if err := json.Unmarshal(respBody, &resp); err != nil {
		applogger.L().Errorf("ChatCompletion: failed to unmarshal response: %v", err)
		return nil, fmt.Errorf("unmarshal chat response: %w", err)
	}

	return &resp, nil
}

// ChatCompletionStream sends a streaming chat completion request and returns chunks via callback.
func (p *OpenAIProvider) ChatCompletionStream(ctx context.Context, req ChatRequest, onChunk func(StreamChunk) error) error {
	if req.Model == "" {
		req.Model = p.model
	}
	req.Stream = true

	body, err := json.Marshal(req)
	if err != nil {
		applogger.L().Errorf("ChatCompletionStream: failed to marshal request: %v", err)
		return fmt.Errorf("marshal chat request: %w", err)
	}

	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		return fmt.Errorf("create stream request: %w", err)
	}
	p.setHeaders(httpReq)

	httpResp, err := p.httpClient.Do(httpReq)
	if err != nil {
		return fmt.Errorf("stream request: %w", err)
	}
	defer httpResp.Body.Close()

	if httpResp.StatusCode != http.StatusOK {
		respBody, _ := io.ReadAll(httpResp.Body)
		apiErr := parseAPIError(httpResp.StatusCode, respBody)
		applogger.L().Errorf("ChatCompletionStream: provider returned status %d", httpResp.StatusCode)
		return apiErr
	}

	return p.parseSSEStream(httpResp.Body, onChunk)
}

// CreateEmbedding sends an embedding request to the OpenAI-compatible API.
func (p *OpenAIProvider) CreateEmbedding(ctx context.Context, req EmbeddingRequest) (*EmbeddingResponse, error) {
	if req.Model == "" {
		req.Model = p.embedModel
	}

	body, err := json.Marshal(req)
	if err != nil {
		applogger.L().Errorf("CreateEmbedding: failed to marshal request: %v", err)
		return nil, fmt.Errorf("marshal embedding request: %w", err)
	}

	respBody, err := p.doRequestWithRetry(ctx, "/embeddings", body)
	if err != nil {
		return nil, fmt.Errorf("embedding request: %w", err)
	}

	var resp EmbeddingResponse
	if err := json.Unmarshal(respBody, &resp); err != nil {
		applogger.L().Errorf("CreateEmbedding: failed to unmarshal response: %v", err)
		return nil, fmt.Errorf("unmarshal embedding response: %w", err)
	}

	return &resp, nil
}

// doRequestWithRetry performs an HTTP request with retry logic.
func (p *OpenAIProvider) doRequestWithRetry(ctx context.Context, path string, body []byte) ([]byte, error) {
	var lastErr error

	for attempt := 0; attempt <= p.maxRetries; attempt++ {
		if attempt > 0 {
			// Exponential backoff: 1s, 2s, 4s
			backoff := time.Duration(1<<uint(attempt-1)) * time.Second
			applogger.L().Infof("Retrying provider request (attempt %d/%d) after %v", attempt, p.maxRetries, backoff)

			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case <-time.After(backoff):
			}
		}

		respBody, err := p.doRequest(ctx, path, body)
		if err == nil {
			return respBody, nil
		}

		// Don't retry on client errors (4xx except 429)
		if isNonRetriableError(err) {
			return nil, err
		}

		lastErr = err
	}

	return nil, fmt.Errorf("max retries (%d) exceeded: %w", p.maxRetries, lastErr)
}

// doRequest performs a single HTTP request to the API.
func (p *OpenAIProvider) doRequest(ctx context.Context, path string, body []byte) ([]byte, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL+path, bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("create request: %w", err)
	}
	p.setHeaders(req)

	resp, err := p.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("execute request: %w", err)
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("read response body: %w", err)
	}

	if resp.StatusCode >= 400 {
		apiErr := parseAPIError(resp.StatusCode, respBody)
		applogger.L().Errorf("Provider API error (status %d)", resp.StatusCode)
		return nil, apiErr
	}

	return respBody, nil
}

// setHeaders sets common headers for API requests.
func (p *OpenAIProvider) setHeaders(req *http.Request) {
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+p.apiKey)
}

// parseSSEStream parses Server-Sent Events from a streaming response body.
func (p *OpenAIProvider) parseSSEStream(body io.Reader, onChunk func(StreamChunk) error) error {
	reader := newSSEReader(body)

	for {
		event, err := reader.Next()
		if err != nil {
			return fmt.Errorf("read SSE event: %w", err)
		}

		if event == nil {
			// Stream ended
			return nil
		}

		// Skip non-data events
		if event.Type != "message" || event.Data == "" {
			continue
		}

		// OpenAI sends "[DONE]" to signal stream end
		if event.Data == "[DONE]" {
			return nil
		}

		var chunk StreamChunk
		if err := json.Unmarshal([]byte(event.Data), &chunk); err != nil {
			applogger.L().Errorf("parseSSEStream: failed to unmarshal provider chunk: %v", err)
			continue
		}

		if err := onChunk(chunk); err != nil {
			return fmt.Errorf("chunk callback: %w", err)
		}
	}
}

// --- Error types ---

// APIError represents an error returned by the OpenAI-compatible API.
type APIError struct {
	StatusCode int
	Message    string
	Type       string
	Code       string
}

func (e *APIError) Error() string {
	return fmt.Sprintf("API error (status %d): %s", e.StatusCode, e.Message)
}

// parseAPIError creates an APIError from HTTP status and response body.
func parseAPIError(statusCode int, body []byte) *APIError {
	apiErr := &APIError{
		StatusCode: statusCode,
		Message:    providerErrorMessage(statusCode),
	}

	// Try to parse OpenAI error structure
	var errResp struct {
		Error struct {
			Message string `json:"message"`
			Type    string `json:"type"`
			Code    string `json:"code"`
		} `json:"error"`
	}
	if err := json.Unmarshal(body, &errResp); err == nil {
		apiErr.Type = errResp.Error.Type
		apiErr.Code = errResp.Error.Code
	}

	return apiErr
}

func providerErrorMessage(statusCode int) string {
	switch statusCode {
	case http.StatusUnauthorized, http.StatusForbidden:
		return "provider authentication failed"
	case http.StatusNotFound:
		return "provider endpoint or model was not found"
	case http.StatusTooManyRequests:
		return "provider rate limit exceeded"
	default:
		return "provider request failed"
	}
}

// isNonRetriableError returns true for errors that should not be retried.
func isNonRetriableError(err error) bool {
	if apiErr, ok := err.(*APIError); ok {
		// Retry on rate limit (429) and server errors (5xx)
		// Don't retry on client errors (400, 401, 403, 404, etc.)
		return apiErr.StatusCode >= 400 && apiErr.StatusCode < 500 && apiErr.StatusCode != 429
	}
	return false
}

// --- SSE Reader ---

// sseEvent represents a parsed SSE event.
type sseEvent struct {
	Type string // event type (default "message" if not specified)
	Data string // data payload
	ID   string // event ID
}

// sseReader reads Server-Sent Events from a stream.
type sseReader struct {
	scanner *sseLineScanner
}

func newSSEReader(body io.Reader) *sseReader {
	return &sseReader{
		scanner: newSSELineScanner(body),
	}
}

// Next reads the next SSE event from the stream.
// Returns nil when the stream is complete.
func (r *sseReader) Next() (*sseEvent, error) {
	var event *sseEvent

	for {
		line, err := r.scanner.Next()
		if err != nil {
			return nil, err
		}
		if line == nil {
			// Stream ended
			return event, nil
		}

		text := *line

		if text == "" {
			// Empty line = event boundary, dispatch current event
			if event != nil {
				return event, nil
			}
			continue
		}

		if strings.HasPrefix(text, ":") {
			// Comment, skip
			continue
		}

		field, value := parseSSEField(text)

		switch field {
		case "event":
			if event == nil {
				event = &sseEvent{}
			}
			event.Type = value
		case "data":
			if event == nil {
				event = &sseEvent{Type: "message"}
			}
			if event.Data != "" {
				event.Data += "\n"
			}
			event.Data += value
		case "id":
			if event == nil {
				event = &sseEvent{}
			}
			event.ID = value
		}
	}
}

func parseSSEField(line string) (field, value string) {
	idx := strings.Index(line, ":")
	if idx == -1 {
		return line, ""
	}
	field = line[:idx]
	value = strings.TrimLeft(line[idx+1:], " ")
	return field, value
}

// sseLineScanner reads lines from an SSE stream efficiently.
type sseLineScanner struct {
	reader  io.Reader
	buffer  []byte
	hasData bool
}

func newSSELineScanner(reader io.Reader) *sseLineScanner {
	return &sseLineScanner{
		reader: reader,
		buffer: make([]byte, 0, 4096),
	}
}

// Next returns the next line from the stream.
// Returns nil when the stream is complete.
func (s *sseLineScanner) Next() (*string, error) {
	for {
		// Check if we have a complete line in the buffer
		idx := bytes.IndexByte(s.buffer, '\n')
		if idx != -1 {
			line := string(s.buffer[:idx])
			s.buffer = s.buffer[idx+1:]
			// Strip \r if present (CRLF)
			line = strings.TrimRight(line, "\r")
			return &line, nil
		}

		// Read more data
		tmp := make([]byte, 4096)
		n, err := s.reader.Read(tmp)
		if n > 0 {
			s.buffer = append(s.buffer, tmp[:n]...)
			s.hasData = true
		}
		if err != nil {
			if err == io.EOF {
				if len(s.buffer) > 0 {
					line := string(s.buffer)
					s.buffer = s.buffer[:0]
					line = strings.TrimRight(line, "\r")
					return &line, nil
				}
				return nil, nil
			}
			return nil, err
		}
	}
}

// --- Utility functions ---

// ParseFloatEmbedding converts a slice of any (JSON numbers) to []float64.
// Useful when embedding responses contain mixed numeric types.
func ParseFloatEmbedding(raw []interface{}) []float64 {
	result := make([]float64, len(raw))
	for i, v := range raw {
		switch n := v.(type) {
		case float64:
			result[i] = n
		case float32:
			result[i] = float64(n)
		case int:
			result[i] = float64(n)
		case int64:
			result[i] = float64(n)
		case string:
			f, err := strconv.ParseFloat(n, 64)
			if err == nil {
				result[i] = f
			}
		}
	}
	return result
}