main.go

package main

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/gptscript-ai/datasets/pkg/tools"
)

func main() {
	if os.Getenv("PORT") == "" {
		fmt.Println("PORT is not set")
		os.Exit(1)
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/addElements", tools.AddElements)
	mux.HandleFunc("/getAllElements", tools.GetAllElements)
	mux.HandleFunc("/listElements", tools.ListElements)
	mux.HandleFunc("/getElement", tools.GetElement)
	mux.HandleFunc("/listDatasets", tools.ListDatasets)
	mux.HandleFunc("/outputFilter", tools.OutputFilter)
	mux.HandleFunc("/", health)

	srv := &http.Server{
		Addr:    "127.0.0.1:" + os.Getenv("PORT"),
		Handler: mux,
	}

	stop := make(chan os.Signal, 1)
	signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)

	go func() {
		if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			fmt.Printf("error: %v\n", err)
		}
	}()

	<-stop
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	if err := srv.Shutdown(ctx); err != nil {
		fmt.Printf("server forced to shut down: %v\n", err)
	}
}

func health(w http.ResponseWriter, _ *http.Request) {
	w.Write([]byte("ok"))
}