package blacktip

import (
	"bufio"
	"context"
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/smtp"
	"net/textproto"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

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

// Mail egress: a site's SMTP relay hands this edge one mail transaction at a
// time through the WireGuard tunnel. The edge finds the recipient's mail server
// (MX lookup, or the configured smarthost), connects from the policy's source
// IP, introduces itself with the policy's HELO name, and relays the exchange,
// passing the receiving server's replies straight back so the sender's own
// queue and retry logic keep working.
//
// Guard rails: only the owning site (by tunnel address) may use a policy;
// destinations on internal address ranges are refused so the edge can't be
// aimed at its own network; and a per-policy rate limit caps how much mail a
// compromised site could push out from the IP.

// MailHost is somewhere to deliver mail.
type MailHost struct {
	Host string
	Port int
}

type smtpErr struct {
	Code int
	Text string
}

func (e *smtpErr) Error() string { return strconv.Itoa(e.Code) + " " + e.Text }

func tempErr(text string) *smtpErr { return &smtpErr{451, "4.4.1 " + text} }

// asSMTPErr converts an error from net/smtp into a reply for the sender.
func asSMTPErr(err error) *smtpErr {
	var te *textproto.Error
	if errors.As(err, &te) {
		return &smtpErr{te.Code, te.Msg}
	}
	return tempErr(err.Error())
}

func domainOf(addr string) string {
	if i := strings.LastIndexByte(addr, '@'); i >= 0 {
		return strings.ToLower(strings.Trim(addr[i+1:], "<> "))
	}
	return ""
}

// lookupMail returns where to deliver mail for domain, best first.
func (b *Blacktip) lookupMail(domain string) ([]MailHost, *smtpErr) {
	if b.cfg.MailLookup != nil {
		hosts, err := b.cfg.MailLookup(domain)
		if err != nil {
			return nil, tempErr(err.Error())
		}
		return hosts, nil
	}
	mxs, err := net.LookupMX(domain)
	var dnsErr *net.DNSError
	if err != nil && !(errors.As(err, &dnsErr) && dnsErr.IsNotFound) && len(mxs) == 0 {
		return nil, tempErr("DNS lookup for " + domain + " failed: " + err.Error())
	}
	if len(mxs) == 1 && mxs[0].Host == "." {
		return nil, &smtpErr{556, "5.1.10 " + domain + " does not accept mail"}
	}
	var out []MailHost
	for _, mx := range mxs {
		out = append(out, MailHost{Host: strings.TrimSuffix(mx.Host, "."), Port: 25})
	}
	if len(out) == 0 { // no MX records: the domain's own address is the mail server (RFC 5321)
		if ips, err := net.LookupHost(domain); err != nil || len(ips) == 0 {
			return nil, &smtpErr{550, "5.1.2 domain " + domain + " not found"}
		}
		out = append(out, MailHost{Host: domain, Port: 25})
	}
	return out, nil
}

// mailRateOK counts a new transaction against the policy's per-minute limit.
func (b *Blacktip) mailRateOK(pol model.MailPolicy) bool {
	limit := pol.RatePerMinute
	if limit <= 0 {
		limit = 60
	}
	b.mailMu.Lock()
	defer b.mailMu.Unlock()
	now := time.Now()
	kept := b.mailRate[pol.ID][:0]
	for _, t := range b.mailRate[pol.ID] {
		if now.Sub(t) < time.Minute {
			kept = append(kept, t)
		}
	}
	if len(kept) >= limit {
		b.mailRate[pol.ID] = kept
		return false
	}
	b.mailRate[pol.ID] = append(kept, now)
	return true
}

type mailSession struct {
	b      *Blacktip
	c      net.Conn
	pol    model.MailPolicy
	peer   string
	cl     *smtp.Client
	from   string
	domain string
	mx     string
	rcpts  int
	dw     io.WriteCloser
	start  time.Time
	sts    bool   // the recipient's MTA-STS policy was enforced
	note   string // what was done to the message (signed...), for the log
}

func (b *Blacktip) serveMail(ln net.Listener) {
	for {
		c, err := ln.Accept()
		if err != nil {
			return
		}
		go b.handleMail(c)
	}
}

func (s *mailSession) close() {
	if s.cl != nil {
		s.cl.Close()
		s.cl = nil
	}
	s.dw, s.rcpts, s.domain = nil, 0, ""
}

func (s *mailSession) reply(code int, text string) {
	b, _ := json.Marshal(tunnel.MailReply{Code: code, Text: text})
	s.c.Write(append(b, '\n'))
}

func (s *mailSession) replyErr(e *smtpErr) { s.reply(e.Code, e.Text) }

func (b *Blacktip) handleMail(c net.Conn) {
	s := &mailSession{b: b, c: c, peer: hostFromAddr(c.RemoteAddr().String()), start: time.Now()}
	defer func() { s.close(); c.Close() }()
	br := bufio.NewReader(c)
	authed := false
	for {
		c.SetReadDeadline(time.Now().Add(5 * time.Minute))
		line, err := br.ReadBytes('\n')
		if err != nil {
			return
		}
		var m tunnel.MailMsg
		if json.Unmarshal(line, &m) != nil {
			return
		}
		if !authed {
			b.mu.Lock()
			pol, ok := b.mail[m.Egress]
			peer := b.peerIP[pol.SiteID]
			b.mu.Unlock()
			if !ok || peer == "" || peer != s.peer {
				s.reply(554, "5.7.1 not allowed: this site may not use that mail egress")
				return
			}
			s.pol, authed = pol, true
		}
		switch m.Op {
		case "rcpt":
			s.rcpt(m)
		case "data":
			if s.cl == nil || s.rcpts == 0 {
				s.reply(503, "5.5.1 no recipients")
				continue
			}
			w, err := s.cl.Data()
			if err != nil {
				s.replyErr(asSMTPErr(err))
				s.close()
				continue
			}
			s.dw = w
			s.reply(354, "start mail input")
		case "body":
			if !s.body(br, m.Len) {
				return
			}
		case "rset", "quit":
			if s.cl != nil {
				s.cl.Quit()
			}
			s.close()
			s.reply(250, "2.0.0 ok")
			if m.Op == "quit" {
				return
			}
		default:
			s.reply(500, "5.5.1 unknown request")
		}
	}
}

// connect opens a connection to a mail server for domain and gets as far as MAIL FROM.
func (s *mailSession) connect(domain, from string) (*smtp.Client, string, *smtpErr) {
	hosts := []MailHost(nil)
	if s.pol.Smarthost != "" {
		h, p, _ := net.SplitHostPort(s.pol.Smarthost)
		port, _ := strconv.Atoi(p)
		hosts = []MailHost{{Host: h, Port: port}}
	} else {
		var e *smtpErr
		if hosts, e = s.b.lookupMail(domain); e != nil {
			return nil, "", e
		}
	}
	// a recipient domain's MTA-STS policy can forbid some servers and demand verified TLS
	requireTLS, why := s.pol.RequireTLS, "this egress requires it"
	if s.pol.Smarthost == "" {
		if sts := s.b.stsFor(domain); sts != nil && sts.Mode == "enforce" {
			var allowed []MailHost
			for _, h := range hosts {
				if sts.Allows(h.Host) {
					allowed = append(allowed, h)
				}
			}
			if len(allowed) == 0 {
				return nil, "", &smtpErr{451, "4.7.5 " + domain + "'s MTA-STS policy doesn't allow delivery to any of its listed mail servers"}
			}
			hosts, requireTLS, why = allowed, true, "the recipient's MTA-STS policy requires it"
			s.sts = true
		}
	}
	src := net.ParseIP(s.pol.SourceIP)
	last := tempErr("could not reach any mail server for " + domain)
	for _, h := range hosts {
		var ips []net.IP
		if ip := net.ParseIP(h.Host); ip != nil {
			ips = []net.IP{ip}
		} else {
			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			res, err := net.DefaultResolver.LookupIPAddr(ctx, h.Host)
			cancel()
			if err != nil {
				last = tempErr("cannot resolve " + h.Host)
				continue
			}
			for _, r := range res {
				ips = append(ips, r.IP)
			}
		}
		for _, ip := range ips {
			if (src.To4() != nil) != (ip.To4() != nil) {
				continue // can't connect to an IPv6 server from an IPv4 source or the reverse
			}
			if !s.b.cfg.MailAllowPrivate && netguard.Blocked(ip) {
				last = &smtpErr{554, "5.7.1 not allowed: " + h.Host + " is an internal address"}
				continue
			}
			d := net.Dialer{Timeout: 15 * time.Second, LocalAddr: &net.TCPAddr{IP: src}}
			conn, err := d.Dial("tcp", net.JoinHostPort(ip.String(), strconv.Itoa(h.Port)))
			if err != nil {
				last = tempErr("cannot connect to " + h.Host + ": " + err.Error())
				continue
			}
			conn.SetDeadline(time.Now().Add(3 * time.Minute))
			cl, err := smtp.NewClient(conn, h.Host)
			if err != nil {
				conn.Close()
				last = asSMTPErr(err)
				continue
			}
			if err := cl.Hello(s.pol.HeloName); err != nil {
				cl.Close()
				last = asSMTPErr(err)
				continue
			}
			if ok, _ := cl.Extension("STARTTLS"); ok {
				if err := cl.StartTLS(&tls.Config{ServerName: h.Host, InsecureSkipVerify: !requireTLS, RootCAs: s.b.cfg.MailRootCAs}); err != nil {
					cl.Close()
					last = tempErr("TLS to " + h.Host + " failed: " + err.Error())
					continue
				}
			} else if requireTLS {
				cl.Close()
				last = tempErr(h.Host + " does not offer TLS and " + why)
				continue
			}
			if err := cl.Mail(from); err != nil {
				cl.Close()
				if e := asSMTPErr(err); e.Code >= 500 {
					return nil, "", e // the server refused this sender for good: another MX won't differ
				} else {
					last = e
				}
				continue
			}
			return cl, h.Host, nil
		}
	}
	return nil, "", last
}

func (s *mailSession) rcpt(m tunnel.MailMsg) {
	if s.cl == nil {
		if !s.b.mailRateOK(s.pol) {
			s.reply(451, "4.7.1 sending rate limit reached for this egress, try again shortly")
			return
		}
		domain := domainOf(m.To)
		if domain == "" {
			s.reply(501, "5.1.3 bad recipient address")
			return
		}
		cl, mx, e := s.connect(domain, m.From)
		if e != nil {
			s.logEntry(0, e.Code, "mail-failed: "+e.Text, domain, "")
			s.replyErr(e)
			return
		}
		s.cl, s.mx, s.domain, s.from = cl, mx, domain, m.From
	} else if domainOf(m.To) != s.domain {
		s.reply(452, "4.5.3 one destination domain per transaction")
		return
	}
	if err := s.cl.Rcpt(m.To); err != nil {
		e := asSMTPErr(err)
		s.b.recordOutcome(s.pol.ID, s.domain, e.Code, e.Text)
		if s.rcpts == 0 {
			s.close()
		}
		s.replyErr(e)
		return
	}
	s.rcpts++
	s.reply(250, "2.1.5 ok")
}

func (s *mailSession) body(br *bufio.Reader, n int) bool {
	if n < 0 || n > tunnel.MaxMessage || s.dw == nil || s.cl == nil {
		s.reply(552, "5.3.4 message size not accepted")
		return false
	}
	s.c.SetReadDeadline(time.Now().Add(5 * time.Minute))
	data := make([]byte, n)
	if _, err := io.ReadFull(br, data); err != nil {
		s.close()
		return false
	}
	data, s.note = s.sign(data)
	if _, err := s.dw.Write(data); err != nil {
		e := asSMTPErr(err)
		s.logEntry(n, e.Code, "mail-failed: "+e.Text, s.domain, s.mx)
		s.replyErr(e)
		s.close()
		return true
	}
	cerr := s.dw.Close()
	s.dw = nil
	domain, mx := s.domain, s.mx
	if cerr != nil {
		e := asSMTPErr(cerr)
		s.logEntry(n, e.Code, "mail-failed: "+e.Text, domain, mx)
		s.replyErr(e)
		s.close()
		return true
	}
	s.cl.Quit()
	s.logEntry(n, 250, "mail", domain, mx)
	s.reply(250, "2.0.0 ok: delivered to "+mx)
	s.close()
	return true
}

func (s *mailSession) logEntry(size, code int, result, domain, mx string) {
	path := "via " + mx
	if s.sts {
		path += " (MTA-STS)"
	}
	if s.note != "" {
		path += " · " + s.note
	}
	s.b.recordOutcome(s.pol.ID, domain, code, strings.TrimPrefix(result, "mail-failed: "))
	s.b.logAccess(model.LogEntry{Time: s.start, Mapping: s.pol.ID, Remote: s.peer, Method: "MAIL", Host: domain, Path: path,
		Status: code, Result: result, BytesIn: int64(size), Millis: time.Since(s.start).Milliseconds()})
}

// ---------- deliverability checks ----------

const defaultMailProbe = "gmail-smtp-in.l.google.com:25"

// mailCheck inspects how receiving servers will see mail from the policy's source IP.
func (b *Blacktip) mailCheck(pol model.MailPolicy) model.MailCheck {
	chk := model.MailCheck{CheckedAt: time.Now(), PTR: []string{}}
	src := net.ParseIP(pol.SourceIP)
	if names, err := net.LookupAddr(pol.SourceIP); err == nil {
		for _, n := range names {
			n = strings.ToLower(strings.TrimSuffix(n, "."))
			chk.PTR = append(chk.PTR, n)
			if n == strings.ToLower(pol.HeloName) {
				chk.PTRMatchesHelo = true
			}
		}
	}
	if ips, err := net.LookupHost(pol.HeloName); err == nil {
		for _, ip := range ips {
			if p := net.ParseIP(ip); p != nil && p.Equal(src) {
				chk.HeloResolves = true
			}
		}
	}
	if len(pol.DKIM) > 0 {
		chk.DKIM = b.checkDKIM(pol.DKIM)
	}
	chk.Domains = b.checkDomains(pol)
	chk.Blocklists, chk.Listable = b.checkBlocklists(pol.SourceIP)
	probe := b.cfg.MailProbe
	if probe == "" {
		probe = defaultMailProbe
	}
	d := net.Dialer{Timeout: 8 * time.Second, LocalAddr: &net.TCPAddr{IP: src}}
	c, err := d.Dial("tcp", probe)
	if err != nil {
		chk.Port25 = "can't connect out on port 25 from this address (" + err.Error() + "); many hosts block it, so ask your provider to open it"
		return chk
	}
	defer c.Close()
	c.SetReadDeadline(time.Now().Add(8 * time.Second))
	banner, err := bufio.NewReader(c).ReadString('\n')
	if err != nil || !strings.HasPrefix(banner, "220") {
		chk.Port25 = "connected, but the server didn't greet us normally"
		return chk
	}
	chk.Port25 = "open"
	return chk
}

// runMailChecks refreshes the report for every current policy.
func (b *Blacktip) runMailChecks() {
	b.mu.Lock()
	pols := make([]model.MailPolicy, 0, len(b.mail))
	for _, p := range b.mail {
		pols = append(pols, p)
	}
	b.mu.Unlock()
	sort.Slice(pols, func(i, j int) bool { return pols[i].ID < pols[j].ID })
	res := map[string]model.MailCheck{}
	var wg sync.WaitGroup
	var mu sync.Mutex
	for _, p := range pols {
		wg.Add(1)
		go func(p model.MailPolicy) {
			defer wg.Done()
			c := b.mailCheck(p)
			mu.Lock()
			res[p.ID] = c
			mu.Unlock()
		}(p)
	}
	wg.Wait()
	b.mu.Lock()
	b.mailChecks = res
	b.mu.Unlock()
	b.report()
}

func (b *Blacktip) mailCheckLoop(ctx context.Context) {
	t := time.NewTicker(15 * time.Minute)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			b.runMailChecks()
		}
	}
}

func mailPolicySig(m map[string]model.MailPolicy) string {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	var sb strings.Builder
	for _, k := range keys {
		p := m[k]
		fmt.Fprintf(&sb, "%s|%s|%s|%s;", p.ID, p.SourceIP, p.HeloName, p.Smarthost)
		for _, d := range p.DKIM {
			fmt.Fprintf(&sb, "%s/%s/%s,", d.KeyID, d.Selector, d.Domain)
		}
	}
	return sb.String()
}