package userbase

import (
	"bytes"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"slices"
	"strings"
	"time"
	"unicode/utf8"

	"tug/internal/platform/sshkey/openssh"
	"tug/internal/storage"
)

const ExportVersion = 4

var (
	ErrUnknownVersion = errors.New("userbase: unknown export_version")
	ErrNotCanonical   = errors.New("userbase: export document is not canonical")
)

type Document struct {
	Version    int
	Snapshot   storage.UserbaseSnapshot
	ExportedAt time.Time
	DexKeys    json.RawMessage
}

type document struct {
	ExportVersion        int                   `json:"export_version"`
	ExportedAt           int64                 `json:"exported_at"`
	NextUIDNumber        int64                 `json:"next_uid_number"`
	NextGIDNumber        int64                 `json:"next_gid_number"`
	Users                []user                `json:"users"`
	Groups               []group               `json:"groups"`
	Memberships          []membership          `json:"memberships"`
	GoogleLinks          []googleLink          `json:"google_links"`
	Invitations          []invitation          `json:"invitations"`
	InvitationGroups     []invitationGroup     `json:"invitation_groups"`
	PendingConfirmations []pendingConfirmation `json:"pending_confirmations"`
	SSHKeys              []sshKey              `json:"ssh_keys"`
	DexKeys              json.RawMessage       `json:"dex_keys,omitempty"`
}

type user struct {
	ID           string `json:"id"`
	Username     string `json:"username"`
	DisplayName  string `json:"display_name"`
	PasswordHash []byte `json:"password_hash"`
	Disabled     bool   `json:"disabled"`
	UIDNumber    int64  `json:"uid_number"`
	CreatedAt    int64  `json:"created_at"`
}

type group struct {
	ID          string `json:"id"`
	Name        string `json:"name"`
	Description string `json:"description"`
	GIDNumber   int64  `json:"gid_number"`
	CreatedAt   int64  `json:"created_at"`
}

type membership struct {
	UserID  string `json:"user_id"`
	GroupID string `json:"group_id"`
}

type googleLink struct {
	UserID    string `json:"user_id"`
	GoogleSub string `json:"google_sub"`
	LinkedAt  int64  `json:"linked_at"`
}

type invitation struct {
	ID        string  `json:"id"`
	TokenHash []byte  `json:"token_hash"`
	CreatedBy *string `json:"created_by"`
	CreatedAt int64   `json:"created_at"`
	ExpiresAt int64   `json:"expires_at"`
	UsedBy    *string `json:"used_by"`
	UsedAt    *int64  `json:"used_at"`
	RevokedAt *int64  `json:"revoked_at"`
	Label     string  `json:"label"`
	BatchID   *string `json:"batch_id"`
}

type invitationGroup struct {
	InvitationID string  `json:"invitation_id"`
	GroupID      *string `json:"group_id"`
	GroupName    string  `json:"group_name"`
}

type pendingConfirmation struct {
	UserID     string `json:"user_id"`
	CodeHash   []byte `json:"code_hash"`
	IssuedAt   int64  `json:"issued_at"`
	ExpiresAt  int64  `json:"expires_at"`
	WrongCodes int64  `json:"wrong_codes"`
}

type sshKey struct {
	ID          string `json:"id"`
	UserID      string `json:"user_id"`
	Fingerprint string `json:"fingerprint"`
	KeyType     string `json:"key_type"`
	KeyData     []byte `json:"key_data"`
	Label       string `json:"label"`
	AddedAt     int64  `json:"added_at"`
}

func Encode(snapshot storage.UserbaseSnapshot, exportedAt time.Time, dexKeys json.RawMessage) ([]byte, error) {
	doc := fromSnapshot(snapshot, exportedAt.UTC().UnixMicro(), dexKeys)
	if err := checkStrings(doc); err != nil {
		return nil, err
	}
	if len(doc.DexKeys) > 0 && doc.DexKeys[0] != '{' {
		return nil, fmt.Errorf("userbase: encode: dex_keys is not a JSON object")
	}
	return marshal(doc)
}

func ContentDigest(snapshot storage.UserbaseSnapshot, dexKeys json.RawMessage) ([32]byte, error) {
	data, err := Encode(snapshot, time.UnixMicro(0), dexKeys)
	if err != nil {
		return [32]byte{}, err
	}
	return sha256.Sum256(data), nil
}

