package auth

import (
	"encoding/json"

	"golang.org/x/crypto/ssh"

	"git.maze.io/maze/conduit/policy/input"
)

const (
	TypeCertificate = "certificate"
	TypeToken       = "token"
	TypePassword    = "password"
	TypePublicKey   = "publickey"
)

type Attributes struct {
	Groups     []string          `json:"groups"`
	Principals []string          `json:"principals"`
	Options    map[string]string `json:"options"`
	Extensions map[string]string `json:"extensions"`
	Source     any               `json:"source"`
}

func MakeAttributes() Attributes {
	return Attributes{
		Groups:     make([]string, 0),
		Principals: make([]string, 0),
		Options:    make(map[string]string),
		Extensions: make(map[string]string),
	}
}

type Principal interface {
	Type() string // "user", "host", "key", etc.

	// Identity is the user identity.
	Identity() string

	// Attributes for the principal.
	Attributes() Attributes
}

type PasswordPrincipal struct {
	User string
	Attr Attributes
}

func MakePasswordPrincipal(meta ssh.ConnMetadata) PasswordPrincipal {
	attr := MakeAttributes()
	attr.Source = meta
	return PasswordPrincipal{
		User: meta.User(),
		Attr: attr,
	}
}

func (p PasswordPrincipal) Type() string           { return TypePassword }
func (p PasswordPrincipal) Identity() string       { return p.User }
func (p PasswordPrincipal) Attributes() Attributes { return p.Attr }
func (p PasswordPrincipal) MarshalJSON() ([]byte, error) {
	return json.Marshal(struct {
		Kind string     `json:"type"`
		User string     `json:"identity"`
		Attr Attributes `json:"attr"`
	}{
		p.Type(),
		p.Identity(),
		p.Attributes(),
	})
}

type TokenPrincipal struct {
	User string     `json:"identity"`
	Attr Attributes `json:"attr"`
}

func (p TokenPrincipal) Type() string           { return TypeToken }
func (p TokenPrincipal) Identity() string       { return p.User }
func (p TokenPrincipal) Attributes() Attributes { return p.Attr }
func (p TokenPrincipal) MarshalJSON() ([]byte, error) {
	return json.Marshal(struct {
		Kind string     `json:"type"`
		User string     `json:"identity"`
		Attr Attributes `json:"attr"`
	}{
		p.Type(),
		p.Identity(),
		p.Attributes(),
	})
}

type PublicKeyPrincipal struct {
	User string
	ssh.PublicKey
}

func MakePublicKeyPrincipal(meta ssh.ConnMetadata, key ssh.PublicKey) PublicKeyPrincipal {
	attr := MakeAttributes()
	attr.Source = meta
	return PublicKeyPrincipal{
		User:      meta.User(),
		PublicKey: key,
	}
}

func (p PublicKeyPrincipal) Type() string           { return TypePublicKey }
func (p PublicKeyPrincipal) Identity() string       { return p.User }
func (p PublicKeyPrincipal) Attributes() Attributes { return MakeAttributes() }
func (p PublicKeyPrincipal) MarshalJSON() ([]byte, error) {
	return json.Marshal(struct {
		Kind string     `json:"type"`
		User string     `json:"identity"`
		Attr Attributes `json:"attr"`
	}{
		p.Type(),
		p.Identity(),
		p.Attributes(),
	})
}

type UserCertificatePrincipal struct {
	*ssh.Certificate
}

func (cert UserCertificatePrincipal) MarshalJSON() ([]byte, error) {
	return json.Marshal(struct {
		Kind string     `json:"type"`
		User string     `json:"identity"`
		Attr Attributes `json:"attr"`
	}{
		cert.Type(),
		cert.Identity(),
		cert.Attributes(),
	})
}

func (cert UserCertificatePrincipal) Type() string     { return TypeCertificate }
func (cert UserCertificatePrincipal) Identity() string { return cert.KeyId }
func (cert UserCertificatePrincipal) Attributes() Attributes {
	return Attributes{
		Principals: cert.ValidPrincipals,
		Options:    cert.CriticalOptions,
		Extensions: cert.Extensions,
		Source:     input.NewCertificate(cert.Certificate),
	}
}