package main

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

	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
)

func main() {
	if err := run(); err != nil {
		slog.Error("vertex-batch stopped", "error", err)
		os.Exit(1)
	}
}

func run() error {
	httpAddr := flag.String("http", ":8080", "HTTP listen address")
	gateway := flag.String("gateway-url", "", "Gateway native conversion URL")
	prefix := flag.String("prefix", "vertex_ai_batch/", "Client model prefix")
	routePrefix := flag.String("route-prefix", "vertex_batch/", "Gateway native route prefix")
	project := flag.String("project", "", "Vertex project")
	bucket := flag.String("bucket", "", "GCS batch bucket")
	credentials := flag.String("credentials", "", "Google service account JSON file")
	vertex := flag.String("vertex-url", "https://aiplatform.googleapis.com", "Vertex API URL")
	storage := flag.String("storage-url", "https://storage.googleapis.com", "GCS API URL")
	poll := flag.Duration("poll", 10*time.Second, "Batch poll interval")
	heartbeat := flag.Duration("heartbeat", 15*time.Second, "SSE keepalive interval")
	flag.Parse()
	if *gateway == "" || *prefix == "" || *routePrefix == "" || *project == "" || *bucket == "" || *credentials == "" || *poll <= 0 || *heartbeat <= 0 {
		return errors.New("gateway-url, prefixes, project, bucket, credentials and positive intervals are required")
	}
	data, err := os.ReadFile(*credentials)
	if err != nil {
		return fmt.Errorf("read credentials file: %w", err)
	}
	config, err := google.JWTConfigFromJSON(data, "https://www.googleapis.com/auth/cloud-platform")
	if err != nil {
		return errors.New("invalid service account configuration")
	}
	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer cancel()
	client := oauth2.NewClient(context.Background(), config.TokenSource(context.Background()))
	backend := &cloud{client: client, vertex: *vertex, storage: *storage, project: *project, bucket: *bucket, poll: *poll}
	mux := http.NewServeMux()
	mux.Handle("POST /front", frontend{client: http.DefaultClient, gateway: *gateway, prefix: *prefix, routePrefix: *routePrefix, heartbeat: *heartbeat})
	mux.Handle("/", handler{runner: backend, heartbeat: *heartbeat})
	httpServer := &http.Server{Addr: *httpAddr, Handler: mux, BaseContext: func(net.Listener) context.Context { return ctx }}
	stopped := make(chan error, 1)
	go func() { stopped <- httpServer.ListenAndServe() }()
	select {
	case err = <-stopped:
	case <-ctx.Done():
	}
	cancel()
	shutdown, stop := context.WithTimeout(context.Background(), 65*time.Second)
	defer stop()
	if e := httpServer.Shutdown(shutdown); e != nil {
		return e
	}
	if errors.Is(err, http.ErrServerClosed) {
		return nil
	}
	return err
}