package smtpserver

import (
	"fmt"
	"io"
	"net"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/emersion/go-smtp"
	"gopkg.in/ini.v1"
	"mailgoserver/internal/db"
	"mailgoserver/internal/dkim"
	"mailgoserver/internal/relay"
	"mailgoserver/internal/toolbox"
)

// Backend holds the shared dependencies every connection's Session uses, mirroring the
// constructor args threaded through smtp_handler.EnhancedCustomSMTPHandler /
// email_server/server_runner.py.
type Backend struct {
	DB                  *db.DB
	DKIM                *dkim.Manager
	Relay               *relay.Relay
	Cfg                 *ini.File
	Logger              *toolbox.Logger
	HeloHostname        string
	AttachmentsBasePath string
}

func (b *Backend) NewSession(c *smtp.Conn) (smtp.Session, error) {
	host, _, _ := net.SplitHostPort(c.Conn().RemoteAddr().String())
	if host == "" {
		host = c.Conn().RemoteAddr().String()
	}
	return &Session{backend: b, conn: c, peerIP: host}, nil
}

// Session implements smtp.Session + smtp.AuthSession for one SMTP connection, mirroring
// EnhancedCustomSMTPHandler's per-connection behavior in smtp_handler.py.
type Session struct {
	backend *Backend
	conn    *smtp.Conn
	peerIP  string

	authenticatedSender *db.Sender
	authType            string // "sender" | "ip" | ""
	authorizedDomain    string
	username            string

	mailFrom string
	rcptTos  []string
}

func (s *Session) Reset() {
	s.mailFrom = ""
	s.rcptTos = nil
}

func (s *Session) Logout() error { return nil }

// Mail mirrors EnhancedCustomSMTPHandler.handle_MAIL, delegating authorization to
// validateSenderAuthorization (== auth.validate_sender_authorization).
func (s *Session) Mail(from string, opts *smtp.MailOptions) error {
	ok, message := s.validateSenderAuthorization(from)
	if !ok {
		return &smtp.SMTPError{Code: 550, EnhancedCode: smtp.NoEnhancedCode, Message: message}
	}
	s.mailFrom = from
	return nil
}

// validateSenderAuthorization mirrors auth.validate_sender_authorization exactly,
// including its two branches (already-authenticated sender vs. IP whitelist fallback)
// and the AuthLog rows each path writes.
func (s *Session) validateSenderAuthorization(mailFrom string) (bool, string) {
	if mailFrom == "" {
		return false, "No sender address provided"
	}
	fromDomain := domainOfAddr(mailFrom)
	if fromDomain == "" {
		return false, "Invalid sender address format"
	}

	// A domain must have its DNS ownership TXT record verified before it can send —
	// otherwise anyone could add a domain they don't control and relay mail as it.
	dom, err := s.backend.DB.GetDomainByName(fromDomain)
	if err != nil {
		s.backend.Logger.Error("domain lookup failed: %v", err)
		return false, fmt.Sprintf("Not authorized to send for domain %s", fromDomain)
	}
	if dom == nil {
		return false, fmt.Sprintf("Domain %s is not configured on this server", fromDomain)
	}
	if !dom.IsVerified {
		return false, fmt.Sprintf("Domain %s has not completed DNS ownership verification yet", fromDomain)
	}

	if s.authenticatedSender != nil {
		sender := s.authenticatedSender
		if sender.CanSendAs(mailFrom) {
			return true, fmt.Sprintf("Sender authorized to send as %s", mailFrom)
		}
		_ = s.backend.DB.LogAuthAttempt("sender_validation", fmt.Sprintf("%s -> %s", sender.Email, mailFrom), s.peerIP, false, "")
		return false, fmt.Sprintf("Sender %s not authorized to send as %s", sender.Email, mailFrom)
	}

	wl, err := s.backend.DB.GetWhitelistedIP(s.peerIP, fromDomain)
	if err != nil {
		s.backend.Logger.Error("IP authorization lookup failed: %v", err)
		return false, fmt.Sprintf("Not authorized to send for domain %s", fromDomain)
	}
	if wl != nil {
		s.authType = "ip"
		s.authorizedDomain = fromDomain
		s.username = "IP:" + s.peerIP
		_ = s.backend.DB.LogAuthAttempt("ip", fmt.Sprintf("%s -> %s", s.peerIP, fromDomain), s.peerIP, true, fmt.Sprintf("IP %s authorized for domain %s", s.peerIP, fromDomain))
		return true, fmt.Sprintf("IP authorized for domain %s", fromDomain)
	}
	_ = s.backend.DB.LogAuthAttempt("ip", fmt.Sprintf("%s -> %s", s.peerIP, fromDomain), s.peerIP, false, fmt.Sprintf("IP %s not authorized for domain %s", s.peerIP, fromDomain))
	return false, fmt.Sprintf("Not authorized to send for domain %s", fromDomain)
}

func domainOfAddr(address string) string {
	i := strings.LastIndex(address, "@")
	if i < 0 {
		return ""
	}
	return strings.ToLower(address[i+1:])
}

// Rcpt mirrors handle_RCPT: accepts any address, no validation.
func (s *Session) Rcpt(to string, opts *smtp.RcptOptions) error {
	s.rcptTos = append(s.rcptTos, to)
	return nil
}

func internalError(msg string) error {
	return &smtp.SMTPError{Code: 550, EnhancedCode: smtp.NoEnhancedCode, Message: msg}
}

