package auth

import (
	"bytes"
	"errors"
	"net"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"

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

// None accepts no authentication, this is only useful for debugging/testing.
type None struct{}

func (None) HandleAccept(conn net.Conn, serverConfig *ssh.ServerConfig) (*ssh.ServerConn, <-chan ssh.NewChannel, <-chan *ssh.Request, error) {
	config := new(ssh.ServerConfig)
	if serverConfig != nil {
		*config = *serverConfig
	}
	config.NoClientAuth = true
	return ssh.NewServerConn(conn, config)
}

// UserCertificate offers user certificate based authentication.
type UserCertificate struct {
	// Loader is the policy loader.
	Loader policy.Loader

	// CA is our trusted user certificate authority public keys.
	CA []ssh.PublicKey

	// VerifyCallback can optionally be used to perform additional checks on the certificate.
	VerifyCallback func(*ssh.Certificate) bool
}

func (auth UserCertificate) HandleAccept(conn net.Conn, serverConfig *ssh.ServerConfig) (*ssh.ServerConn, <-chan ssh.NewChannel, <-chan *ssh.Request, error) {
	log := logger.StandardLog.Value("client", conn.RemoteAddr().String())

	policy, err := auth.Loader.LoadPolicy("conduit/auth/user_certificate")
	if err != nil {
		return nil, nil, nil, err
	}
	_ = policy

	checker := &ssh.CertChecker{
		Clock: time.Now,
	}
	checker.IsUserAuthority = func(key ssh.PublicKey) bool {
		blob := key.Marshal()
		for _, ca := range auth.CA {
			if bytes.Equal(ca.Marshal(), blob) {
				return true
			}
		}
		return false
	}

	config := new(ssh.ServerConfig)
	if serverConfig != nil {
		*config = *serverConfig
	}
	config.PublicKeyCallback = func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
		if _, err := checker.Authenticate(conn, key); err != nil {
			log.Err(err).Debug("Certificate check failed")
			return nil, err
		}

		cert := key.(*ssh.Certificate)
		principal := UserCertificatePrincipal{Certificate: cert}
		var (
			query = "data." + policy.Package()
			input = struct {
				Conn      *input.ConnMetadata `json:"conn"`
				Principal Principal           `json:"principal"`
			}{input.NewConnMetadata(conn), principal}
			result struct {
				Permit bool `mapstructure:"permit"`
			}
		)
		if err = policy.Query(query, input, &result); err != nil {
			log.Err(err).Warn("Policy query returned error")
			return nil, err
		}
		if !result.Permit {
			log.Debug("Policy rejected certificate")
			return nil, ErrUnauthorized
		}
		return &cert.Permissions, nil
	}

	return ssh.NewServerConn(conn, config)
}

// MultiFactor offers policy-based MFA (multi factor authentication).
type MultiFactor struct {
	// Loader for the policy rules.
	Loader policy.Loader

	// PublicKey enables public key authentication.
	PublicKey PublicKey

	// UserCA enables user certificate based authentication.
	UserCA []ssh.PublicKey

	// Password enables password authentication.
	Password Password

	// Token enabled token authentication.
	Token Token
}

