package main

import (
	"encoding/base64"
	"flag"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"

	"github.com/coreos/go-systemd/v22/activation"
)

var (
	addr       = flag.String("addr", "127.0.0.1:11434", "address to listen on")
	useSystemd = flag.Bool("systemd-activation", false, "use systemd activation")
	token      = flag.String("token", "", "bearer token for authentication")
	target     = flag.String("target", "https://ollama.ula.inventati.org", "target hostname")
)

func newSystemdListener() (net.Listener, error) {
	listeners, err := activation.Listeners()
	if err != nil {
		return nil, err
	}
	return listeners[0], nil
}

func newListener() (net.Listener, error) {
	if *useSystemd {
		return newSystemdListener()
	}
	return net.Listen("tcp", *addr)
}

func main() {
	flag.Parse()

	authHdr := "Bearer " + base64.StdEncoding.EncodeToString([]byte(*token))
	targetURI, _ := url.Parse(*target)

	proxy := &httputil.ReverseProxy{
		Rewrite: func(req *httputil.ProxyRequest) {
			req.SetURL(targetURI)
			req.Out.Header.Set("Authorization", authHdr)
		},
	}

	l, err := newListener()
	if err != nil {
		log.Fatal(err)
	}

	if err := http.Serve(l, proxy); err != nil {
		log.Fatal(err)
	}
}