package main

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"flag"
	"git.gammaspectra.live/S.O.N.G/FinalCommander/content"
	"git.gammaspectra.live/S.O.N.G/FinalCommander/utilities"
	"git.gammaspectra.live/S.O.N.G/MakyuuIchaival"
	"git.gammaspectra.live/S.O.N.G/MakyuuIchaival/httputils"
	"git.gammaspectra.live/S.O.N.G/MakyuuIchaival/tlsutils"
	"github.com/cloudflare/circl/sign/ed25519"
	"github.com/mroth/weightedrand"
	"github.com/multiformats/go-multihash"
	"log"
	"net/http"
	"os"
	"path"
	"strings"
	"time"
)

var publicKey ed25519.PublicKey

var debugOutput = false
var contentServers []*content.Server

var db *content.Database

var redirectListenAddress string

func selectNextContentServer(skip []int) *content.Server {
	inSkip := func(i int) bool {
		for _, c := range skip {
			if i == c {
				return true
			}
		}
		return false
	}

	var choosingList []weightedrand.Choice
	for _, c := range contentServers {
		if !inSkip(c.Index) && c.GetCheckResult() {
			choosingList = append(choosingList, weightedrand.NewChoice(c, c.Weight))
		}
	}

	chooser, err := weightedrand.NewChooser(choosingList...)

	if err != nil {
		return nil
	}

	return chooser.Pick().(*content.Server)
}

func setOtherHeaders(ctx httputils.RequestContext) {
	ctx.SetResponseHeader("Cache-Control", "no-store")
	ctx.SetResponseHeader("Server", "FinalCommander")
	ctx.SetResponseHeader("Vary", "Content-Encoding")
	ctx.SetResponseHeader("X-Content-Type-Options", "nosniff")
	ctx.SetResponseHeader("X-Robots-Tags", "noindex, nofollow, notranslate")
	ctx.SetResponseHeader("Referrer-Policy", "origin")
}
func setCORSHeaders(ctx httputils.RequestContext) {
	ctx.SetResponseHeader("Access-Control-Allow-Credentials", "true")
	ctx.SetResponseHeader("Access-Control-Max-Age", "7200") //Firefox caps this to 86400, Chrome to 7200. Default is 5 seconds (!!!)
	ctx.SetResponseHeader("Access-Control-Allow-Methods", "GET,HEAD,OPTIONS")
	ctx.SetResponseHeader("Access-Control-Allow-Headers", "DNT,ETag,Origin,Accept,Accept-Language,X-Requested-With,Range")
	ctx.SetResponseHeader("Access-Control-Allow-Origin", "*")
	ctx.SetResponseHeader("Access-Control-Expose-Headers", "*")

	//CORP, COEP, COOP
	ctx.SetResponseHeader("Cross-Origin-Embedder-Policy", "require-corp")
	ctx.SetResponseHeader("Cross-Origin-Resource-Policy", "cross-origin")
	ctx.SetResponseHeader("Cross-Origin-Opener-Policy", "unsafe-none")
}

