package mailruntime

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"net/mail"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/extension/providers/mail/connection"
	"github.com/tidbcloud/AgentTeam/internal/extension/providers/mail/mailstore"
	"github.com/tidbcloud/AgentTeam/internal/extension/providers/mail/protocol"
	"github.com/tidbcloud/AgentTeam/internal/extension/providers/mail/transport"
	"github.com/tidbcloud/AgentTeam/internal/store/app"
	"golang.org/x/net/idna"
)

const storedSenderSubjectDomain = "agentteam/email/external-subject/v1\x00"

type storedAuth struct {
	Username string `json:"username"`
	Password string `json:"password"`
}

type storedCredentials struct {
	SchemaVersion int        `json:"schemaVersion"`
	IMAP          storedAuth `json:"imap"`
	SMTP          storedAuth `json:"smtp"`
}

type storedEndpoint = connection.Endpoint
type storedConnectionConfig = connection.Config

type connectionMaterial struct {
	Connection    mailstore.Connection
	SenderAddress string
	IMAP          protocol.Endpoint
	SMTP          protocol.Endpoint
	Capabilities  protocol.Capabilities
}

func decodeConnectionMaterial(
	currentKey []byte,
	legacyKeys [][]byte,
	connection mailstore.Connection,
) (connectionMaterial, error) {
	plaintext, err := app.DecryptStoredCredentialsJSONWithFallback(
		currentKey,
		legacyKeys,
		connection.EncryptedCredentials,
	)
	if err != nil {
		return connectionMaterial{}, fmt.Errorf(
			"decrypt Email Connection credentials: %w",
			err,
		)
	}
	var credentials storedCredentials
	if err := decodeStoredJSON(plaintext, &credentials); err != nil {
		return connectionMaterial{}, fmt.Errorf(
			"decode Email Connection credentials: %w",
			err,
		)
	}
	var config storedConnectionConfig
	if err := decodeStoredJSON(connection.Config, &config); err != nil {
		return connectionMaterial{}, fmt.Errorf(
			"decode Email Connection config: %w",
			err,
		)
	}
	return materialFromStoredConnection(
		currentKey,
		legacyKeys,
		connection,
		credentials,
		config,
	)
}

func materialFromStoredConnection(
	currentKey []byte,
	legacyKeys [][]byte,
	connection mailstore.Connection,
	credentials storedCredentials,
	config storedConnectionConfig,
) (connectionMaterial, error) {
	if credentials.SchemaVersion != 1 ||
		config.SchemaVersion != 1 ||
		!config.SenderAddressAccepted ||
		strings.TrimSpace(config.SenderAddress) == "" ||
		len(config.SenderAddress) > 320 ||
		strings.ContainsAny(config.SenderAddress, "\r\n") {
		return connectionMaterial{}, fmt.Errorf(
			"Email Connection material is invalid",
		)
	}
	normalizedSender, err := normalizeStoredSenderAddress(config.SenderAddress)
	if err != nil || normalizedSender != config.SenderAddress {
		return connectionMaterial{}, fmt.Errorf(
			"Email Connection sender is invalid",
		)
	}
	if !matchesStoredExternalSubject(
		currentKey,
		legacyKeys,
		config.SenderAddress,
		connection.ExternalSubjectRef,
	) {
		return connectionMaterial{}, fmt.Errorf(
			"Email Connection sender identity is inconsistent",
		)
	}
	if err := validateStoredAuth(credentials.IMAP); err != nil {
		return connectionMaterial{}, fmt.Errorf(
			"Email IMAP credentials are invalid",
		)
	}
	if err := validateStoredAuth(credentials.SMTP); err != nil {
		return connectionMaterial{}, fmt.Errorf(
			"Email SMTP credentials are invalid",
		)
	}
	imapHost, err := validateStoredEndpoint(
		transport.ProtocolIMAP,
		config.IMAP,
	)
	if err != nil {
		return connectionMaterial{}, err
	}
	smtpHost, err := validateStoredEndpoint(
		transport.ProtocolSMTP,
		config.SMTP,
	)
	if err != nil {
		return connectionMaterial{}, err
	}
	if !isKnownCapabilitySet(config.IMAP.Capabilities) ||
		!isKnownCapabilitySet(config.SMTP.Capabilities) {
		return connectionMaterial{}, fmt.Errorf(
			"Email Connection capability set is invalid",
		)
	}
	if err := config.VerifyDigest(); err != nil {
		return connectionMaterial{}, fmt.Errorf(
			"Email Connection capability digest is invalid",
		)
	}
	values := make(
		map[string]bool,
		len(config.IMAP.Capabilities),
	)
	for _, value := range config.IMAP.Capabilities {
		values[value] = true
	}
	return connectionMaterial{
		Connection:    connection,
		SenderAddress: config.SenderAddress,
		IMAP: protocol.Endpoint{
			Host:     imapHost,
			Port:     config.IMAP.Port,
			TLSMode:  config.IMAP.TLSMode,
			Username: credentials.IMAP.Username,
			Password: credentials.IMAP.Password,
		},
		SMTP: protocol.Endpoint{
			Host:     smtpHost,
			Port:     config.SMTP.Port,
			TLSMode:  config.SMTP.TLSMode,
			Username: credentials.SMTP.Username,
			Password: credentials.SMTP.Password,
		},
		Capabilities: protocol.Capabilities{
			Values:       values,
			DraftMailbox: config.IMAP.DraftMailbox,
		},
	}, nil
}

