package bamboo

import (
	"bufio"
	"context"
	"crypto/subtle"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"net/textproto"
	"strconv"
	"strings"
	"time"

	"git.veillax.com/Veillax/reef/model"
	"git.veillax.com/Veillax/reef/tunnel"
)

// Mail egress: a small SMTP relay on this machine. Point your mail server's
// relayhost at it (e.g. Postfix: relayhost = [127.0.0.1]:2525). Each message is
// handed through the WireGuard tunnel to the edge, which delivers it to the
// recipient's server from the edge's public IP and passes the replies back.
//
// One transaction is relayed to one destination domain. If a message has
// recipients in several domains, the extra ones get "452 too many recipients"
// and your MTA resends them in a separate transaction, which is standard behavior.

type mailRun struct {
	spec model.MailSpec
	ln   net.Listener
}

const maxRecipients = 100

// syncMail starts, restarts and stops relays to match specs.
func (b *Bamboo) syncMail(specs []model.MailSpec) {
	b.mu.Lock()
	defer b.mu.Unlock()
	want := map[string]model.MailSpec{}
	for _, sp := range specs {
		want[sp.ID] = sp
	}
	for id, run := range b.mail {
		if sp, ok := want[id]; !ok || sp != run.spec {
			run.ln.Close()
			delete(b.mail, id)
			delete(b.mailErr, id)
		}
	}
	for id, sp := range want {
		if _, running := b.mail[id]; running {
			continue
		}
		ln, err := net.Listen("tcp", sp.Listen)
		if err != nil {
			b.mailErr[id] = err.Error()
			log.Printf("mail egress %s: %v", sp.Listen, err)
			continue
		}
		run := &mailRun{spec: sp, ln: ln}
		b.mail[id] = run
		b.mailErr[id] = ""
		log.Printf("mail egress relay listening on %s (via edge %s)", sp.Listen, sp.EdgeID)
		go b.serveMail(run)
	}
}

func (b *Bamboo) mailStatus() map[string]string {
	b.mu.Lock()
	defer b.mu.Unlock()
	out := map[string]string{}
	for id, e := range b.mailErr {
		out[id] = e
	}
	return out
}

func (b *Bamboo) serveMail(r *mailRun) {
	for {
		c, err := r.ln.Accept()
		if err != nil {
			return
		}
		go b.handleSMTP(c, r.spec)
	}
}

type smtpSession struct {
	b      *Bamboo
	c      net.Conn
	tp     *textproto.Conn
	spec   model.MailSpec
	authed bool

	hasFrom bool
	from    string
	domain  string
	rcpts   int

	stream net.Conn
	sr     *bufio.Reader
}

func (s *smtpSession) say(code int, text string) {
	text = strings.NewReplacer("\r", " ", "\n", " ").Replace(text)
	fmt.Fprintf(s.c, "%d %s\r\n", code, text)
}

func (s *smtpSession) sayMulti(code int, lines ...string) {
	for i, l := range lines {
		sep := "-"
		if i == len(lines)-1 {
			sep = " "
		}
		fmt.Fprintf(s.c, "%d%s%s\r\n", code, sep, l)
	}
}

// openStream connects to the edge's mail port through the tunnel.
func (s *smtpSession) openStream() error {
	b := s.b
	b.mu.Lock()
	tun, edgeIP := b.tun, ""
	for _, e := range b.edges {
		if e.ID == s.spec.EdgeID {
			edgeIP = e.IP
		}
	}
	b.mu.Unlock()
	if tun == nil || edgeIP == "" {
		return errors.New("the tunnel to the edge isn't up")
	}
	ctx, cancel := context.WithTimeout(b.ctx, 12*time.Second)
	defer cancel()
	c, err := tun.DialContext(ctx, "tcp", net.JoinHostPort(edgeIP, strconv.Itoa(tunnel.MailTCPPort)))
	if err != nil {
		return errors.New("can't reach the edge through the tunnel")
	}
	s.stream, s.sr = c, bufio.NewReader(c)
	return nil
}

func (s *smtpSession) closeStream() {
	if s.stream != nil {
		s.stream.Close()
		s.stream, s.sr = nil, nil
	}
}

func (s *smtpSession) resetTxn() {
	s.closeStream()
	s.hasFrom, s.from, s.domain, s.rcpts = false, "", "", 0
}

