main.go

// teploy-gateway — the OSS AI gateway: one project key, every model,
// tracked, on hardware you own. Provider keys never leave this daemon.
package main

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/useteploy/teploy-gateway/internal/gateway"
)

// version is set by goreleaser at build time.
var version = "dev"

func main() {
	if len(os.Args) < 2 {
		usage()
		os.Exit(2)
	}
	var err error
	switch os.Args[1] {
	case "serve":
		err = serve(os.Args[2:])
	case "key":
		err = keyCommand(os.Args[2:])
	case "usage":
		err = usageCommand(os.Args[2:])
	case "version":
		fmt.Println(version)
	default:
		usage()
		os.Exit(2)
	}
	if err != nil {
		fmt.Fprintln(os.Stderr, "teploy-gateway:", err)
		os.Exit(1)
	}
}

func usage() {
	fmt.Fprintln(os.Stderr, `teploy-gateway — one project key, every model, tracked

Usage:
  teploy-gateway serve [--addr 127.0.0.1:7440] [--config config.json] [--data-dir /deployments/gateway]
  teploy-gateway key create --project <name>
  teploy-gateway key list
  teploy-gateway key rotate <key-id>
  teploy-gateway key revoke <key-id>
  teploy-gateway usage [--project <name>]
  teploy-gateway version

Provider keys come from the config file or environment
(ANTHROPIC_API_KEY, OPENAI_API_KEY, GROQ_API_KEY, ...). Point the AI
SDK at the gateway: createAnthropic({ baseURL, apiKey: projectKey }).`)
}

func serve(args []string) error {
	flags := flag.NewFlagSet("serve", flag.ExitOnError)
	addr := flags.String("addr", "127.0.0.1:7440", "listen address")
	configPath := flags.String("config", "", "providers/fallbacks JSON (optional; env keys work without it)")
	dataDir := flags.String("data-dir", "/deployments/gateway", "keys, usage, and admin token directory")
	usageSinkFlag := flags.String("usage-sink", "", "usage sink: file (default) or nucleus (requires DATABASE_URL)")
	if err := flags.Parse(args); err != nil {
		return err
	}

	log := slog.New(slog.NewJSONHandler(os.Stdout, nil))
	config, err := gateway.LoadConfig(*configPath)
	if err != nil {
		return err
	}
	keys, err := gateway.OpenKeystore(*dataDir)
	if err != nil {
		return err
	}

	// Sink selection: flag wins over config; Nucleus requires DATABASE_URL,
	// otherwise fall back to the file sink (the OSS default).
	sinkChoice := *usageSinkFlag
	if sinkChoice == "" {
		sinkChoice = config.UsageSink
	}
	dbURL := os.Getenv("DATABASE_URL")
	var usageSink gateway.UsageSink
	if sinkChoice == "nucleus" && dbURL != "" {
		usageSink, err = gateway.OpenNucleusUsage(context.Background(), dbURL, 5*time.Second, log)
		if err != nil {
			return err
		}
		log.Info("usage sink: nucleus")
	} else {
		if sinkChoice == "nucleus" {
			log.Warn("usage sink nucleus requested but DATABASE_URL is unset — using file sink")
		}
		usageSink, err = gateway.OpenFileUsage(*dataDir, 5*time.Second)
		if err != nil {
			return err
		}
	}
	adminToken, err := loadOrMintAdminToken(filepath.Join(*dataDir, "admin-token"))
	if err != nil {
		return err
	}

	server := &gateway.Server{
		Config:     config,
		Keys:       keys,
		Usage:      usageSink,
		Proxy:      &gateway.Proxy{Config: config, Usage: usageSink, Log: log},
		AdminToken: adminToken,
		Version:    version,
		Log:        log,
	}

	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()
	return gateway.Serve(ctx, *addr, server)
}

func loadOrMintAdminToken(path string) (string, error) {
	if data, err := os.ReadFile(path); err == nil && len(strings.TrimSpace(string(data))) >= 32 {
		return strings.TrimSpace(string(data)), nil
	}
	var raw [24]byte
	if _, err := rand.Read(raw[:]); err != nil {
		return "", err
	}
	token := "tgwa_" + hex.EncodeToString(raw[:])
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return "", err
	}
	if err := os.WriteFile(path, []byte(token+"\n"), 0o600); err != nil {
		return "", err
	}
	return token, nil
}

// ---------------------------------------------------------------------------
// Admin CLI verbs — thin calls to the running daemon's admin API.
// ---------------------------------------------------------------------------

type adminClient struct {
	base  string
	token string
}

