package main

import (
	"context"
	"encoding/json"
	"io"
	"log/slog"
	"net/http"
	"sync"
	"time"
)

type batchRunner interface {
	run(context.Context, string, json.RawMessage) (json.RawMessage, error)
}

type handler struct {
	runner    batchRunner
	heartbeat time.Duration
}

func (h handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method == http.MethodGet && r.URL.Path == "/healthz" {
		w.WriteHeader(http.StatusOK)
		return
	}
	if r.Method != http.MethodPost || r.URL.Path != "/v1/chat/completions" {
		http.NotFound(w, r)
		return
	}
	model := r.Header.Get("x-batch-model")
	streaming := r.Header.Get("x-batch-stream")
	if model == "" || (streaming != "true" && streaming != "false") {
		http.Error(w, "invalid batch model or stream headers", http.StatusBadRequest)
		return
	}
	native, err := io.ReadAll(r.Body)
	if err != nil || !json.Valid(native) {
		http.Error(w, "invalid Gemini request", http.StatusBadRequest)
		return
	}
	ctx, cancel := context.WithCancel(r.Context())
	type outcome struct {
		body json.RawMessage
		err  error
	}
	done := make(chan outcome, 1)
	var workers sync.WaitGroup
	workers.Go(func() {
		body, err := h.runner.run(ctx, model, native)
		done <- outcome{body: body, err: err}
	})
	defer workers.Wait()
	defer cancel()
	if streaming != "true" {
		select {
		case <-ctx.Done():
			return
		case result := <-done:
			if result.err != nil {
				writeError(w)
				return
			}
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write(result.body)
			return
		}
	}
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	flush := http.NewResponseController(w)
	if err := flush.Flush(); err != nil {
		return
	}
	ticker := time.NewTicker(h.heartbeat)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			if _, err := io.WriteString(w, ": batch pending\n\n"); err != nil {
				return
			}
			if err := flush.Flush(); err != nil {
				return
			}
		case result := <-done:
			if result.err != nil {
				slog.Error("batch failed", "kind", "cloud")
				return
			}
			_, _ = io.WriteString(w, "data: "+string(result.body)+"\n\n")
			return
		}
	}
}

func writeError(w http.ResponseWriter) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusBadGateway)
	_, _ = io.WriteString(w, `{"error":{"code":502,"status":"INTERNAL","message":"Batch did not produce a result"}}`)
}