const dns = require('dns').promises;
const logger = require('./logger');

const IP_CONNECTION_COUNTS = new Map();
const BANNED_IPS = new Set();

const IP_RATE_LIMIT = 20; // 每个IP(服务器)每分钟最多20次连接
const IP_TIME_WINDOW = 60 * 1000; // 1分钟
const IP_BAN_DURATION = 5 * 60 * 1000; // 5分钟

/**
 * 检查IP地址是否因为连接频率过高而被限制。
 * @param {string} remoteAddress 客户端IP地址。
 * @returns {boolean} 如果被限制则返回true，否则返回false。
 */
function isIpRateLimited(remoteAddress) {
    if (BANNED_IPS.has(remoteAddress)) {
        logger.warn(`Connection from banned IP ${remoteAddress} rejected.`);
        return true;
    }

    const now = Date.now();
    const requests = IP_CONNECTION_COUNTS.get(remoteAddress) || [];
    const recentRequests = requests.filter(timestamp => now - timestamp < IP_TIME_WINDOW);

    if (recentRequests.length >= IP_RATE_LIMIT) {
        logger.warn(`IP ${remoteAddress} has exceeded the connection rate limit. Banning for ${IP_BAN_DURATION / 1000} seconds.`);
        BANNED_IPS.add(remoteAddress);
        setTimeout(() => {
            BANNED_IPS.delete(remoteAddress);
            logger.info(`IP ${remoteAddress} has been unbanned.`);
        }, IP_BAN_DURATION);
        IP_CONNECTION_COUNTS.delete(remoteAddress);
        return true;
    }

    recentRequests.push(now);
    IP_CONNECTION_COUNTS.set(remoteAddress, recentRequests);
    return false;
}

/**
 * 检查IP地址是否有有效的反向DNS（PTR）记录。
 * 正规的邮件服务器通常都有PTR记录。
 * @param {string} remoteAddress 客户端IP地址。
 * @returns {Promise<boolean>} 如果验证通过则返回true，否则返回false。
 */
async function hasValidPtrRecord(remoteAddress) {
    // 对于本地和私有地址，我们跳过检查，因为它们通常没有公共PTR记录
    if (remoteAddress.startsWith('127.') || remoteAddress.startsWith('192.168.') || remoteAddress.startsWith('10.') || remoteAddress.startsWith('::1')) {
        return true;
    }

    try {
        const hostnames = await dns.reverse(remoteAddress);
        if (hostnames && hostnames.length > 0) {
            logger.info(`PTR record for ${remoteAddress} found: ${hostnames.join(', ')}`);
            return true;
        }
        logger.warn(`No PTR record found for ${remoteAddress}.`);
        return false;
    } catch (error) {
        // 'ENOTFOUND' 是最常见的错误，意味着没有找到PTR记录。
        if (error.code === 'ENOTFOUND') {
            logger.warn(`No PTR record found for ${remoteAddress}.`);
        } else {
            logger.error(`Error during PTR lookup for ${remoteAddress}:`, error);
        }
        return false;
    }
}

/**
 * 在连接建立时验证客户端。
 * @param {object} session SMTP会话对象。
 * @param {function} callback 回调函数。
 */
async function validateConnection(session, callback) {
    const { remoteAddress } = session;

    // 1. IP频率限制检查
    if (isIpRateLimited(remoteAddress)) {
        const err = new Error('Connection rejected due to high frequency. Please try again later.');
        err.responseCode = 421;
        return callback(err);
    }

    // 2. 反向DNS检查
    const hasPtr = await hasValidPtrRecord(remoteAddress);
    if (!hasPtr) {
        const err = new Error('Connection rejected: The IP address has no PTR record.');
        err.responseCode = 550; // 550表示请求的操作未执行，邮箱不可用（在这里引申为连接源不可信）
        return callback(err);
    }

    // 所有检查通过
    callback();
}

module.exports = { validateConnection };