// main.go
// Web PGP Key Server with Bootstrap UI, HTMX live search,
// WKD (Web Key Directory) + minimal HKP-compatible lookup,
// and automatic fingerprint parsing from uploaded ASCII-armored keys.
//
// Build & Run
//
//	go mod init example.com/pgp-keyserver
//	go get github.com/ProtonMail/go-crypto/openpgp@v0.0.0-20230828082145-5dc2f9f7b5c1
//	go get github.com/ProtonMail/go-crypto/openpgp/armor@v0.0.0-20230828082145-5dc2f9f7b5c1
//	go mod tidy
//	go run .
//
// Notes
// - This server publishes PUBLIC keys only.
// - WKD uses z-base-32(SHA1(strings.ToLower(addr-spec))) per spec.
// - HKP here is minimal: /pks/lookup?op=get|index&search=<term> (email/fpr/substring).
// - Protect /upload behind auth in production.
package main

import (
	"bytes"
	"crypto/sha1"
	"embed"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"html/template"
	"io"
	"log"
	"mime"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	"github.com/ProtonMail/go-crypto/openpgp"
)

//go:embed assets/** templates/*
var contentFS embed.FS

const (
	maxUploadSize = 2 << 20 // 2 MiB
	dataDir       = "data"
	keysDir       = "data/keys"
	indexFile     = "data/index.json"
)

type KeyRecord struct {
	ID          string    `json:"id"`
	Name        string    `json:"name"`
	Email       string    `json:"email"`
	Fingerprint string    `json:"fingerprint"`
	Filename    string    `json:"filename"`
	CreatedAt   time.Time `json:"created_at"`
}

type store struct {
	mu   sync.RWMutex
	list []KeyRecord
	byID map[string]KeyRecord
}

func newStore() *store { return &store{byID: make(map[string]KeyRecord)} }

// persistSnapshot schreibt eine bereits kopierte Liste auf Disk.
func (s *store) persistSnapshot(list []KeyRecord) error {
	b, err := json.MarshalIndent(list, "", "  ")
	if err != nil {
		return err
	}
	if err := os.MkdirAll(dataDir, 0o755); err != nil {
		return err
	}
	tmp := indexFile + ".tmp"
	if err := os.WriteFile(tmp, b, 0o644); err != nil {
		return err
	}
	return os.Rename(tmp, indexFile)
}

func (s *store) upsert(rec KeyRecord) error {
	s.mu.Lock()
	// mutate
	idx := -1
	for i, it := range s.list {
		if it.ID == rec.ID {
			idx = i
			break
		}
	}
	if idx >= 0 {
		s.list[idx] = rec
	} else {
		s.list = append(s.list, rec)
	}
	s.byID[rec.ID] = rec
	sort.SliceStable(s.list, func(i, j int) bool {
		return s.list[i].CreatedAt.After(s.list[j].CreatedAt)
	})
	// Snapshot erstellen, bevor wir den Lock freigeben
	snap := make([]KeyRecord, len(s.list))
	copy(snap, s.list)
	s.mu.Unlock()

	// Jetzt ohne Lock auf Disk schreiben
	return s.persistSnapshot(snap)
}

func (s *store) delete(id string) error {
	s.mu.Lock()
	idx := -1
	for i, it := range s.list {
		if it.ID == id {
			idx = i
			break
		}
	}
	if idx < 0 {
		s.mu.Unlock()
		return os.ErrNotExist
	}
	rec := s.list[idx]
	_ = os.Remove(filepath.Join(keysDir, rec.Filename))
	s.list = append(s.list[:idx], s.list[idx+1:]...)
	delete(s.byID, id)
	snap := make([]KeyRecord, len(s.list))
	copy(snap, s.list)
	s.mu.Unlock()

	return s.persistSnapshot(snap)
}

// Optional: beim Start laden (unverändert, aber ohne Locks von außen aufzurufen)
func (s *store) load() error {
	_ = os.MkdirAll(keysDir, 0o755)
	b, err := os.ReadFile(indexFile)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return err
	}
	var items []KeyRecord
	if err := json.Unmarshal(b, &items); err != nil {
		return err
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	s.list = items
	s.byID = make(map[string]KeyRecord, len(items))
	for _, it := range items {
		s.byID[it.ID] = it
	}
	return nil
}