func Decode(data []byte) (Document, error) {
	var head struct {
		ExportVersion int `json:"export_version"`
	}
	if err := json.Unmarshal(data, &head); err != nil {
		return Document{}, fmt.Errorf("userbase: decode: %w", err)
	}
	switch head.ExportVersion {
	case 1:
		return decodeV1(data)
	case 2:
		return decodeV2(data)
	case 3:
		return decodeV3(data)
	case ExportVersion:
	default:
		return Document{}, fmt.Errorf("%w: export_version %d, this binary reads 1 to %d", ErrUnknownVersion, head.ExportVersion, ExportVersion)
	}
	decoder := json.NewDecoder(bytes.NewReader(data))
	decoder.DisallowUnknownFields()
	var doc document
	if err := decoder.Decode(&doc); err != nil {
		return Document{}, fmt.Errorf("userbase: decode: %w", err)
	}
	out := Document{Version: ExportVersion, Snapshot: toSnapshot(doc), ExportedAt: time.UnixMicro(doc.ExportedAt).UTC(), DexKeys: doc.DexKeys}
	again, err := Encode(out.Snapshot, out.ExportedAt, out.DexKeys)
	if err != nil {
		return Document{}, err
	}
	if !bytes.Equal(again, data) {
		return Document{}, ErrNotCanonical
	}
	if err := checkSSHKeys(doc.SSHKeys); err != nil {
		return Document{}, err
	}
	return out, nil
}

func marshal(doc any) ([]byte, error) {
	var buf bytes.Buffer
	encoder := json.NewEncoder(&buf)
	encoder.SetEscapeHTML(false)
	if err := encoder.Encode(doc); err != nil {
		return nil, fmt.Errorf("userbase: encode: %w", err)
	}
	return buf.Bytes(), nil
}

func fromSnapshot(s storage.UserbaseSnapshot, exportedAt int64, dexKeys json.RawMessage) document {
	doc := document{
		ExportVersion:        ExportVersion,
		ExportedAt:           exportedAt,
		NextUIDNumber:        s.NextUIDNumber,
		NextGIDNumber:        s.NextGIDNumber,
		Users:                make([]user, 0, len(s.Users)),
		Groups:               make([]group, 0, len(s.Groups)),
		Memberships:          make([]membership, 0, len(s.Memberships)),
		GoogleLinks:          make([]googleLink, 0, len(s.GoogleLinks)),
		Invitations:          make([]invitation, 0, len(s.Invitations)),
		InvitationGroups:     make([]invitationGroup, 0, len(s.InvitationGroups)),
		PendingConfirmations: make([]pendingConfirmation, 0, len(s.PendingConfirmations)),
		SSHKeys:              make([]sshKey, 0, len(s.SSHKeys)),
		DexKeys:              dexKeys,
	}
	for _, u := range s.Users {
		doc.Users = append(doc.Users, user(u))
	}
	for _, g := range s.Groups {
		doc.Groups = append(doc.Groups, group(g))
	}
	for _, m := range s.Memberships {
		doc.Memberships = append(doc.Memberships, membership(m))
	}
	for _, l := range s.GoogleLinks {
		doc.GoogleLinks = append(doc.GoogleLinks, googleLink(l))
	}
	for _, i := range s.Invitations {
		doc.Invitations = append(doc.Invitations, invitation(i))
	}
	for _, g := range s.InvitationGroups {
		doc.InvitationGroups = append(doc.InvitationGroups, invitationGroup(g))
	}
	for _, p := range s.PendingConfirmations {
		doc.PendingConfirmations = append(doc.PendingConfirmations, pendingConfirmation(p))
	}
	for _, k := range s.SSHKeys {
		doc.SSHKeys = append(doc.SSHKeys, sshKey(k))
	}
	doc.Users = sortedRows(doc.Users, func(u user) (string, string) { return u.ID, "" })
	doc.Groups = sortedRows(doc.Groups, groupKey)
	doc.Memberships = sortedRows(doc.Memberships, membershipKey)
	doc.GoogleLinks = sortedRows(doc.GoogleLinks, func(l googleLink) (string, string) { return l.UserID, "" })
	doc.Invitations = sortedRows(doc.Invitations, func(i invitation) (string, string) { return i.ID, "" })
	doc.InvitationGroups = sortedRows(doc.InvitationGroups, invitationGroupKey)
	doc.PendingConfirmations = sortedRows(doc.PendingConfirmations, pendingKey)
	doc.SSHKeys = sortedRows(doc.SSHKeys, sshKeyKey)
	return doc
}

