proxy.go

// Package internal provides proxy service logic for github-copilot-svcs.
package internal

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"
)

var (
	copilotAPIBase   = "https://api.githubcopilot.com"
	copilotAPIBaseMu sync.RWMutex
)
var chatCompletionsPath = "/chat/completions"
var responsesPath = "/responses"

// SetCopilotAPIBase overrides the upstream API base URL for testing.
func SetCopilotAPIBase(base string) {
	copilotAPIBaseMu.Lock()
	defer copilotAPIBaseMu.Unlock()
	copilotAPIBase = base
}

// ResetCopilotAPIBase restores the default upstream API base URL.
func ResetCopilotAPIBase() {
	copilotAPIBaseMu.Lock()
	defer copilotAPIBaseMu.Unlock()
	copilotAPIBase = "https://api.githubcopilot.com"
}

// GetCopilotAPIBase returns the current upstream API base URL.
func GetCopilotAPIBase() string {
	copilotAPIBaseMu.RLock()
	defer copilotAPIBaseMu.RUnlock()
	return copilotAPIBase
}

const (
	maxChatRetries     = 3
	baseChatRetryDelay = 1 // seconds

	// Circuit breaker configuration
	circuitBreakerFailureThreshold = 5

	// Request configuration
	maxRequestBodySize  = 5 * 1024 * 1024 // 5MB
	streamingBufferSize = 1024

	// Status code ranges
	statusCodeServerError     = 500
	statusCodeTooManyRequests = 429
	statusCodeRequestTimeout  = 408
)

const (
	// ProxyCBStateClosed indicates the circuit breaker is closed.
	ProxyCBStateClosed = 0
	// ProxyCBStateOpen indicates the circuit breaker is open.
	ProxyCBStateOpen = 1
	// ProxyCBStateHalfOpen indicates the circuit breaker is half-open.
	ProxyCBStateHalfOpen = 2
)

// CircuitBreakerState represents the state of the circuit breaker
type CircuitBreakerState int

const (
	// CircuitClosed allows all requests through
	CircuitClosed CircuitBreakerState = iota
	// CircuitOpen rejects all requests
	CircuitOpen
	// CircuitHalfOpen allows limited requests through
	CircuitHalfOpen
)

// CircuitBreaker implements circuit breaker pattern for upstream API calls
type CircuitBreaker struct {
	failureCount    int64
	lastFailureTime time.Time
	state           CircuitBreakerState
	timeout         time.Duration
	mutex           sync.RWMutex
}

// CoalescingCache handles request coalescing for identical requests
type CoalescingCache struct {
	requests map[string]chan interface{}
	mutex    sync.RWMutex
}

// ProxyService provides proxy functionality
type ProxyService struct {
	config         *Config
	httpClient     *http.Client
	authService    *AuthService
	workerPool     WorkerPoolInterface
	circuitBreaker *CircuitBreaker
	bufferPool     *sync.Pool
}

// WorkerPoolInterface interface for background processing
type WorkerPoolInterface interface {
	Submit(job func())
}

// responseWrapper tracks if headers have been sent
type responseWrapper struct {
	http.ResponseWriter
	headersSent bool
}

// NewCoalescingCache creates a new coalescing cache
func NewCoalescingCache() *CoalescingCache {
	return &CoalescingCache{
		requests: make(map[string]chan interface{}),
	}
}

// GetRequestKey generates a cache key for request coalescing
func (cc *CoalescingCache) GetRequestKey(method, url string, body interface{}) string {
	h := sha256.New()
	h.Write([]byte(method))
	h.Write([]byte(url))
	if body != nil {
		if bodyBytes, ok := body.([]byte); ok {
			h.Write(bodyBytes)
		}
	}
	return hex.EncodeToString(h.Sum(nil))
}

// CoalesceRequest executes a function only once for identical concurrent requests
func (cc *CoalescingCache) CoalesceRequest(key string, fn func() interface{}) interface{} {
	cc.mutex.Lock()

	// Check if request is already in progress
	if ch, exists := cc.requests[key]; exists {
		cc.mutex.Unlock()
		// Wait for the existing request to complete
		return <-ch
	}

	// Create new channel for this request
	ch := make(chan interface{}, 1)
	cc.requests[key] = ch
	cc.mutex.Unlock()

	// Execute the request
	result := fn()

	// Broadcast result to all waiting goroutines
	ch <- result
	close(ch)

	// Clean up
	cc.mutex.Lock()
	delete(cc.requests, key)
	cc.mutex.Unlock()

	return result
}