func (s *store) all() []KeyRecord {
	s.mu.RLock()
	defer s.mu.RUnlock()
	cp := make([]KeyRecord, len(s.list))
	copy(cp, s.list)
	return cp
}
func (s *store) get(id string) (KeyRecord, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	it, ok := s.byID[id]
	return it, ok
}

func (s *store) search(q string) []KeyRecord {
	s.mu.RLock()
	defer s.mu.RUnlock()
	q = strings.ToLower(strings.TrimSpace(q))
	if q == "" {
		return append([]KeyRecord(nil), s.list...)
	}
	normFPR := normalizeFPR(q)
	var out []KeyRecord
	for _, it := range s.list {
		if containsFold(it.Name, q) || containsFold(it.Email, q) || strings.Contains(strings.ToLower(it.Fingerprint), q) || normalizeFPR(it.Fingerprint) == normFPR {
			out = append(out, it)
		}
	}
	return out
}

// --- helpers ---
var safeFilenameRe = regexp.MustCompile(`[^a-zA-Z0-9_.@+-]+`)

func sanitizeFilename(name string) string {
	if name == "" {
		name = fmt.Sprintf("key-%d", time.Now().Unix())
	}
	name = safeFilenameRe.ReplaceAllString(name, "-")
	if len(name) > 120 {
		name = name[:120]
	}
	if !strings.HasSuffix(name, ".asc") {
		name += ".asc"
	}
	return name
}

func genID(email, fingerprint string) string {
	base := strings.ToLower(strings.TrimSpace(email))
	fp := strings.ToUpper(strings.TrimSpace(fingerprint))
	if fp == "" {
		fp = fmt.Sprintf("%d", time.Now().UnixNano())
	}
	id := base + "--" + fp
	id = strings.ReplaceAll(id, "@", "_at_")
	id = safeFilenameRe.ReplaceAllString(id, "-")
	return id
}

func containsFold(h, n string) bool { return strings.Contains(strings.ToLower(h), strings.ToLower(n)) }

func isASCII(s string) bool { return utf8.ValidString(s) && !strings.ContainsRune(s, '�') }

func normalizeFPR(s string) string {
	s = strings.ToUpper(strings.TrimSpace(s))
	s = strings.TrimPrefix(s, "0X")
	s = strings.ReplaceAll(s, " ", "")
	return s
}

// --- WKD utilities ---
var zbase32Alphabet = []byte("ybndrfg8ejkmcpqxot1uwisza345h769")

func zbase32Encode(b []byte) string {
	var out []byte
	var bits, val uint
	for _, by := range b {
		val = (val << 8) | uint(by)
		bits += 8
		for bits >= 5 {
			idx := (val >> (bits - 5)) & 31
			out = append(out, zbase32Alphabet[idx])
			bits -= 5
		}
	}
	if bits > 0 {
		idx := (val << (5 - bits)) & 31
		out = append(out, zbase32Alphabet[idx])
	}
	return string(out)
}

// wkdHash: z-base-32(SHA1(strings.ToLower(addr-spec))) + domain
func wkdHash(email string) (hash string, domain string) {
	email = strings.ToLower(strings.TrimSpace(email))
	parts := strings.Split(email, "@")
	if len(parts) != 2 {
		return "", ""
	}
	domain = parts[1]
	s := sha1.Sum([]byte(email))
	return zbase32Encode(s[:]), domain
}