func (mfa *MultiFactor) HandleAccept(conn net.Conn, serverConfig *ssh.ServerConfig) (*ssh.ServerConn, <-chan ssh.NewChannel, <-chan *ssh.Request, error) {
	log := logger.StandardLog.Value("client", conn.RemoteAddr().String())

	defer func() {
		if r := recover(); r != nil {
			if err, ok := r.(error); ok {
				log.Err(err).Error("Recovered from panic in accept handler! This is a bug!")
			}
		}
	}()

	rules, err := mfa.Loader.LoadPolicy("conduit/auth/mfa")
	if err != nil {
		return nil, nil, nil, err
	}

	config := new(ssh.ServerConfig)
	if serverConfig != nil {
		*config = *serverConfig
	}

	var (
		meta            ssh.ConnMetadata
		principals      []Principal
		tokenAuth       = mfa.Token
		passwordAuth    = mfa.Password
		publicKeyAuth   = mfa.PublicKey
		certificateAuth []ssh.PublicKey
	)

	checkPolicy := func() (perms *ssh.Permissions, err error) {
		var (
			query = "data." + rules.Package()
			input = struct {
				Conn       *input.ConnMetadata `json:"conn"`
				Principals []Principal         `json:"principals"`
			}{input.NewConnMetadata(meta), principals}
			result struct {
				Permit            bool `mapstructure:"permit"`
				PermitPassword    bool `mapstructure:"permit_password"`
				PermitToken       bool `mapstructure:"permit_token"`
				PermitCertificate bool `mapstructure:"permit_certificate"`
				PermitPublicKey   bool `mapstructure:"permit_publickey"`
			}
		)
		if err = rules.Query(query, input, &result); err != nil {
			log.Err(err).Warn("Policy query returned error")
			return
		}

		if !result.PermitPassword && passwordAuth != nil {
			log.Debug("Policy disabled password authentication")
			passwordAuth = nil
		}
		if passwordAuth == nil && config.PasswordCallback != nil {
			// Disable password authentication.
			log.Debug("Policy disabled password callback")
			config.PasswordCallback = nil
		}

		if !result.PermitToken && tokenAuth != nil {
			log.Debug("Policy disabled token authentication")
			tokenAuth = nil
		}
		if tokenAuth == nil && config.KeyboardInteractiveCallback != nil {
			// Disable token authentication.
			log.Debug("Policy disabled token authentication")
			config.KeyboardInteractiveCallback = nil
		}

		if !result.PermitCertificate && len(certificateAuth) != 0 {
			log.Debug("Policy disabled certificate authentication")
			certificateAuth = nil
		}
		if !result.PermitPublicKey && publicKeyAuth != nil {
			log.Debug("Policy disabled public key authentication")
			publicKeyAuth = nil
		}
		if publicKeyAuth == nil && len(certificateAuth) == 0 && config.PublicKeyCallback != nil {
			// Disable pubkey authentication.
			log.Debug("Policy disabled public key callback")
			config.PublicKeyCallback = nil
		}

		if !result.Permit {
			log.Debug("Policy rejected principals; not authorized")
			return nil, ErrUnauthorized
		}

		if len(principals) == 0 {
			log.Debug("No valid principals; not authorized")
			return nil, ErrUnauthorized
		}

		attr := principals[0].Attributes()
		return &ssh.Permissions{
			CriticalOptions: attr.Options,
			Extensions:      attr.Extensions,
		}, nil
	}

	if mfa.PublicKey != nil || len(mfa.UserCA) > 0 {
		log = log.Values(logger.Values{
			"permit_pubkey":      mfa.PublicKey != nil,
			"permit_certificate": len(mfa.UserCA) > 0,
		})
		log.Trace("MFA enabling public key authentication")

		var checker *ssh.CertChecker
		if len(mfa.UserCA) > 0 {
			checker = &ssh.CertChecker{
				Clock: time.Now,
			}
			checker.IsUserAuthority = func(key ssh.PublicKey) bool {
				blob := key.Marshal()
				for _, ca := range mfa.UserCA {
					if bytes.Equal(ca.Marshal(), blob) {
						return true
					}
				}
				return false
			}
		}

		config.PublicKeyCallback = func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			var (
				principal Principal
				errs      []error
			)

			if checker != nil {
				if _, err := checker.Authenticate(conn, key); err != nil {
					log.Err(err).Debug("Certificate check failed")
					errs = append(errs, err)
				} else {
					principal = UserCertificatePrincipal{Certificate: key.(*ssh.Certificate)}
				}
			}
			if principal == nil && mfa.PublicKey != nil {
				if _, err := mfa.PublicKey.VerifyPublicKey(conn, key); err != nil {
					log.Err(err).Debug("Public key check failed")
					errs = append(errs, err)
				} else {
					principal = PublicKeyPrincipal{User: conn.User(), PublicKey: key}
				}
			}

			if principal == nil {
				if len(errs) > 0 {
					return nil, errors.Join(errs...)
				}
				return nil, ErrUnauthorized
			}

			meta = conn
			principals = append(principals, principal)
			return checkPolicy()
		}
	}

	if mfa.Password != nil {
		log.Trace("MFA enabling password authentication")

		config.PasswordCallback = func(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
			if _, err := mfa.Password.VerifyPassword(conn, string(password)); err != nil {
				return nil, err
			}

			meta = conn
			principals = append(principals, PasswordPrincipal{User: conn.User()})
			return checkPolicy()
		}
	}

	if mfa.Token != nil {
		log.Trace("MFA enabling token authentication")

		config.KeyboardInteractiveCallback = func(conn ssh.ConnMetadata, challenge ssh.KeyboardInteractiveChallenge) (*ssh.Permissions, error) {
			var (
				multiline = mfa.Token.Multiline()
				token     string
			)

			for {
				answers, err := challenge("", mfa.Token.Instruction(), []string{"Token: "}, []bool{false})
				if err != nil {
					return nil, err
				} else if len(answers) != 1 {
					return nil, ErrUnauthorized
				}
				if multiline {
					if answers[0] == "" {
						break
					}
					token += strings.TrimSpace(answers[0])
				} else {
					token = answers[0]
					break
				}
			}

			if _, err := mfa.Token.VerifyToken(conn, token); err != nil {
				return nil, err
			}

			meta = conn
			principals = append(principals, TokenPrincipal{User: conn.User()})
			return checkPolicy()
		}
	}

	return ssh.NewServerConn(conn, config)
}