// Data mirrors EnhancedCustomSMTPHandler.handle_DATA end to end: Message-ID
// extraction/rehost, full header rebuild, DKIM signing, attachment extraction/storage,
// relay delivery, and EmailLog/EmailRecipientLog/EmailAttachment persistence.
func (s *Session) Data(r io.Reader) error {
	raw, err := io.ReadAll(r)
	if err != nil {
		return internalError("Internal server error")
	}
	content := string(raw)

	messageID := extractMessageID(content, s.backend.HeloHostname)
	senderDomain := domainOfAddr(s.mailFrom)

	var customHeaders [][2]string
	if senderDomain != "" {
		customHeaders, _ = s.backend.DKIM.GetActiveCustomHeaders(senderDomain)
	}
	customHeaders = append(customHeaders,
		[2]string{"X-Originating-IP", "[" + s.peerIP + "]"},
		[2]string{"X-Mailer", "NetBro Mail Server 1.0"},
		[2]string{"X-Priority", "3"},
	)

	rebuilt := ensureRequiredHeaders(content, messageID, s.rcptTos, s.mailFrom, customHeaders)

	signedContent := rebuilt
	dkimSigned := false
	if senderDomain != "" {
		signedContent = s.backend.DKIM.Sign(rebuilt, senderDomain)
		dkimSigned = signedContent != rebuilt
	}

	rebuiltHeaders := existingHeaders(rebuilt)
	toHeader := rebuiltHeaders["to"]
	ccHeader := rebuiltHeaders["cc"]
	subject := rebuiltHeaders["subject"]

	// Attachment storage: only if the authenticated sender or whitelisted IP opted in.
	storeMessage := false
	if sender, _ := s.backend.DB.GetSenderByEmail(s.mailFrom); sender != nil && sender.StoreMessageContent {
		storeMessage = true
	} else if wl, _ := s.backend.DB.GetWhitelistedIP(s.peerIP, senderDomain); wl != nil && wl.StoreMessageContent {
		storeMessage = true
	}

	parsed, parseErr := parseMessage(raw)

	type savedAttachment struct {
		Filename, ContentType, FilePath string
		Size                            int64
	}
	var toSave []savedAttachment
	if storeMessage && parseErr == nil && len(parsed.Attachments) > 0 {
		usernameOrIP := s.username
		if usernameOrIP == "" && s.peerIP != "" {
			usernameOrIP = sanitizePathSegment(s.peerIP, ":")
		} else {
			usernameOrIP = sanitizePathSegment(usernameOrIP, "/\\")
		}
		storagePath := attachmentStoragePath(s.backend.AttachmentsBasePath, senderDomain, usernameOrIP, time.Now())
		if err := os.MkdirAll(storagePath, 0o755); err == nil {
			prefix := cleanMessageIDPrefix(messageID)
			for _, a := range parsed.Attachments {
				filename := prefix + "_" + a.Filename
				fullPath := filepath.Join(storagePath, filename)
				if err := os.WriteFile(fullPath, a.Data, 0o644); err == nil {
					toSave = append(toSave, savedAttachment{Filename: a.Filename, ContentType: a.ContentType, FilePath: fullPath, Size: int64(len(a.Data))})
				} else {
					s.backend.Logger.Error("Failed to write attachment %s: %v", filename, err)
				}
			}
		}
	}

	// Classify each envelope recipient as to/cc/bcc by presence in the To/Cc headers —
	// anything not literally present in either is inferred BCC.
	toList := parseAddressList(toHeader)
	ccList := parseAddressList(ccHeader)
	recipientTypes := make([]string, len(s.rcptTos))
	for i, rcpt := range s.rcptTos {
		lower := strings.ToLower(rcpt)
		switch {
		case containsStr(toList, lower):
			recipientTypes[i] = "to"
		case containsStr(ccList, lower):
			recipientTypes[i] = "cc"
		default:
			recipientTypes[i] = "bcc"
		}
	}

	results := s.backend.Relay.RelayEmailAsync(s.mailFrom, s.rcptTos, signedContent, recipientTypes)

	allSucceeded := len(results) > 0
	for _, res := range results {
		if res.Status != "success" {
			allSucceeded = false
		}
	}

	var emailHeaders, messageBody string
	if parseErr == nil {
		emailHeaders = strings.Join(parsed.HeaderLines, "\n")
		messageBody = parsed.BodyText
	}

	logID, logErr := s.backend.Relay.LogEmail(s.backend.Cfg, s.peerIP, s.mailFrom, toHeader, ccHeader, "", subject, emailHeaders, messageBody, messageID, s.username, dkimSigned, results)
	if logErr != nil {
		s.backend.Logger.Error("Failed to log email: %v", logErr)
	} else {
		for _, a := range toSave {
			if err := s.backend.DB.InsertEmailAttachment(db.EmailAttachment{
				EmailLogID: logID, Filename: a.Filename, ContentType: a.ContentType, FilePath: a.FilePath, Size: a.Size,
			}); err != nil {
				s.backend.Logger.Error("Failed to record attachment %s: %v", a.Filename, err)
			}
		}
	}

	if allSucceeded {
		return &smtp.SMTPError{Code: 250, EnhancedCode: smtp.NoEnhancedCode, Message: "Message accepted for delivery"}
	}
	return &smtp.SMTPError{Code: 550, EnhancedCode: smtp.NoEnhancedCode, Message: "Message relay failed"}
}

func containsStr(list []string, s string) bool {
	for _, v := range list {
		if v == s {
			return true
		}
	}
	return false
}