func main() {
	_ = os.MkdirAll(keysDir, 0o755)
	_ = mime.AddExtensionType(".asc", "application/pgp-keys")

	pageTmpl, err := template.ParseFS(
		contentFS,
		"templates/layout.html",
		"templates/index.html",
		"templates/rows.html", // <— hinzufügen!
	)
	if err != nil {
		log.Fatal(err)
	}

	st := newStore()
	if err := st.load(); err != nil {
		log.Fatal(err)
	}

	mux := http.NewServeMux()

	// Home
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		data := map[string]any{"Items": st.all()}
		if err := pageTmpl.ExecuteTemplate(w, "layout", data); err != nil {
			http.Error(w, err.Error(), 500)
		}
	})
	// Live search (HTMX)
	mux.HandleFunc("/search", func(w http.ResponseWriter, r *http.Request) {
		q := r.URL.Query().Get("q")
		items := st.search(q)
		// Nutze das bereits geladene Set:
		if err := pageTmpl.ExecuteTemplate(w, "rows", items); err != nil {
			http.Error(w, err.Error(), 500)
		}
	})
	// Upload with automatic fingerprint parsing
	mux.HandleFunc("/upload", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}
		if err := r.ParseMultipartForm(maxUploadSize); err != nil {
			http.Error(w, "invalid form", http.StatusBadRequest)
			return
		}
		name := strings.TrimSpace(r.FormValue("name"))
		email := strings.TrimSpace(r.FormValue("email"))
		userFPR := strings.TrimSpace(r.FormValue("fingerprint")) // optional override
		file, hdr, err := r.FormFile("file")
		if err != nil {
			http.Error(w, "missing file", http.StatusBadRequest)
			return
		}
		defer file.Close()
		var buf bytes.Buffer
		lr := io.LimitedReader{R: file, N: maxUploadSize}
		if _, err := io.Copy(&buf, &lr); err != nil {
			http.Error(w, "read error", http.StatusBadRequest)
			return
		}
		b := buf.Bytes()
		if !isASCII(string(b)) || !bytes.Contains(b, []byte("-----BEGIN PGP PUBLIC KEY BLOCK-----")) {
			http.Error(w, "file must be ASCII-armored PGP public key (.asc)", http.StatusBadRequest)
			return
		}
		// Parse fingerprint
		autoFPR, parseErr := parseFingerprintFromASCII(b)
		fpr := userFPR
		if fpr == "" && parseErr == nil {
			fpr = autoFPR
		}
		if fpr == "" {
			http.Error(w, "could not parse fingerprint; please provide it manually", http.StatusBadRequest)
			return
		}
		fpr = strings.ToUpper(strings.ReplaceAll(fpr, " ", ""))

		base := sanitizeFilename(hdr.Filename)
		if base == ".asc" || base == "" {
			base = sanitizeFilename(email)
		}
		path := filepath.Join(keysDir, base)
		if err := os.WriteFile(path, b, 0o644); err != nil {
			http.Error(w, "save error", 500)
			return
		}

		rec := KeyRecord{ID: genID(email, fpr), Name: name, Email: email, Fingerprint: fpr, Filename: base, CreatedAt: time.Now()}
		if err := st.upsert(rec); err != nil {
			http.Error(w, "index error", 500)
			return
		}
		http.Redirect(w, r, "/", http.StatusSeeOther)
	})

	// Serve/download by ID
	mux.HandleFunc("/keys/", func(w http.ResponseWriter, r *http.Request) {
		id := strings.TrimPrefix(r.URL.Path, "/keys/")
		rec, ok := st.get(id)
		if !ok {
			http.NotFound(w, r)
			return
		}
		p := filepath.Join(keysDir, rec.Filename)
		data, err := os.ReadFile(p)
		if err != nil {
			http.NotFound(w, r)
			return
		}
		w.Header().Set("Content-Type", "application/pgp-keys; charset=utf-8")
		w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"%s\"", rec.Filename))
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write(data)
	})

	// Inline preview
	mux.HandleFunc("/view/", func(w http.ResponseWriter, r *http.Request) {
		id := strings.TrimPrefix(r.URL.Path, "/view/")
		rec, ok := st.get(id)
		if !ok {
			http.NotFound(w, r)
			return
		}
		p := filepath.Join(keysDir, rec.Filename)
		data, err := os.ReadFile(p)
		if err != nil {
			http.NotFound(w, r)
			return
		}
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		fmt.Fprintf(w, "<html><head><meta charset=\"utf-8\"><title>%s</title><link href=\"https://cdn.jsdelivr.net/npm/bootstrap@5.3.3/dist/css/bootstrap.min.css\" rel=\"stylesheet\"></head><body class=\"p-4\"><div class=\"container\"><h1 class=\"h4 mb-3\">%s &lt;%s&gt;</h1><pre class=\"bg-light p-3 rounded border\">%s</pre><a class=\"btn btn-primary mt-3\" href=\"/keys/%s\">Download .asc</a> <a class=\"btn btn-outline-secondary mt-3\" href=\"/\">Back</a></div></body></html>", template.HTMLEscapeString(rec.Email), template.HTMLEscapeString(rec.Name), template.HTMLEscapeString(rec.Email), template.HTMLEscapeString(string(data)), template.HTMLEscapeString(rec.ID))
	})

	// --- WKD (Web Key Directory) ---
	mux.HandleFunc("/.well-known/openpgpkey/policy", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("# WKD policy"))
	})
	// direct method: /.well-known/openpgpkey/hu/<hash>
	mux.HandleFunc("/.well-known/openpgpkey/hu/", func(w http.ResponseWriter, r *http.Request) {
		pathHash := strings.TrimPrefix(r.URL.Path, "/.well-known/openpgpkey/hu/")
		var match *KeyRecord
		for _, it := range st.all() {
			h, _ := wkdHash(it.Email)
			if h == pathHash {
				match = &it
				break
			}
		}
		if match == nil {
			http.NotFound(w, r)
			return
		}
		data, err := os.ReadFile(filepath.Join(keysDir, match.Filename))
		if err != nil {
			http.NotFound(w, r)
			return
		}
		w.Header().Set("Content-Type", "application/pgp-keys")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write(data)
	})
	// advanced method: /openpgpkey/<domain>/hu/<hash>
	mux.HandleFunc("/openpgpkey/", func(w http.ResponseWriter, r *http.Request) {
		parts := strings.Split(strings.TrimPrefix(r.URL.Path, "/openpgpkey/"), "/")
		if len(parts) != 3 || parts[1] != "hu" {
			http.NotFound(w, r)
			return
		}
		domain, hash := parts[0], parts[2]
		var match *KeyRecord
		for _, it := range st.all() {
			h, d := wkdHash(it.Email)
			if h == hash && strings.EqualFold(d, domain) {
				match = &it
				break
			}
		}
		if match == nil {
			http.NotFound(w, r)
			return
		}
		data, err := os.ReadFile(filepath.Join(keysDir, match.Filename))
		if err != nil {
			http.NotFound(w, r)
			return
		}
		w.Header().Set("Content-Type", "application/pgp-keys")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write(data)
	})

	// --- Minimal HKP-compatible lookup ---
	mux.HandleFunc("/pks/lookup", func(w http.ResponseWriter, r *http.Request) {
		q := r.URL.Query()
		op := strings.ToLower(q.Get("op"))
		term := strings.TrimSpace(q.Get("search"))
		term, _ = url.QueryUnescape(term)
		if term == "" {
			http.Error(w, "missing search", http.StatusBadRequest)
			return
		}
		items := st.search(term)
		switch op {
		case "get":
			w.Header().Set("Content-Type", "application/pgp-keys")
			for i, it := range items {
				data, err := os.ReadFile(filepath.Join(keysDir, it.Filename))
				if err == nil {
					_, _ = w.Write(data)
					if i < len(items)-1 {
						_, _ = w.Write([]byte(""))
					}
				}
			}
		case "index":
			w.Header().Set("Content-Type", "text/plain; charset=utf-8")
			for _, it := range items {
				fmt.Fprintf(w, "pub:%s:%s:%s", it.Fingerprint, it.Email, it.Name)
			}
		default:
			http.Error(w, "unsupported op", http.StatusBadRequest)
		}
	})

	// Static assets
	fs := http.FileServer(http.FS(contentFS))
	mux.Handle("/assets/", http.StripPrefix("/", fs))

	addr := ":8080"
	log.Printf("PGP Key Server listening on %s", addr)
	log.Fatal(http.ListenAndServe(addr, withSecurityHeaders(mux)))
}

// parseFingerprintFromASCII parses the first entity in an armored key and returns the hex fingerprint.
func parseFingerprintFromASCII(b []byte) (string, error) {
	ents, err := openpgp.ReadArmoredKeyRing(bytes.NewReader(b))
	if err != nil {
		return "", err
	}
	if len(ents) == 0 || ents[0].PrimaryKey == nil {
		return "", fmt.Errorf("no primary key in block")
	}
	fp := ents[0].PrimaryKey.Fingerprint
	return strings.ToUpper(hex.EncodeToString(fp[:])), nil
}

func withSecurityHeaders(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("X-Content-Type-Options", "nosniff")
		w.Header().Set("X-Frame-Options", "DENY")
		w.Header().Set("Referrer-Policy", "no-referrer")
		next.ServeHTTP(w, r)
	})
}