package main

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

type frontend struct {
	client      *http.Client
	gateway     string
	prefix      string
	routePrefix string
	heartbeat   time.Duration
}

func (f frontend) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	data, err := io.ReadAll(r.Body)
	var body map[string]json.RawMessage
	if err != nil || json.Unmarshal(data, &body) != nil {
		http.Error(w, "invalid JSON", http.StatusBadRequest)
		return
	}
	var model string
	var streaming bool
	if json.Unmarshal(body["model"], &model) != nil || !strings.HasPrefix(model, f.prefix) || strings.TrimPrefix(model, f.prefix) == "" {
		http.Error(w, "invalid batch model", http.StatusBadRequest)
		return
	}
	if raw, ok := body["stream"]; ok {
		if json.Unmarshal(raw, &streaming) != nil {
			http.Error(w, "invalid stream flag", http.StatusBadRequest)
			return
		}
	}
	body["model"], err = json.Marshal(f.routePrefix + strings.TrimPrefix(model, f.prefix))
	if err != nil {
		writeError(w)
		return
	}
	data, err = json.Marshal(body)
	if err != nil {
		writeError(w)
		return
	}
	ctx, cancel := context.WithCancel(r.Context())
	defer cancel()
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, f.gateway, bytes.NewReader(data))
	if err != nil {
		writeError(w)
		return
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", r.Header.Get("Authorization"))
	if streaming {
		f.stream(w, req)
		return
	}
	res, err := f.client.Do(req)
	if err != nil {
		writeError(w)
		return
	}
	defer func() { _ = res.Body.Close() }()
	w.Header().Set("Content-Type", res.Header.Get("Content-Type"))
	w.WriteHeader(res.StatusCode)
	_, _ = io.Copy(w, res.Body)
}

type fragment struct {
	data []byte
	end  bool
}

func (f frontend) receive(req *http.Request, incoming chan<- fragment) {
	send := func(part fragment) bool {
		select {
		case incoming <- part:
			return true
		case <-req.Context().Done():
			return false
		}
	}
	res, err := f.client.Do(req)
	if err != nil {
		send(fragment{end: true})
		return
	}
	defer func() { _ = res.Body.Close() }()
	if res.StatusCode < http.StatusOK || res.StatusCode >= http.StatusMultipleChoices {
		send(fragment{end: true})
		return
	}
	buffer := make([]byte, 4096)
	for {
		n, err := res.Body.Read(buffer)
		if !send(fragment{data: bytes.Clone(buffer[:n]), end: err != nil}) || err != nil {
			return
		}
	}
}

func (f frontend) stream(w http.ResponseWriter, req *http.Request) {
	incoming := make(chan fragment)
	go f.receive(req, incoming)
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	flush := http.NewResponseController(w)
	if flush.Flush() != nil {
		return
	}
	ticker := time.NewTicker(f.heartbeat)
	defer ticker.Stop()
	var parser events
	for {
		var output []byte
		select {
		case <-req.Context().Done():
			return
		case <-ticker.C:
			output = []byte(": batch pending\n\n")
		case part := <-incoming:
			output = parser.feed(part.data, part.end)
		}
		if _, err := w.Write(output); err != nil {
			return
		}
		if flush.Flush() != nil || parser.done {
			return
		}
	}
}