/* SPDX-License-Identifier: GPL-2.0
 *
 * Copyright (C) 2021 Jason A. Donenfeld. All Rights Reserved.
 */

package main

import (
	"bufio"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"log"
	"net/http"
	"os"
	"path"
	"strings"
	"sync"
	"time"
)

type ShortLink struct {
	mu            sync.RWMutex
	diskCachePath string
	links         map[string]string
	dirty         bool
	cacher        *time.Timer
}

func NewShortLink(diskCachePath string) *ShortLink {
	sl := &ShortLink{
		diskCachePath: diskCachePath,
		links:         make(map[string]string, 4096),
	}
	f, err := os.Open(sl.diskCachePath)
	if err == nil {
		defer f.Close()
		scanner := bufio.NewScanner(f)
		for scanner.Scan() {
			parts := strings.SplitN(scanner.Text(), "\t", 2)
			if len(parts) != 2 || len(parts[0]) == 0 || len(parts[1]) == 0 {
				continue
			}
			sl.links[parts[0]] = parts[1]
		}
	}
	return sl
}

func (sl *ShortLink) Set(k, v string, overwrite bool) bool {
	sl.mu.Lock()
	defer sl.mu.Unlock()
	if wantV, ok := sl.links[k]; ok && (!overwrite || wantV == v) {
		return false
	}
	sl.links[k] = v
	if !sl.dirty {
		sl.dirty = true
		time.AfterFunc(time.Second*10, sl.SaveToDisk)
	}
	return true
}

func (sl *ShortLink) New(v string) string {
	for {
		var bytes [3]byte
		_, err := rand.Read(bytes[:])
		if err != nil {
			log.Panicf("RNG is broken: %v", err)
		}
		k := base64.RawURLEncoding.EncodeToString(bytes[:])
		sl.mu.RLock()
		_, exists := sl.links[k]
		sl.mu.RUnlock()
		if !exists && sl.Set(k, v, false) {
			return k
		}
	}
}

func (sl *ShortLink) HandleRequest(w http.ResponseWriter, r *http.Request) {
	w.Header().Add("Server", "WurGurBoo/1.0")
	target := "https://www.wireguard.com/"
	wantL, key := path.Split(r.URL.Path)
	if wantL == "/l/" {
		sl.mu.RLock()
		v, ok := sl.links[key]
		sl.mu.RUnlock()
		if ok {
			target = v
		}
	}
	http.Redirect(w, r, target, 302)
}

func (sl *ShortLink) SaveToDisk() {
	sl.mu.Lock()
	dirty := sl.dirty
	sl.dirty = false
	sl.mu.Unlock()
	if !dirty {
		return
	}
	f, err := os.Create(sl.diskCachePath + ".tmp")
	if err != nil {
		log.Printf("Unable to backup short link db: %v", err)
		return
	}
	defer f.Close()
	sl.mu.RLock()
	for k, v := range sl.links {
		fmt.Fprintf(f, "%s\t%s\n", k, v)
	}
	sl.mu.RUnlock()
	err = os.Rename(f.Name(), sl.diskCachePath)
	if err != nil {
		log.Printf("Unable to backup short link db: %v", err)
		return
	}
}