// NewProxyService creates a new proxy service
func NewProxyService(cfg *Config, httpClient *http.Client, authService *AuthService, workerPool WorkerPoolInterface) *ProxyService {
	circuitBreaker := &CircuitBreaker{
		state:   CircuitClosed,
		timeout: time.Duration(cfg.Timeouts.CircuitBreaker) * time.Second,
	}

	bufferPool := &sync.Pool{
		New: func() interface{} {
			return new(bytes.Buffer)
		},
	}

	return &ProxyService{
		config:         cfg,
		httpClient:     httpClient,
		authService:    authService,
		workerPool:     workerPool,
		circuitBreaker: circuitBreaker,
		bufferPool:     bufferPool,
	}
}

// Handler returns an HTTP handler for the proxy endpoint
func (s *ProxyService) Handler() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		// Create context with extended timeout for long-lived streaming responses
		ctx, cancel := context.WithTimeout(r.Context(), time.Duration(s.config.Timeouts.ProxyContext)*time.Second)
		defer cancel()

		// Check circuit breaker
		if !s.circuitBreaker.canExecute() {
			Warn("Circuit breaker is open, rejecting request")
			http.Error(w, "Service temporarily unavailable", http.StatusServiceUnavailable)
			return
		}

		// Limit request body size
		r.Body = http.MaxBytesReader(w, r.Body, maxRequestBodySize)

		// Use a response wrapper to track if headers have been sent
		respWrapper := &responseWrapper{ResponseWriter: w, headersSent: false}

		// Create a done channel to track completion
		done := make(chan error, 1)

		// Submit request to worker pool
		s.workerPool.Submit(func() {
			defer func() {
				if recovery := recover(); recovery != nil {
					Error("Worker panic recovered", "panic", recovery)
					done <- NewProxyError("request_processing", "worker panic during request processing", fmt.Errorf("panic: %v", recovery))
				}
			}()

			err := s.processProxyRequest(ctx, respWrapper, r)
			done <- err
		})

		// Wait for worker to complete or context timeout
		select {
		case err := <-done:
			if err != nil {
				Error("Worker error", "error", err)
				// Only write error if headers haven't been sent
				if !respWrapper.headersSent {
					switch {
					case strings.Contains(err.Error(), "authentication error"):
						http.Error(w, err.Error(), http.StatusUnauthorized)
					case strings.Contains(err.Error(), "token validation failed"):
						http.Error(w, err.Error(), http.StatusUnauthorized)
					case strings.Contains(err.Error(), "bad request"):
						http.Error(w, err.Error(), http.StatusBadRequest)
					case strings.Contains(err.Error(), "method not allowed"):
						http.Error(w, err.Error(), http.StatusMethodNotAllowed)
					default:
						http.Error(w, err.Error(), http.StatusInternalServerError)
					}
				}
			}
		case <-ctx.Done():
			Warn("Request timeout in worker pool")
			// Only write timeout error if headers haven't been sent
			if !respWrapper.headersSent {
				http.Error(w, "Request timeout", http.StatusRequestTimeout)
			}
		}
	}
}

func (rw *responseWrapper) WriteHeader(statusCode int) {
	if !rw.headersSent {
		rw.headersSent = true
		rw.ResponseWriter.WriteHeader(statusCode)
	}
}

func (rw *responseWrapper) Write(data []byte) (int, error) {
	if !rw.headersSent {
		rw.headersSent = true
	}
	return rw.ResponseWriter.Write(data)
}

func (cb *CircuitBreaker) canExecute() bool {
	cb.mutex.RLock()
	defer cb.mutex.RUnlock()

	// No metrics to update for circuit breaker state changes

	if cb.state == CircuitClosed {
		return true
	}

	if cb.state == CircuitOpen {
		if time.Since(cb.lastFailureTime) > cb.timeout {
			cb.mutex.RUnlock()
			cb.mutex.Lock()
			cb.state = CircuitHalfOpen
			cb.mutex.Unlock()
			cb.mutex.RLock()
			return true
		}
		return false
	}

	// CircuitHalfOpen
	return true
}

func (cb *CircuitBreaker) onSuccess() {
	cb.mutex.Lock()
	defer cb.mutex.Unlock()

	cb.failureCount = 0
	cb.state = CircuitClosed
}

func (cb *CircuitBreaker) onFailure() {
	cb.mutex.Lock()
	defer cb.mutex.Unlock()

	cb.failureCount++
	cb.lastFailureTime = time.Now()

	if cb.failureCount >= circuitBreakerFailureThreshold {
		cb.state = CircuitOpen
	}
}

