package proxy

import (
	"bufio"
	"io"
	"net/http"

	"github.com/gin-gonic/gin"
	"github.com/rose_cat707/luminary/internal/gateway"
	"github.com/rose_cat707/luminary/internal/middleware"
	"github.com/rose_cat707/luminary/internal/model"
	"github.com/rose_cat707/luminary/internal/provider"
)

type Handler struct {
	gateway *gateway.Service
}

func New(gw *gateway.Service) *Handler {
	return &Handler{gateway: gw}
}

func (h *Handler) Register(r *gin.RouterGroup) {
	r.GET("/models", h.ListModels)
	r.POST("/chat/completions", h.forward(provider.EndpointChatCompletion))
	r.POST("/completions", h.forward(provider.EndpointTextCompletion))
	r.POST("/responses", h.forward(provider.EndpointResponses))
	r.POST("/embeddings", h.forward(provider.EndpointEmbeddings))
	r.POST("/rerank", h.forward(provider.EndpointRerank))
}

func (h *Handler) ingressKey(c *gin.Context) (*model.IngressKey, bool) {
	v, ok := c.Get("ingress_key")
	if !ok {
		return nil, false
	}
	k, ok := v.(*model.IngressKey)
	return k, ok
}

func (h *Handler) ListModels(c *gin.Context) {
	ingress, ok := h.ingressKey(c)
	if !ok {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "unauthorized"})
		return
	}
	models, err := h.gateway.ListModels(c.Request.Context(), ingress)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"object": "list", "data": models})
}

func (h *Handler) forward(endpoint provider.EndpointType) gin.HandlerFunc {
	return func(c *gin.Context) {
		ingress, ok := h.ingressKey(c)
		if !ok {
			c.JSON(http.StatusUnauthorized, gin.H{"error": "unauthorized"})
			return
		}
		body, err := io.ReadAll(c.Request.Body)
		if err != nil {
			c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"})
			return
		}
		if provider.IsStreamRequest(body) && provider.SupportsStreaming(endpoint) {
			h.handleStream(c, ingress, endpoint, body)
			return
		}
		clientIP := middleware.ClientIP(c)
		resp, err := h.gateway.ForwardEndpoint(c.Request.Context(), ingress, endpoint, body, clientIP)
		if err != nil {
			c.JSON(gatewayHTTPStatus(err), gin.H{"error": gin.H{"message": err.Error(), "type": "gateway_error"}})
			return
		}
		c.Data(resp.StatusCode, "application/json", resp.Body)
	}
}

func (h *Handler) handleStream(c *gin.Context, ingress *model.IngressKey, endpoint provider.EndpointType, body []byte) {
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Status(http.StatusOK)

	flusher, ok := c.Writer.(http.Flusher)
	if !ok {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "streaming not supported"})
		return
	}
	w := bufio.NewWriter(c.Writer)
	clientIP := middleware.ClientIP(c)
	err := h.gateway.ForwardStream(c.Request.Context(), ingress, endpoint, body, clientIP, w, flusher)
	w.Flush()
	if err != nil {
		_, _ = c.Writer.Write([]byte("data: {\"error\":\"" + err.Error() + "\"}\n\n"))
		flusher.Flush()
	}
}