package submission

import (
	"bytes"
	"context"
	"crypto/tls"
	"database/sql"
	"errors"
	"io"
	"log/slog"
	"sync"

	"postern/internal/db"
	"postern/internal/dkim"
	"postern/internal/model"
	"postern/internal/policy"

	"github.com/emersion/go-sasl"
	smtpserver "github.com/emersion/go-smtp"
	"github.com/google/uuid"
)

type Config struct {
	DB                *db.DB
	DKIM              *dkim.DKIM
	Listen            string
	TLS               *tls.Config
	PolicyEngine      *policy.Engine
	IsLocalSubmission bool // If set, will disable TLS and Auth and accept all local mails for delivery.
}
type Server struct {
	db           *db.DB
	dkim         *dkim.DKIM
	listen       string
	tlsConfig    *tls.Config
	policyEngine *policy.Engine

	mu                sync.Mutex
	smtpSrv           *smtpserver.Server
	status            string
	isLocalSubmission bool
}

func NewServer(config *Config) (*Server, error) {
	return &Server{
		db:                config.DB,
		dkim:              config.DKIM,
		listen:            config.Listen,
		tlsConfig:         config.TLS,
		policyEngine:      config.PolicyEngine,
		isLocalSubmission: config.IsLocalSubmission,
	}, nil
}

func (i *Server) Start() error {
	be := &backend{
		db:                i.db,
		dkim:              i.dkim,
		isLocalSubmission: i.isLocalSubmission,
		policyEngine:      i.policyEngine,
	}
	s := smtpserver.NewServer(be)

	s.Addr = i.listen
	s.MaxMessageBytes = 20_000_000

	i.mu.Lock()
	i.smtpSrv = s
	i.status = "running"
	i.mu.Unlock()

	if i.isLocalSubmission {
		err := s.ListenAndServe()
		if err != nil {
			return err
		}
	} else {
		s.TLSConfig = i.tlsConfig
		err := s.ListenAndServeTLS()
		if err != nil {
			return err
		}
	}

	return nil
}

func (i *Server) Stop() error {
	i.mu.Lock()
	defer i.mu.Unlock()
	i.status = "stopped"
	if i.smtpSrv != nil {
		return i.smtpSrv.Close()
	}
	return nil
}

func (i *Server) Status() string {
	i.mu.Lock()
	defer i.mu.Unlock()
	return i.status
}

type smtpSession struct {
	db                *db.DB
	dkim              *dkim.DKIM
	conn              *smtpserver.Conn
	authenticatedUser model.User
	sessionID         string
	mailFrom          string
	rcptTo            []string
	logger            *slog.Logger
	isLocalSubmission bool
	policyEngine      *policy.Engine
}

// The backend implements SMTP server methods.
type backend struct {
	db                *db.DB
	dkim              *dkim.DKIM
	isLocalSubmission bool
	policyEngine      *policy.Engine
}

// NewSession is called after client greeting (EHLO, HELO).
func (bkd *backend) NewSession(c *smtpserver.Conn) (smtpserver.Session, error) {
	sessionID := uuid.NewString()
	l := slog.With("component", "submission", "remote_ip", c.Conn().RemoteAddr().String(), "session_id", sessionID)
	l.Info("new session")
	return &smtpSession{
		db:                bkd.db,
		dkim:              bkd.dkim,
		conn:              c,
		authenticatedUser: model.User{},
		sessionID:         sessionID,
		logger:            l,
		rcptTo:            make([]string, 0),
		isLocalSubmission: bkd.isLocalSubmission,
		policyEngine:      bkd.policyEngine,
	}, nil
}

// AuthMechanisms returns a slice of available auth mechanisms
func (s *smtpSession) AuthMechanisms() []string {
	return []string{sasl.Plain}
}

// Auth is the handler for supported authenticators.
func (s *smtpSession) Auth(_ string) (sasl.Server, error) {
	return sasl.NewPlainServer(func(identity, username, password string) error {
		u, err := s.db.GetUser(context.Background(), username)
		if err != nil {
			if errors.Is(err, sql.ErrNoRows) {
				s.logger.Info("user not found")
			} else {
				s.logger.Error("getting user", "error", err)
			}
			return smtpserver.ErrAuthFailed
		}

		err = u.VerifyPassword([]byte(password))
		if err != nil {
			slog.Info("invalid password", "username", username)
			return smtpserver.ErrAuthFailed
		}

		s.authenticatedUser = u
		s.logger.Info("authenticated", "user_id", u.ID, "username", username)
		return nil
	}), nil
}

func (s *smtpSession) Reset() {
	s.mailFrom = ""
	s.rcptTo = make([]string, 0)
}

func (s *smtpSession) Logout() error {
	s.authenticatedUser = model.User{}
	return nil
}

func (s *smtpSession) Mail(from string, _ *smtpserver.MailOptions) error {
	ctx := context.Background()
	s.logger.InfoContext(ctx, "MAIL FROM", "from", from)

	if !s.isLocalSubmission && s.authenticatedUser.ID == 0 {
		s.logger.InfoContext(ctx, "no authenticated user")
		return smtpserver.ErrAuthFailed
	}

	s.mailFrom = from
	return nil
}

func (s *smtpSession) Rcpt(to string, _ *smtpserver.RcptOptions) error {
	ctx := context.Background()
	s.logger.InfoContext(ctx, "RCPT TO", "to", to)

	if !s.isLocalSubmission && s.authenticatedUser.ID == 0 {
		s.logger.InfoContext(ctx, "no authenticated user")
		return smtpserver.ErrAuthFailed
	}

	s.rcptTo = append(s.rcptTo, to)
	return nil
}

func (s *smtpSession) Data(r io.Reader) error {
	ctx := context.Background()
	s.logger.InfoContext(ctx, "DATA")

	if !s.isLocalSubmission && s.authenticatedUser.ID == 0 {
		s.logger.InfoContext(ctx, "no authenticated user")
		return smtpserver.ErrAuthFailed
	}

	// TODO Check if this mail is for us or for the relay

	var buffer bytes.Buffer
	_, err := io.Copy(&buffer, r)
	if err != nil {
		s.logger.ErrorContext(ctx, "Copy", "error", err)
		return errInternalServerError
	}

	message := buffer.Bytes()

	signed, err := s.dkim.Sign(ctx, message)
	if err != nil {
		s.logger.ErrorContext(ctx, "dkim sign", "error", err)
		return errInternalServerError
	}

	// Call policy hook to determine relay config
	msgCtx := policy.MessageContextFromEnvelope(s.mailFrom, "", int64(len(signed)))
	user := policy.UserToStarlark(s.authenticatedUser, s.db)
	relayConfig, err := s.policyEngine.OnMessageSubmit(user, msgCtx)
	if err != nil {
		s.logger.ErrorContext(ctx, "policy on_message_submit", "error", err)
		return errInternalServerError
	}

	err = shipMessage(s.mailFrom, s.rcptTo, bytes.NewReader(signed), relayConfig)
	if err != nil {
		s.logger.ErrorContext(ctx, "ship message", "error", err)
	}
	return err
}