func (s *ProxyService) processProxyRequest(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
	Debug("Starting proxy request", "method", r.Method, "path", r.URL.Path)

	// Validate method
	if r.Method != http.MethodPost {
		return fmt.Errorf("method not allowed: %s", r.Method)
	}

	// Read the request body
	body, err := io.ReadAll(r.Body)
	if err != nil {
		Error("Error reading request body", "error", err)
		// Check for "http: request body too large" error and return 413
		if strings.Contains(err.Error(), "http: request body too large") {
			return fmt.Errorf("payload too large: %w", err)
		}
		return fmt.Errorf("bad request: failed to read request body: %w", err)
	}
	defer func() {
		if err := r.Body.Close(); err != nil {
			Warn("Error closing request body", "error", err)
		}
	}()

	// Basic body validation (for demonstration: consider empty body an error)
	if len(body) == 0 {
		return fmt.Errorf("bad request: empty request body")
	}

	var input struct {
		Model string `json:"model"`
	}
	if jsonErr := json.Unmarshal(body, &input); jsonErr != nil {
		return fmt.Errorf("bad request: invalid JSON: %w", jsonErr)
	}

	// AllowedModels validation
	if len(s.config.AllowedModels) > 0 {
		allowed := false
		for _, m := range s.config.AllowedModels {
			if input.Model == m {
				allowed = true
				break
			}
		}
		if !allowed {
			return fmt.Errorf("bad request: model '%s' is not allowed by allowed_models in config", input.Model)
		}
	}

	// Ensure we have a valid token before making the request
	if tokenErr := s.authService.EnsureValidToken(s.config); tokenErr != nil {
		Error("Failed to ensure valid token", "error", tokenErr)
		return NewAuthError("token validation failed", tokenErr)
	}

	// Create new request to GitHub Copilot
	var targetURL string
	base := GetCopilotAPIBase()
	switch r.URL.Path {
	case "/v1/chat/completions":
		targetURL = base + chatCompletionsPath
	case "/v1/responses":
		targetURL = base + responsesPath
	default:
		return fmt.Errorf("unsupported proxy path: %s", r.URL.Path)
	}
	Debug("Sending request to target", "url", targetURL, "body_length", len(body))

	req, err := http.NewRequestWithContext(ctx, r.Method, targetURL, bytes.NewBuffer(body))
	if err != nil {
		Error("Error creating request", "error", err)
		return NewProxyError("create_request", "failed to create proxy request", err)
	}

	// Set headers
	// Forward content/negotiation headers from client if present; use defaults if missing
	headersToProxy := []string{"Content-Type", "Accept", "Accept-Encoding", "TE"}
	defaults := map[string]string{
		"Content-Type": "application/json",
		"Accept":       "application/json",
	}
	for _, h := range headersToProxy {
		if v := r.Header.Get(h); v != "" {
			req.Header.Set(h, v)
		} else if def, ok := defaults[h]; ok {
			req.Header.Set(h, def)
		}
	}

	req.Header.Set("Authorization", "Bearer "+s.config.CopilotToken)
	req.Header.Set("User-Agent", s.config.Headers.UserAgent)
	req.Header.Set("Editor-Version", s.config.Headers.EditorVersion)
	req.Header.Set("Editor-Plugin-Version", s.config.Headers.EditorPluginVersion)
	req.Header.Set("Copilot-Integration-Id", s.config.Headers.CopilotIntegrationID)
	req.Header.Set("Openai-Intent", s.config.Headers.OpenaiIntent)
	req.Header.Set("X-Initiator", s.config.Headers.XInitiator)

	resp, err := s.makeRequestWithRetry(req, body)
	if err != nil {
		s.circuitBreaker.onFailure()
		Error("Error making request after retries", "error", err)
		return NewNetworkError("proxy_request", targetURL, "failed to complete request after retries", err)
	}
	defer func() {
		if err := resp.Body.Close(); err != nil {
			Warn("Error closing response body", "error", err)
		}
	}()

	// Update circuit breaker based on response
	if resp.StatusCode < statusCodeServerError {
		s.circuitBreaker.onSuccess()
	} else {
		s.circuitBreaker.onFailure()
	}

	Debug("Received response", "status", resp.StatusCode, "content_type", resp.Header.Get("Content-Type"))

	// Copy response headers
	for key, values := range resp.Header {
		for _, value := range values {
			w.Header().Add(key, value)
		}
	}

	// Add configurable CORS headers
	if len(s.config.CORS.AllowedOrigins) > 0 {
		w.Header().Set("Access-Control-Allow-Origin", strings.Join(s.config.CORS.AllowedOrigins, ", "))
	}
	if len(s.config.CORS.AllowedHeaders) > 0 {
		w.Header().Set("Access-Control-Allow-Headers", strings.Join(s.config.CORS.AllowedHeaders, ", "))
	}

	// Copy status code
	w.WriteHeader(resp.StatusCode)

	// Handle streaming vs regular responses
	if resp.Header.Get("Content-Type") == "text/event-stream" {
		return s.handleStreamingResponse(w, resp)
	}
	return s.handleRegularResponse(w, resp)
}