// ask sends one request to the edge and reads its reply.
func (s *smtpSession) ask(m tunnel.MailMsg, body []byte) (tunnel.MailReply, error) {
	var rep tunnel.MailReply
	line, _ := json.Marshal(m)
	s.stream.SetDeadline(time.Now().Add(4 * time.Minute))
	if _, err := s.stream.Write(append(line, '\n')); err != nil {
		return rep, err
	}
	if body != nil {
		if _, err := s.stream.Write(body); err != nil {
			return rep, err
		}
	}
	resp, err := s.sr.ReadBytes('\n')
	if err != nil {
		return rep, err
	}
	return rep, json.Unmarshal(resp, &rep)
}

// angle extracts the address from "FROM:<a@b> SIZE=1" or "TO:<a@b>".
func angle(arg, key string) (addr, params string, ok bool) {
	if len(arg) < len(key) || !strings.EqualFold(arg[:len(key)], key) {
		return "", "", false
	}
	rest := strings.TrimSpace(arg[len(key):])
	if !strings.HasPrefix(rest, "<") {
		return "", "", false
	}
	end := strings.IndexByte(rest, '>')
	if end < 0 {
		return "", "", false
	}
	return rest[1:end], strings.TrimSpace(rest[end+1:]), true
}

func (b *Bamboo) handleSMTP(c net.Conn, spec model.MailSpec) {
	s := &smtpSession{b: b, c: c, tp: textproto.NewConn(c), spec: spec, authed: spec.Username == ""}
	defer func() {
		recover()
		s.closeStream()
		c.Close()
	}()
	s.say(220, "mail-egress ESMTP ready")
	badAuth := 0
	for {
		c.SetReadDeadline(time.Now().Add(5 * time.Minute))
		line, err := s.tp.ReadLine()
		if err != nil {
			return
		}
		cmd, arg, _ := strings.Cut(line, " ")
		arg = strings.TrimSpace(arg)
		switch strings.ToUpper(cmd) {
		case "EHLO", "HELO":
			lines := []string{"mail-egress at your service", "SIZE " + strconv.Itoa(tunnel.MaxMessage), "8BITMIME", "SMTPUTF8", "ENHANCEDSTATUSCODES"}
			if spec.Username != "" {
				lines = append(lines, "AUTH PLAIN LOGIN")
			}
			if strings.ToUpper(cmd) == "HELO" {
				s.say(250, "mail-egress")
			} else {
				s.sayMulti(250, lines...)
			}
		case "AUTH":
			if spec.Username == "" {
				s.say(503, "5.5.1 authentication isn't needed here")
				continue
			}
			if !s.auth(arg) {
				s.say(535, "5.7.8 authentication failed")
				if badAuth++; badAuth >= 3 {
					return
				}
				continue
			}
			s.authed = true
			s.say(235, "2.7.0 authenticated")
		case "MAIL":
			if !s.authed {
				s.say(530, "5.7.0 authentication required")
				continue
			}
			if s.hasFrom {
				s.say(503, "5.5.1 a transaction is already in progress")
				continue
			}
			addr, params, ok := angle(arg, "FROM:")
			if !ok {
				s.say(501, "5.1.7 syntax: MAIL FROM:<address>")
				continue
			}
			if strings.Contains(strings.ToUpper(params), "SIZE=") {
				for _, p := range strings.Fields(params) {
					if strings.HasPrefix(strings.ToUpper(p), "SIZE=") {
						if n, _ := strconv.Atoi(p[5:]); n > tunnel.MaxMessage {
							s.say(552, "5.3.4 message too large")
							addr = "\x00"
						}
					}
				}
				if addr == "\x00" {
					continue
				}
			}
			s.hasFrom, s.from = true, addr
			s.say(250, "2.1.0 ok")
		case "RCPT":
			if !s.hasFrom {
				s.say(503, "5.5.1 send MAIL FROM first")
				continue
			}
			addr, _, ok := angle(arg, "TO:")
			if !ok || addr == "" {
				s.say(501, "5.1.3 syntax: RCPT TO:<address>")
				continue
			}
			s.rcpt(addr)
		case "DATA":
			s.data()
		case "RSET":
			if s.stream != nil {
				s.ask(tunnel.MailMsg{Op: "rset"}, nil)
			}
			s.resetTxn()
			s.say(250, "2.0.0 ok")
		case "NOOP":
			s.say(250, "2.0.0 ok")
		case "QUIT":
			s.say(221, "2.0.0 bye")
			return
		case "VRFY", "EXPN":
			s.say(252, "2.5.2 cannot verify, but will accept and attempt delivery")
		case "STARTTLS":
			s.say(502, "5.5.1 TLS isn't offered: this relay is meant to be reached on this machine")
		default:
			s.say(500, "5.5.2 command not recognized")
		}
	}
}