func handleHexHash(pathElements []string, ctx httputils.RequestContext, host string) {

	if len(pathElements) < 3 {
		ctx.SetResponseCode(http.StatusBadRequest)
		return
	}

	hashType := strings.ToLower(pathElements[1])

	hashString := pathElements[2]
	if len(path.Ext(hashString)) > 0 {
		hashString = strings.TrimSuffix(hashString, path.Ext(hashString))
	}
	hash, err := hex.DecodeString(hashString)
	if err != nil {
		ctx.SetResponseCode(http.StatusBadRequest)
		return
	}

	var mh multihash.Multihash
	if hashType == "sha256" && len(hash) == 32 {
		mh, _ = multihash.Encode(hash, multihash.SHA2_256)
	} else if hashType == "md5" && len(hash) == 16 {
		mh, _ = multihash.Encode(hash, multihash.MD5)
	} else {
		ctx.SetResponseCode(http.StatusNotImplemented)
		return
	}
	key := content.NewHashIdentifierFromMultihash(mh)

	var skip []int

	entry := getContentEntry(key)

	if len(pathElements) > 3 {

		if pathElements[3] == "information" {
			ctx.SetResponseHeader("Content-Type", "application/json")
			ctx.SetResponseCode(http.StatusOK)
			if entry != nil {
				b, _ := json.Marshal(struct {
					Known      bool   `json:"known"`
					CID        string `json:"cid"`
					AccessTime int64  `json:"accessTime"`
					CheckTime  int64  `json:"checkTime"`
					Invalid    []int  `json:"invalid,omitempty"`
				}{
					Known:      true,
					CID:        entry.CID().String(),
					AccessTime: entry.AccessTime,
					CheckTime:  entry.CheckTime,
					Invalid:    entry.InvalidList,
				})

				ctx.ServeBytes(b)
			} else {
				b, _ := json.Marshal(struct {
					Known bool   `json:"known"`
					CID   string `json:"cid"`
				}{
					Known: false,
					CID:   key.CID().String(),
				})

				ctx.ServeBytes(b)
			}
			return
		} else if pathElements[3] == "drop" { //Drop key from cache
			ctx.SetResponseHeader("Content-Type", "text/plain")
			ctx.SetResponseCode(http.StatusOK)
			_ = db.RemoveEntry(key)
			if entry != nil && !entry.Key.Equals(key) {
				_ = db.RemoveEntry(&entry.Key)
			}
			ctx.ServeBytes([]byte{})
			return
		}

		//TODO: update entries with these instant returns
		data, err := MakyuuIchaival.Bech32Encoding.DecodeString(pathElements[3])
		if err != nil {
			ctx.SetResponseCode(http.StatusBadRequest)
			return
		}
		skip = utilities.DecodeIntegerList(data)
	}

	if entry != nil {
		oldSkip := skip
		skip = entry.InvalidList
		for _, ci := range oldSkip {
			if !entry.InInvalidList(ci) {
				skip = append(skip, ci)
			}
		}
		contentServer := selectNextContentServer(skip)
		if contentServer == nil {
			ctx.SetResponseCode(http.StatusNotFound)
			return
		}

		ctx.DoRedirect(contentServer.GetContentURL(entry, skip)+host, http.StatusFound)
	} else {
		contentServer := selectNextContentServer(skip)
		if contentServer == nil {
			ctx.SetResponseCode(http.StatusNotFound)
			return
		}

		//TODO: only trigger this when we don't get a 404
		go func() {
			var newInvalidList []int
			var e *content.Entry

			for _, c := range contentServers {
				if !c.GetCheckResult() {
					continue
				}

				result, err := c.CheckEntryKey(key)
				if result != nil {
					if e == nil {
						e = &content.Entry{
							Key:        *result,
							Version:    0,
							AccessTime: time.Now().UTC().Unix(),
							CheckTime:  time.Now().UTC().Unix() + 3600*24*3, // Check sooner after addition, not all servers might have it yet
						}
					}
				} else if err == nil {
					newInvalidList = append(newInvalidList, c.Index)
				}
			}

			if e != nil {
				if !key.IsKey() { //Check for entry already existing
					entry := getContentEntry(&e.Key)

					if entry != nil {
						e = entry
					}
					//Add alias mapping
					_ = db.SetAlias(&content.Alias{
						Key:        *key,
						Identifier: e.Key,
					})
				}

				e.AccessTime = time.Now().UTC().Unix()
				e.InvalidList = newInvalidList
				_ = db.SetEntry(e)
			}
		}()

		ctx.DoRedirect(contentServer.GetHashURL(mh, skip)+host, http.StatusFound)
	}
}

func handle(ctx httputils.RequestContext) {
	if len(ctx.GetHost()) > 0 && len(ctx.GetTLSServerName()) > 0 && strings.Split(ctx.GetHost(), ":")[0] != ctx.GetTLSServerName() { //Prevents rebinding / DNS stuff
		ctx.SetResponseCode(http.StatusNotFound)
		return
	}

	host := ctx.GetHost()
	if host == "" && ctx.GetTLSServerName() != "" {
		host = ctx.GetTLSServerName() + redirectListenAddress
	}

	if host != "" {
		host = "/" + host
	}

	if ctx.IsGet() || ctx.IsHead() {
		if debugOutput {
			log.Printf("Serve %s", ctx.GetPath())
		}
		setOtherHeaders(ctx)
		setCORSHeaders(ctx)

		if ctx.GetPath() == "/status" {
			ctx.SetResponseHeader("Content-Type", "application/json")
			ctx.SetResponseCode(http.StatusOK)
			b, _ := json.MarshalIndent(struct {
				PublicKey string            `json:"public_key,omitempty"`
				Servers   []*content.Server `json:"servers" json:"servers,omitempty"`
			}{
				PublicKey: MakyuuIchaival.Bech32Encoding.EncodeToString(publicKey),
				Servers:   contentServers,
			}, "", " ")

			ctx.ServeBytes(b)
			return
		}
		pathElements := strings.Split(ctx.GetPath(), "/")

		//TODO: handle ni RFC 6920
		handleHexHash(pathElements, ctx, host)
	} else if ctx.IsOptions() {
		setOtherHeaders(ctx)
		setCORSHeaders(ctx)
		ctx.SetResponseCode(http.StatusNoContent)
	} else {
		ctx.SetResponseCode(http.StatusNotImplemented)
	}
}