func (s *ProxyService) handleStreamingResponse(w http.ResponseWriter, resp *http.Response) error {
	Debug("Starting streaming response copy")

	if flusher, ok := w.(http.Flusher); ok {
		// Copy in chunks and flush periodically for better streaming
		buf := make([]byte, streamingBufferSize)
		for {
			n, readErr := resp.Body.Read(buf)
			if n > 0 {
				_, writeErr := w.Write(buf[:n])
				if writeErr != nil {
					Error("Error writing streaming chunk", "error", writeErr)
					return writeErr
				}
				flusher.Flush()
			}
			if readErr == io.EOF {
				Debug("Streaming response completed successfully")
				break
			}
			if readErr != nil {
				Error("Error reading streaming response", "error", readErr)
				return readErr
			}
		}
	} else {
		// Fallback to direct copy if no flusher available
		_, err := io.Copy(w, resp.Body)
		if err != nil {
			Error("Error copying streaming response", "error", err)
			return err
		}
	}
	return nil
}

func (s *ProxyService) handleRegularResponse(w http.ResponseWriter, resp *http.Response) error {
	Debug("Starting regular response copy")

	// Use buffer pool for regular responses
	buf := s.bufferPool.Get().(*bytes.Buffer)
	buf.Reset()
	defer s.bufferPool.Put(buf)

	_, err := io.CopyBuffer(w, resp.Body, buf.Bytes()[:0])
	if err != nil {
		Error("Error copying response", "error", err)
		return err
	}

	Debug("Regular response completed successfully")
	return nil
}

func (s *ProxyService) makeRequestWithRetry(req *http.Request, body []byte) (*http.Response, error) {
	var lastResp *http.Response
	var lastErr error

	for attempt := 1; attempt <= maxChatRetries; attempt++ {
		// Create a new request for each attempt with the original context
		retryReq, err := http.NewRequestWithContext(req.Context(), req.Method, req.URL.String(), bytes.NewBuffer(body))
		if err != nil {
			return nil, err
		}

		// Copy all headers
		for key, values := range req.Header {
			for _, value := range values {
				retryReq.Header.Add(key, value)
			}
		}

		Debug("Making request attempt", "attempt", attempt, "max_attempts", maxChatRetries)

		resp, err := s.httpClient.Do(retryReq)
		if err != nil {
			lastErr = err
			if attempt == maxChatRetries {
				Error("Request failed after max attempts", "attempts", maxChatRetries, "error", err)
				return nil, err
			}

			// Context-aware waiting instead of blocking sleep
			waitTime := time.Duration(baseChatRetryDelay*attempt*attempt) * time.Second
			Warn("Request failed, retrying", "attempt", attempt, "wait_time", waitTime, "error", err)

			timer := time.NewTimer(waitTime)
			select {
			case <-timer.C:
				// Continue with retry
			case <-req.Context().Done():
				timer.Stop()
				return nil, req.Context().Err()
			}
			continue
		}

		lastResp = resp

		// Check if we should retry based on status code
		if !s.isRetriableError(resp.StatusCode, nil) {
			Debug("Request successful", "attempt", attempt, "status", resp.StatusCode)
			return resp, nil
		}

		// Close the response body before retrying
		if closeErr := resp.Body.Close(); closeErr != nil {
			Warn("Failed to close response body during retry", "error", closeErr)
		}

		if attempt == maxChatRetries {
			Warn("Request failed after max attempts", "attempts", maxChatRetries, "status", resp.StatusCode)
			return resp, nil // Return the last response even if it failed
		}

		// Context-aware waiting for status code retries
		waitTime := time.Duration(baseChatRetryDelay*attempt*attempt) * time.Second
		Warn("Request failed, retrying", "status", resp.StatusCode, "attempt", attempt, "wait_time", waitTime)

		timer := time.NewTimer(waitTime)
		select {
		case <-timer.C:
			// Continue with retry
		case <-req.Context().Done():
			timer.Stop()
			return nil, req.Context().Err()
		}
	}

	return lastResp, lastErr
}

func (s *ProxyService) isRetriableError(statusCode int, err error) bool {
	if err != nil {
		return true // Network errors are retriable
	}

	// Retry on server errors and rate limiting
	return statusCode >= statusCodeServerError || statusCode == statusCodeTooManyRequests || statusCode == statusCodeRequestTimeout
}