func (s *smtpSession) auth(arg string) bool {
	mech, initial, _ := strings.Cut(arg, " ")
	read := func() string { // one base64 line from the client
		l, err := s.tp.ReadLine()
		if err != nil {
			return ""
		}
		b, _ := base64.StdEncoding.DecodeString(l)
		return string(b)
	}
	var user, pass string
	switch strings.ToUpper(mech) {
	case "PLAIN":
		raw := ""
		if initial != "" {
			b, _ := base64.StdEncoding.DecodeString(initial)
			raw = string(b)
		} else {
			s.say(334, "")
			raw = read()
		}
		parts := strings.Split(raw, "\x00") // authzid \0 user \0 password
		if len(parts) != 3 {
			return false
		}
		user, pass = parts[1], parts[2]
	case "LOGIN":
		s.say(334, base64.StdEncoding.EncodeToString([]byte("Username:")))
		user = read()
		s.say(334, base64.StdEncoding.EncodeToString([]byte("Password:")))
		pass = read()
	default:
		return false
	}
	okU := subtle.ConstantTimeCompare([]byte(user), []byte(s.spec.Username)) == 1
	okP := subtle.ConstantTimeCompare([]byte(pass), []byte(s.spec.Password)) == 1
	return okU && okP
}

func (s *smtpSession) rcpt(addr string) {
	at := strings.LastIndexByte(addr, '@')
	if at < 1 || at == len(addr)-1 {
		s.say(501, "5.1.3 bad recipient address")
		return
	}
	domain := strings.ToLower(addr[at+1:])
	if s.rcpts > 0 && domain != s.domain {
		s.say(452, "4.5.3 one destination domain per transaction; send the others separately")
		return
	}
	if s.rcpts >= maxRecipients {
		s.say(452, "4.5.3 too many recipients")
		return
	}
	first := s.stream == nil
	if first {
		if err := s.openStream(); err != nil {
			s.say(451, "4.4.1 "+err.Error())
			return
		}
	}
	rep, err := s.ask(tunnel.MailMsg{Op: "rcpt", Egress: s.spec.ID, From: s.from, To: addr}, nil)
	if err != nil {
		s.closeStream()
		s.say(451, "4.4.2 lost contact with the edge")
		return
	}
	if rep.Code != 250 {
		if first {
			s.closeStream()
		}
		s.say(rep.Code, rep.Text)
		return
	}
	s.rcpts++
	s.domain = domain
	s.say(250, "2.1.5 ok")
}

func (s *smtpSession) data() {
	if s.rcpts == 0 || s.stream == nil {
		s.say(503, "5.5.1 need a recipient first")
		return
	}
	rep, err := s.ask(tunnel.MailMsg{Op: "data"}, nil)
	if err != nil {
		s.resetTxn()
		s.say(451, "4.4.2 lost contact with the edge")
		return
	}
	if rep.Code != 354 {
		s.say(rep.Code, rep.Text)
		s.resetTxn()
		return
	}
	s.say(354, "end data with <CR><LF>.<CR><LF>")
	s.c.SetReadDeadline(time.Now().Add(10 * time.Minute))
	dr := s.tp.DotReader()
	body, err := io.ReadAll(io.LimitReader(dr, tunnel.MaxMessage+1))
	if err != nil {
		s.resetTxn()
		return
	}
	if len(body) > tunnel.MaxMessage {
		io.Copy(io.Discard, dr)
		s.ask(tunnel.MailMsg{Op: "rset"}, nil)
		s.resetTxn()
		s.say(552, "5.3.4 message too large")
		return
	}
	rep, err = s.ask(tunnel.MailMsg{Op: "body", Len: len(body)}, body)
	s.resetTxn()
	if err != nil {
		// we can't know whether the far side accepted it; a temporary failure
		// makes the MTA retry, which risks a duplicate but never loses mail
		s.say(451, "4.4.2 lost contact with the edge while sending")
		return
	}
	s.say(rep.Code, rep.Text)
}