func getContentEntry(key *content.HashIdentifier) *content.Entry {
	entry := db.GetEntry(*key)

	if entry != nil {
		entry.UpdateAccessTime(db)

		if entry.NeedsCheck() {
			go func() {
				b, _ := entry.Encode() //Encode/decode to copy object
				e := content.DecodeEntry(entry.Key, b)
				e.CheckTime = time.Now().UTC().Unix() + 3600*24*30 //force a check every 30 days as they get fetched

				var newInvalidList []int

				for _, c := range contentServers {
					if !c.GetCheckResult() {
						if e.InInvalidList(c.Index) { //Keep old result if down
							newInvalidList = append(newInvalidList, c.Index)
						}
						continue
					}

					h, err := c.CheckEntryKey(&e.Key)
					if h == nil && err == nil {
						newInvalidList = append(newInvalidList, c.Index)
					}
				}

				e.InvalidList = newInvalidList
				_ = db.SetEntry(e)
			}()
		}
	}
	return entry
}

func checkContentServers() {
	for _, c := range contentServers {
		c.Check()
	}
}

func main() {
	debugOption := flag.Bool("debug", false, "Enable debug output.")
	certificatePath := flag.String("certificate", "", "Path to SSL certificate file.")
	keypairPath := flag.String("keypair", "", "Path to SSL key file.")
	databasePath := flag.String("dbpath", "database", "Path to key/value database.")

	listenAddress := flag.String("listen", ":7777", "address/port to listen on.")

	weightedServerList := flag.String("servers", "", "Weighted list of servers to use. All will use HTTPs. Allowed protocols: orbt. Format [protocol=]address:PORT/WEIGHT,[...]")

	sniAddressOption := flag.String("sni", "", "Define SNI address if desired. Empty will serve any requests regardless.")

	flag.Parse()

	var err error

	debugOutput = *debugOption

	privateKeyEnv := os.Getenv("PRIVATE_KEY")

	if privateKeyEnv != "" {
		log.Print("PRIVATE_KEY is deprecated. Use PRIVATE_SEED instead with seed")
		privateKey, _ := MakyuuIchaival.Bech32Encoding.DecodeString(privateKeyEnv)

		log.Printf("Private Ed25519 seed (keep safe!): %s", MakyuuIchaival.Bech32Encoding.EncodeToString(ed25519.PrivateKey(privateKey).Seed()))
		return
	}

	privateSeedEnv := os.Getenv("PRIVATE_SEED")

	if privateSeedEnv == "" {
		log.Print("No PRIVATE_SEED environment variable specified, generating new identity")
		publicKey, privateKey, _ := ed25519.GenerateKey(rand.Reader)

		log.Printf("Public Ed25519 key (share this): %s", MakyuuIchaival.Bech32Encoding.EncodeToString(publicKey))
		log.Printf("Private Ed25519 seed (keep safe!): %s", MakyuuIchaival.Bech32Encoding.EncodeToString(privateKey.Seed()))
		return
	}

	privateSeed, err := MakyuuIchaival.Bech32Encoding.DecodeString(privateSeedEnv)
	if err != nil {
		log.Fatal(err)
	}

	if len(privateSeed) != ed25519.SeedSize {
		log.Fatal("Wrong Private key length")
	}

	privateKey := ed25519.NewKeyFromSeed(privateSeed)
	publicKey = make([]byte, ed25519.PublicKeySize)
	copy(publicKey, privateKey[ed25519.PublicKeySize:])
	log.Printf("Loaded Private Ed25519 key, Public %s", MakyuuIchaival.Bech32Encoding.EncodeToString(publicKey))

	db, err = content.OpenDatabase(*databasePath)
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	for i, s := range strings.Split(*weightedServerList, ",") {
		cs, err := content.NewContentServerFromArgument(s, i, privateKey)

		if err != nil {
			log.Fatal(err)
		}
		contentServers = append(contentServers, cs)
	}

	checkContentServers()

	go func() {
		ticker := time.NewTicker(1 * time.Minute)
		for range ticker.C {
			checkContentServers()
		}
	}()

	tlsConfiguration, err := tlsutils.NewTLSConfiguration(*certificatePath, *keypairPath, *sniAddressOption)
	if err != nil {
		log.Fatal(err)
	}

	server := &httputils.Server{
		ListenAddress: *listenAddress,
		TLSConfig:     tlsConfiguration,
		EnableHTTP2:   true,
		EnableHTTP3:   true,
		Handler:       handle,
		Debug:         debugOutput,
	}

	redirectListenAddress = *listenAddress

	server.Serve()
}