func newAdminClient(addr, dataDir string) (*adminClient, error) {
	data, err := os.ReadFile(filepath.Join(dataDir, "admin-token"))
	if err != nil {
		return nil, fmt.Errorf("admin token (is the daemon running with this --data-dir?): %w", err)
	}
	return &adminClient{base: "http://" + addr, token: strings.TrimSpace(string(data))}, nil
}

func (c *adminClient) do(method, path string, body string) ([]byte, error) {
	request, err := http.NewRequest(method, c.base+path, strings.NewReader(body))
	if err != nil {
		return nil, err
	}
	request.Header.Set("Authorization", "Bearer "+c.token)
	if body != "" {
		request.Header.Set("Content-Type", "application/json")
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()
	payload, _ := io.ReadAll(response.Body)
	if response.StatusCode >= 400 {
		return nil, errors.New(strings.TrimSpace(string(payload)))
	}
	return payload, nil
}

func adminFlags(flags *flag.FlagSet) (addr, dataDir *string) {
	addr = flags.String("addr", "127.0.0.1:7440", "daemon address")
	dataDir = flags.String("data-dir", "/deployments/gateway", "daemon data directory")
	return addr, dataDir
}

func keyCommand(args []string) error {
	if len(args) < 1 {
		return errors.New("usage: key create|list|rotate|revoke")
	}
	verb, rest := args[0], args[1:]
	flags := flag.NewFlagSet("key "+verb, flag.ExitOnError)
	addr, dataDir := adminFlags(flags)
	project := flags.String("project", "", "project name (create)")

	switch verb {
	case "create":
		if err := flags.Parse(rest); err != nil {
			return err
		}
		if *project == "" {
			return errors.New("--project is required")
		}
		client, err := newAdminClient(*addr, *dataDir)
		if err != nil {
			return err
		}
		payload, err := client.do("POST", "/admin/v1/keys", fmt.Sprintf(`{"project":%q}`, *project))
		if err != nil {
			return err
		}
		var minted struct {
			ID    string `json:"id"`
			Token string `json:"token"`
		}
		if err := json.Unmarshal(payload, &minted); err != nil {
			return err
		}
		fmt.Printf("id:    %s\ntoken: %s\n\nShown once — store it now. Point the AI SDK at the gateway with this as apiKey.\n", minted.ID, minted.Token)
		return nil
	case "list":
		if err := flags.Parse(rest); err != nil {
			return err
		}
		client, err := newAdminClient(*addr, *dataDir)
		if err != nil {
			return err
		}
		payload, err := client.do("GET", "/admin/v1/keys", "")
		if err != nil {
			return err
		}
		fmt.Println(string(payload))
		return nil
	case "rotate":
		if len(rest) < 1 || strings.HasPrefix(rest[0], "-") {
			return errors.New("usage: key rotate <key-id>")
		}
		id := rest[0]
		if err := flags.Parse(rest[1:]); err != nil {
			return err
		}
		client, err := newAdminClient(*addr, *dataDir)
		if err != nil {
			return err
		}
		payload, err := client.do("POST", "/admin/v1/keys/"+id+"/rotate", "")
		if err != nil {
			return err
		}
		var minted struct {
			ID    string `json:"id"`
			Token string `json:"token"`
		}
		if err := json.Unmarshal(payload, &minted); err != nil {
			return err
		}
		fmt.Printf("id:    %s\ntoken: %s\n\nRotated from %s (now revoked). Shown once — store it now and update the project's apiKey.\n", minted.ID, minted.Token, id)
		return nil
	case "revoke":
		if len(rest) < 1 || strings.HasPrefix(rest[0], "-") {
			return errors.New("usage: key revoke <key-id>")
		}
		id := rest[0]
		if err := flags.Parse(rest[1:]); err != nil {
			return err
		}
		client, err := newAdminClient(*addr, *dataDir)
		if err != nil {
			return err
		}
		if _, err := client.do("DELETE", "/admin/v1/keys/"+id, ""); err != nil {
			return err
		}
		fmt.Println("revoked", id)
		return nil
	default:
		return fmt.Errorf("unknown key verb %q", verb)
	}
}

func usageCommand(args []string) error {
	flags := flag.NewFlagSet("usage", flag.ExitOnError)
	addr, dataDir := adminFlags(flags)
	project := flags.String("project", "", "filter to one project")
	if err := flags.Parse(args); err != nil {
		return err
	}
	client, err := newAdminClient(*addr, *dataDir)
	if err != nil {
		return err
	}
	path := "/admin/v1/usage"
	if *project != "" {
		// url.Values encodes the project name: a literal &, #, + or ? in a
		// name must survive the round trip instead of being read as query
		// structure by the server.
		q := url.Values{}
		q.Set("project", *project)
		path += "?" + q.Encode()
	}
	payload, err := client.do("GET", path, "")
	if err != nil {
		return err
	}
	fmt.Println(string(payload))
	return nil
}