// Copyright 2026 Rose Hogenson <r@roseh.moe>
//
// This file is part of roseh.moe.
//
// roseh.moe is free software: you can redistribute it and/or modify it
// under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License,
// or (at your option) any later version.
//
// roseh.moe is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty
// of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with roseh.moe. If not, see
// <https://www.gnu.org/licenses/>.

package auth

import (
	"bytes"
	"crypto"
	"crypto/ecdsa"
	"crypto/hmac"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json/v2"
	"fmt"
	"io"
	"log"
	"net/url"
	"slices"
	"strings"
	"sync"
	"time"
	"unicode"
)

var SecretKey []byte

func MAC(msg []byte, purpose string) []byte {
	mac := hmac.New(sha256.New, SecretKey)
	io.WriteString(mac, purpose)
	io.WriteString(mac, ":")
	mac.Write(msg)
	return mac.Sum(nil)
}

func SignMsg(msg, purpose string) string {
	return msg + base64.RawURLEncoding.EncodeToString(MAC([]byte(msg), purpose))
}

func VerifyHMAC(signedMsg, purpose string) (string, bool) {
	macSize := base64.RawURLEncoding.EncodedLen(sha256.Size)
	if len(signedMsg) < macSize {
		return "", false
	}
	msg, msgMAC := signedMsg[:len(signedMsg)-macSize], signedMsg[len(signedMsg)-macSize:]
	expectedMAC := base64.RawURLEncoding.EncodeToString(MAC([]byte(msg), purpose))
	if !hmac.Equal([]byte(msgMAC), []byte(expectedMAC)) {
		return "", false
	}
	return msg, true
}

type Passkey struct {
	ID  string `xml:",attr"`
	Key string `xml:",attr"`
}

func LoadPasskeys(passkeys []Passkey) (map[string][]byte, error) {
	names := make(map[string]bool)
	out := make(map[string][]byte)
	namesByID := make(map[string]string)
	for _, passkey := range passkeys {
		if names[passkey.ID] {
			return nil, fmt.Errorf("duplicate name %q", passkey.ID)
		}
		names[passkey.ID] = true
		keyBytes, err := base64.RawURLEncoding.DecodeString(strings.Map(func(r rune) rune {
			if unicode.IsSpace(r) {
				return -1
			}
			return r
		}, passkey.Key))
		if err != nil {
			return nil, fmt.Errorf("invalid key for passkey %q: %s", passkey.ID, err)
		}
		if len(keyBytes) < 16 {
			return nil, fmt.Errorf("invalid key for passkey %q: bad user ID: too short", passkey.ID)
		}
		id, key := keyBytes[:16], keyBytes[16:]
		if _, err := sigVerifier(key); err != nil {
			return nil, fmt.Errorf("invalid key for passkey %q: %s", passkey.ID, err)
		}
		if _, ok := out[string(id)]; ok {
			return nil, fmt.Errorf("passkeys %q and %q have duplicate ID %d", namesByID[string(id)], passkey.ID, id)
		}
		out[string(id)] = key
		namesByID[string(id)] = passkey.ID
	}
	return out, nil
}

func sigVerifier(keyBytes []byte) (func(msg, sig []byte) bool, error) {
	key, err := x509.ParsePKIXPublicKey(keyBytes)
	if err != nil {
		return nil, err
	}
	switch key := key.(type) {
	case *ecdsa.PublicKey:
		return func(msg, sig []byte) bool {
			hash := sha256.Sum256(msg)
			return ecdsa.VerifyASN1(key, hash[:], sig)
		}, nil
	case *rsa.PublicKey:
		return func(msg, sig []byte) bool {
			hash := sha256.Sum256(msg)
			return rsa.VerifyPKCS1v15(key, crypto.SHA256, hash[:], sig) == nil
		}, nil
	default:
		return nil, fmt.Errorf("unsupported key type")
	}
}

func VerifyPubKey(pubKey, msg, sig []byte) error {
	verify, err := sigVerifier(pubKey)
	if err != nil {
		return err
	}
	if !verify(msg, sig) {
		return fmt.Errorf("message failed validation")
	}
	return nil
}

const challengeExpiry = 10 * time.Minute

var (
	challengeMu sync.Mutex
	challenges  = make(map[string]time.Time)
)

func NewChallenge() ([]byte, bool) {
	buf := make([]byte, 32)
	rand.Read(buf)
	now := time.Now()
	challengeMu.Lock()
	if len(challenges) >= 100 {
		challengeMu.Unlock()
		return nil, false
	}
	challenges[string(buf)] = now
	challengeMu.Unlock()
	return buf, true
}

func Login(user, authData, clientData, sig []byte, passkeys map[string][]byte, selfURL *url.URL) error {
	pubKey, ok := passkeys[string(user)]
	if !ok {
		return fmt.Errorf("no user with the given ID")
	}
	clientDataHash := sha256.Sum256(clientData)
	if err := VerifyPubKey(pubKey, slices.Concat(authData, clientDataHash[:]), sig); err != nil {
		return fmt.Errorf("challenge failed: %s", err)
	}
	var cd struct {
		Type        string `json:"type"`
		Challenge   string `json:"challenge"`
		Origin      string `json:"origin"`
		CrossOrigin bool   `json:"crossOrigin"`
	}
	if err := json.Unmarshal(clientData, &cd); err != nil {
		return fmt.Errorf("bad client data: invalid JSON: %s", err)
	}
	challenge, err := base64.RawURLEncoding.DecodeString(cd.Challenge)
	if err != nil {
		return fmt.Errorf("bad client data: bad challenge: %s", err)
	}
	challengeMu.Lock()
	initiated, ok := challenges[string(challenge)]
	delete(challenges, string(challenge))
	challengeMu.Unlock()
	if !ok {
		return fmt.Errorf("request already consumed")
	}
	if time.Since(initiated) > challengeExpiry {
		return fmt.Errorf("challenge is too old; please try again")
	}
	if cd.Type != "webauthn.get" || cd.Origin != selfURL.String() || cd.CrossOrigin {
		return fmt.Errorf("bad client data")
	}
	expectedRPHash := sha256.Sum256([]byte(selfURL.Hostname()))
	const flagUserPresence, flagUserVerification = 0x01, 0x04
	if len(authData) < 37 || !bytes.Equal(authData[:32], expectedRPHash[:]) || authData[32]&(flagUserPresence|flagUserVerification) != (flagUserPresence|flagUserVerification) {
		return fmt.Errorf("bad authenticator data")
	}
	return nil
}

func CleanChallenges() {
	log.Printf("Cleaning expired challenges")
	now := time.Now()
	challengeMu.Lock()
	for challenge, initiated := range challenges {
		if now.Sub(initiated) > challengeExpiry {
			delete(challenges, challenge)
		}
	}
	challengeMu.Unlock()
}