func sortedRows[T any](rows []T, key func(T) (string, string)) []T {
	out := append(make([]T, 0, len(rows)), rows...)
	slices.SortFunc(out, func(a, b T) int {
		a1, a2 := key(a)
		b1, b2 := key(b)
		if c := strings.Compare(a1, b1); c != 0 {
			return c
		}
		return strings.Compare(a2, b2)
	})
	return out
}

func groupKey(g group) (string, string)                     { return g.ID, "" }
func membershipKey(m membership) (string, string)           { return m.UserID, m.GroupID }
func invitationGroupKey(g invitationGroup) (string, string) { return g.InvitationID, g.GroupName }
func pendingKey(p pendingConfirmation) (string, string)     { return p.UserID, "" }
func sshKeyKey(k sshKey) (string, string)                   { return k.ID, "" }

func toSnapshot(doc document) storage.UserbaseSnapshot {
	s := storage.UserbaseSnapshot{NextUIDNumber: doc.NextUIDNumber, NextGIDNumber: doc.NextGIDNumber}
	for _, u := range doc.Users {
		s.Users = append(s.Users, storage.UserbaseUser(u))
	}
	for _, g := range doc.Groups {
		s.Groups = append(s.Groups, storage.UserbaseGroup(g))
	}
	for _, m := range doc.Memberships {
		s.Memberships = append(s.Memberships, storage.UserbaseMembership(m))
	}
	for _, l := range doc.GoogleLinks {
		s.GoogleLinks = append(s.GoogleLinks, storage.UserbaseGoogleLink(l))
	}
	for _, i := range doc.Invitations {
		s.Invitations = append(s.Invitations, storage.UserbaseInvitation(i))
	}
	for _, g := range doc.InvitationGroups {
		s.InvitationGroups = append(s.InvitationGroups, storage.UserbaseInvitationGroup(g))
	}
	for _, p := range doc.PendingConfirmations {
		s.PendingConfirmations = append(s.PendingConfirmations, storage.UserbasePendingConfirmation(p))
	}
	for _, k := range doc.SSHKeys {
		s.SSHKeys = append(s.SSHKeys, storage.UserbaseSSHKey(k))
	}
	return s
}

func checkStrings(doc document) error {
	var bad []string
	check := func(field string, values ...*string) {
		for _, value := range values {
			if value != nil && !utf8.ValidString(*value) {
				bad = append(bad, field)
				return
			}
		}
	}
	for i := range doc.Users {
		u := &doc.Users[i]
		check(fmt.Sprintf("users[%d]", i), &u.ID, &u.Username, &u.DisplayName)
	}
	for i := range doc.Groups {
		g := &doc.Groups[i]
		check(fmt.Sprintf("groups[%d]", i), &g.ID, &g.Name, &g.Description)
	}
	for i := range doc.Memberships {
		m := &doc.Memberships[i]
		check(fmt.Sprintf("memberships[%d]", i), &m.UserID, &m.GroupID)
	}
	for i := range doc.GoogleLinks {
		l := &doc.GoogleLinks[i]
		check(fmt.Sprintf("google_links[%d]", i), &l.UserID, &l.GoogleSub)
	}
	for i := range doc.Invitations {
		v := &doc.Invitations[i]
		check(fmt.Sprintf("invitations[%d]", i), &v.ID, v.CreatedBy, v.UsedBy, &v.Label, v.BatchID)
	}
	for i := range doc.InvitationGroups {
		g := &doc.InvitationGroups[i]
		check(fmt.Sprintf("invitation_groups[%d]", i), &g.InvitationID, g.GroupID, &g.GroupName)
	}
	for i := range doc.PendingConfirmations {
		p := &doc.PendingConfirmations[i]
		check(fmt.Sprintf("pending_confirmations[%d]", i), &p.UserID)
	}
	for i := range doc.SSHKeys {
		k := &doc.SSHKeys[i]
		check(fmt.Sprintf("ssh_keys[%d]", i), &k.ID, &k.UserID, &k.Fingerprint, &k.KeyType, &k.Label)
	}
	if len(bad) > 0 {
		return fmt.Errorf("userbase: encode: not valid UTF-8: %s", strings.Join(bad, ", "))
	}
	return nil
}

func checkSSHKeys(keys []sshKey) error {
	for i, k := range keys {
		parsed, err := openssh.FromData(k.KeyType, k.KeyData)
		if err != nil {
			return fmt.Errorf("userbase: decode: ssh_keys[%d]: %w", i, err)
		}
		if parsed.Fingerprint != k.Fingerprint {
			return fmt.Errorf("userbase: decode: ssh_keys[%d]: fingerprint %s is not that of the key's data", i, k.Fingerprint)
		}
	}
	return nil
}