func (material connectionMaterial) supports(action string) bool {
	hasIMAP := material.Capabilities.Has("IMAP4REV1") ||
		material.Capabilities.Has("IMAP4REV2")
	switch action {
	case "search-messages", "get-message", "list-attachments",
		"get-attachment":
		return hasIMAP
	case "create-draft":
		return hasIMAP &&
			material.Capabilities.DraftMailbox != "" &&
			material.Capabilities.Has("SPECIAL-USE") &&
			material.Capabilities.Has("UIDPLUS")
	case "update-draft", "create-reply-draft",
		"create-forward-draft", "send-draft":
		return hasIMAP && material.Capabilities.SafeDraftMutation()
	default:
		return false
	}
}

func validateStoredAuth(auth storedAuth) error {
	if strings.TrimSpace(auth.Username) != auth.Username ||
		auth.Username == "" || len(auth.Username) > 512 ||
		auth.Password == "" || len(auth.Password) > 4096 ||
		strings.ContainsAny(auth.Username+auth.Password, "\r\n") {
		return fmt.Errorf("invalid Email authentication")
	}
	return nil
}

func validateStoredEndpoint(
	protocol string,
	endpoint storedEndpoint,
) (string, error) {
	host, err := transport.NormalizeHostname(endpoint.Host)
	if err != nil {
		return "", err
	}
	if host != endpoint.Host {
		return "", fmt.Errorf("Email endpoint host is not normalized")
	}
	if err := transport.ValidateTarget(transport.Target{
		Protocol: protocol,
		Host:     host,
		Port:     endpoint.Port,
		TLSMode:  endpoint.TLSMode,
	}); err != nil {
		return "", err
	}
	if protocol == transport.ProtocolSMTP &&
		endpoint.DraftMailbox != "" {
		return "", fmt.Errorf(
			"SMTP endpoint cannot name an IMAP Drafts mailbox",
		)
	}
	if len(endpoint.DraftMailbox) > 1024 ||
		strings.ContainsAny(endpoint.DraftMailbox, "\r\n") {
		return "", fmt.Errorf("IMAP Drafts mailbox is invalid")
	}
	return host, nil
}

func isKnownCapabilitySet(values []string) bool {
	if len(values) > 128 {
		return false
	}
	for index, value := range values {
		if value == "" || len(value) > 128 ||
			value != strings.ToUpper(value) ||
			strings.ContainsAny(value, "\r\n ") {
			return false
		}
		if index > 0 && values[index-1] >= value {
			return false
		}
	}
	return true
}

func decodeStoredJSON(payload []byte, destination any) error {
	decoder := json.NewDecoder(bytes.NewReader(payload))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(destination); err != nil {
		return err
	}
	var trailing json.RawMessage
	if err := decoder.Decode(&trailing); err != io.EOF {
		if err == nil {
			return fmt.Errorf("JSON must contain one document")
		}
		return fmt.Errorf("JSON has trailing data")
	}
	return nil
}

func normalizeStoredSenderAddress(value string) (string, error) {
	parsed, err := mail.ParseAddress(value)
	if err != nil || parsed.Address != value {
		return "", fmt.Errorf("invalid Email sender")
	}
	at := strings.LastIndexByte(parsed.Address, '@')
	if at <= 0 || at == len(parsed.Address)-1 {
		return "", fmt.Errorf("invalid Email sender")
	}
	domain, err := idna.Lookup.ToASCII(parsed.Address[at+1:])
	if err != nil || domain == "" {
		return "", fmt.Errorf("invalid Email sender domain")
	}
	return parsed.Address[:at] + "@" + strings.ToLower(domain), nil
}

func matchesStoredExternalSubject(
	currentKey []byte,
	legacyKeys [][]byte,
	sender string,
	stored string,
) bool {
	keys := make([][]byte, 0, len(legacyKeys)+1)
	keys = append(keys, currentKey)
	keys = append(keys, legacyKeys...)
	for _, key := range keys {
		expected := storedExternalSubjectRef(key, sender)
		if expected != "" &&
			hmac.Equal([]byte(expected), []byte(stored)) {
			return true
		}
	}
	return false
}

func storedExternalSubjectRef(key []byte, sender string) string {
	if len(key) != 32 {
		return ""
	}
	mac := hmac.New(sha256.New, key)
	_, _ = mac.Write([]byte(storedSenderSubjectDomain))
	_, _ = mac.Write([]byte(sender))
	return "email:h1:" + hex.EncodeToString(mac.Sum(nil))
}