package api

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"

	"monica-go-completion-backend/internal/completion"
)

var wsUpgrader = websocket.Upgrader{
	ReadBufferSize:  4096,
	WriteBufferSize: 4096,
	CheckOrigin: func(_ *http.Request) bool {
		return true
	},
}

type wsCompletionRequest struct {
	ID        string `json:"id"`
	URI       string `json:"uri"`
	Text      string `json:"text"`
	Line      int    `json:"line"`
	Character int    `json:"character"`
}

type wsCompletionResponse struct {
	ID           string            `json:"id"`
	Items        []completion.Item `json:"items,omitempty"`
	IsIncomplete bool              `json:"isIncomplete,omitempty"`
	Error        string            `json:"error,omitempty"`
}

func registerWSRoutes(router *gin.Engine, service CompletionService) {
	router.GET("/ws/completions/go", func(c *gin.Context) {
		conn, err := wsUpgrader.Upgrade(c.Writer, c.Request, nil)
		if err != nil {
			return
		}
		defer conn.Close()

		var writeMu sync.Mutex

		for {
			_, payload, err := conn.ReadMessage()
			if err != nil {
				break
			}

			var req wsCompletionRequest
			if err := json.Unmarshal(payload, &req); err != nil {
				sendWSResponse(conn, &writeMu, wsCompletionResponse{
					ID:    "",
					Error: "invalid JSON payload",
				})
				continue
			}

			go func(r wsCompletionRequest) {
				ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
				defer cancel()

				resp, err := service.Complete(ctx, completion.Request{
					URI:       r.URI,
					Text:      r.Text,
					Line:      r.Line,
					Character: r.Character,
				})
				if err != nil {
					msg := "completion failed"
					if errors.Is(err, completion.ErrInvalidRequest) {
						msg = err.Error()
					}
					sendWSResponse(conn, &writeMu, wsCompletionResponse{
						ID:    r.ID,
						Error: msg,
					})
					return
				}

				sendWSResponse(conn, &writeMu, wsCompletionResponse{
					ID:           r.ID,
					Items:        resp.Items,
					IsIncomplete: resp.IsIncomplete,
				})
			}(req)
		}
	})
}

func sendWSResponse(conn *websocket.Conn, writeMu *sync.Mutex, resp wsCompletionResponse) {
	writeMu.Lock()
	defer writeMu.Unlock()
	_ = conn.WriteJSON(resp)
}