import { registerPlugin, loadConfig, rootDir } from '../index.js';
import fs from 'fs';
import { dirname, join } from 'path';
import { fileURLToPath } from 'url';
import maxmind from 'maxmind';
import { AhoCorasick } from 'string-dsa';
import { getRealIP } from '../utils/network.js';
import { createGunzip } from 'zlib';
import tarStream from 'tar-stream';
import { Buffer } from 'buffer';
import * as logs from '../utils/logs.js';
import { recordEvent } from './stats.js';

const cfg = {};
await loadConfig('ipfilter', cfg);

// Map configuration to internal structure
const enabled = cfg.Core.Enabled;
const accountId = cfg.Core.AccountID || process.env.MAXMIND_ACCOUNT_ID;
const licenseKey = cfg.Core.LicenseKey || process.env.MAXMIND_LICENSE_KEY;
const dbUpdateInterval = cfg.Core.DBUpdateIntervalHours;

const ipBlockCacheTTL = cfg.Cache.IPBlockCacheTTLSec * 1000;
const ipBlockCacheMaxEntries = cfg.Cache.IPBlockCacheMaxEntries;

const blockedCountryCodes = new Set(cfg.Blocking.CountryCodes);
const blockedContinentCodes = new Set(cfg.Blocking.ContinentCodes);
const defaultBlockPage = cfg.Blocking.DefaultBlockPage;

// Process ASN blocks
const blockedASNs = {};
const asnGroupBlockPages = {};
for (const [group, config] of Object.entries(cfg.ASN || {})) {
  blockedASNs[group] = config.Numbers || [];
  asnGroupBlockPages[group] = config.BlockPage;
}

// Process ASN name blocks
const blockedASNNames = {};
for (const [group, config] of Object.entries(cfg.ASNNames || {})) {
  blockedASNNames[group] = config.Patterns || [];
  if (config.BlockPage) {
    asnGroupBlockPages[group] = config.BlockPage;
  }
}

const countryBlockPages = cfg.CountryBlockPages || {};
const continentBlockPages = cfg.ContinentBlockPages || {};

const ipBlockCache = new Map();

const blockPageCache = new Map();
async function loadBlockPage(filePath) {
  if (!blockPageCache.has(filePath)) {
    try {
      const txt = await fs.promises.readFile(filePath, 'utf8');
      blockPageCache.set(filePath, txt);
    } catch {
      blockPageCache.set(filePath, null);
    }
  }
  return blockPageCache.get(filePath);
}

const __dirname = dirname(fileURLToPath(import.meta.url));

const geoIPCountryDBPath = join(rootDir, 'data/GeoLite2-Country.mmdb');
const geoIPASNDBPath = join(rootDir, 'data/GeoLite2-ASN.mmdb');
const updateTimestampPath = join(rootDir, 'data/ipfilter_update.json');

let geoipCountryReader, geoipASNReader;

let isReloading = false;
let reloadLock = Promise.resolve();

async function getLastUpdateTimestamp() {
  try {
    if (fs.existsSync(updateTimestampPath)) {
      const data = await fs.promises.readFile(updateTimestampPath, 'utf8');
      const json = JSON.parse(data);
      return json.lastUpdated || 0;
    }
  } catch (err) {
    logs.warn('ipfilter', `Failed to read last update timestamp: ${err}`);
  }
  return 0;
}

async function saveUpdateTimestamp() {
  try {
    const timestamp = Date.now();
    await fs.promises.writeFile(
      updateTimestampPath,
      JSON.stringify({ lastUpdated: timestamp }),
      'utf8',
    );
    return timestamp;
  } catch (err) {
    logs.error('ipfilter', `Failed to save update timestamp: ${err}`);
    return Date.now();
  }
}

// Ensure the update timestamp file exists on first run
if (!fs.existsSync(updateTimestampPath)) {
  try {
    await saveUpdateTimestamp();
  } catch (err) {
    logs.error('ipfilter', `Failed to initialize update timestamp file: ${err}`);
  }
}

// Download GeoIP databases if missing
async function downloadGeoIPDatabases() {
  if (!licenseKey || !accountId) {
    logs.warn(
      'ipfilter',
      'No MaxMind credentials found; skipping GeoIP database download. Please set MAXMIND_ACCOUNT_ID and MAXMIND_LICENSE_KEY environment variables or add AccountID and LicenseKey to config/ipfilter.toml',
    );
    return;
  }
  const editions = [
    { id: 'GeoLite2-Country', filePath: geoIPCountryDBPath },
    { id: 'GeoLite2-ASN', filePath: geoIPASNDBPath },
  ];
  for (const { id, filePath } of editions) {
    if (!fs.existsSync(filePath)) {
      logs.plugin('ipfilter', `Downloading ${id} database...`);
      const url = `https://download.maxmind.com/app/geoip_download?edition_id=${id}&license_key=${licenseKey}&suffix=tar.gz`;
      const res = await fetch(url);
      if (!res.ok) {
        logs.error(
          'ipfilter',
          `Failed to download ${id} database: ${res.status} ${res.statusText}`,
        );
        continue;
      }
      const tempTar = join(rootDir, 'data', `${id}.tar.gz`);
      // write response body into a .tar.gz file
      const arrayBuf = await res.arrayBuffer();
      await fs.promises.writeFile(tempTar, Buffer.from(arrayBuf));
      // extract .mmdb files from the downloaded tar.gz
      const extract = tarStream.extract();
      extract.on('entry', (header, stream, next) => {
        if (header.name.endsWith('.mmdb')) {
          const filename = header.name.split('/').pop();
          const outPath = join(rootDir, 'data', filename);
          const ws = fs.createWriteStream(outPath);
          stream
            .pipe(ws)
            .on('finish', next)
            .on('error', (err) => {
              logs.error('ipfilter', `Extraction error: ${err}`);
              next();
            });
        } else {
          stream.resume();
          next();
        }
      });
      await new Promise((resolve, reject) => {
        fs.createReadStream(tempTar)
          .pipe(createGunzip())
          .pipe(extract)
          .on('finish', resolve)
          .on('error', reject);
      });
      await fs.promises.unlink(tempTar);
      logs.plugin('ipfilter', `${id} database downloaded and extracted.`);
    }
  }
}

await downloadGeoIPDatabases();

async function loadGeoDatabases() {
  if (isReloading) {
    await reloadLock;
    return true;
  }

  isReloading = true;
  let lockResolve;
  reloadLock = new Promise((resolve) => {
    lockResolve = resolve;
  });

  try {
    const countryStats = fs.statSync(geoIPCountryDBPath);
    const asnStats = fs.statSync(geoIPASNDBPath);

    if (countryStats.size > 1024 && asnStats.size > 1024) {
      logs.plugin('ipfilter', 'Initializing GeoIP databases from disk...');
      const newCountryReader = await maxmind.open(geoIPCountryDBPath);
      const newASNReader = await maxmind.open(geoIPASNDBPath);

      try {
        const testIP = '8.8.8.8';
        const countryTest = newCountryReader.get(testIP);
        const asnTest = newASNReader.get(testIP);

        if (!countryTest || !asnTest) {
          throw new Error('Database validation failed: test lookups returned empty results');
        }
      } catch (validationErr) {
        logs.error('ipfilter', `GeoIP database validation failed: ${validationErr}`);

        try {
          await newCountryReader.close();
        } catch (e) {}
        try {
          await newASNReader.close();
        } catch (e) {}
        throw new Error('Database validation failed');
      }

      const oldCountryReader = geoipCountryReader;
      const oldASNReader = geoipASNReader;

      geoipCountryReader = newCountryReader;
      geoipASNReader = newASNReader;
      if (oldCountryReader || oldASNReader) {
        logs.plugin('ipfilter', 'GeoIP databases reloaded and active');
      } else {
        logs.plugin('ipfilter', 'GeoIP databases loaded and active');
      }

      ipBlockCache.clear();

      await saveUpdateTimestamp();

      if (oldCountryReader || oldASNReader) {
        setTimeout(async () => {
          if (oldCountryReader) {
            try {
              await oldCountryReader.close();
            } catch (e) {}
          }
          if (oldASNReader) {
            try {
              await oldASNReader.close();
            } catch (e) {}
          }
          logs.plugin('ipfilter', 'Old GeoIP database instances closed successfully');
        }, 5000);
      }

      return true;
    } else {
      logs.warn(
        'ipfilter',
        'GeoIP database files are empty or too small. IP filtering will be disabled.',
      );
      return false;
    }
  } catch (err) {
    logs.error('ipfilter', `Failed to load GeoIP databases: ${err}`);
    return false;
  } finally {
    isReloading = false;
    lockResolve();
  }
}

async function checkAndUpdateDatabases() {
  if (isReloading) return false;

  const lastUpdate = await getLastUpdateTimestamp();
  const now = Date.now();
  const hoursSinceUpdate = (now - lastUpdate) / (1000 * 60 * 60);

  if (hoursSinceUpdate >= dbUpdateInterval) {
    logs.plugin(
      'ipfilter',
      `GeoIP databases last updated ${hoursSinceUpdate.toFixed(1)} hours ago, reloading...`,
    );
    return await loadGeoDatabases();
  }

  return false;
}

function startPeriodicDatabaseUpdates() {
  // Calculate interval in milliseconds
  const intervalMs = dbUpdateInterval * 60 * 60 * 1000;

  // Schedule periodic updates
  setInterval(async () => {
    try {
      await checkAndUpdateDatabases();
    } catch (err) {
      logs.error('ipfilter', `Failed during periodic database update: ${err}`);
    }
  }, intervalMs);

  logs.plugin('ipfilter', `Scheduled GeoIP database updates every ${dbUpdateInterval} hours`);
}

await loadGeoDatabases();

startPeriodicDatabaseUpdates();

const asnNameMatchers = new Map();
for (const [group, names] of Object.entries(blockedASNNames)) {
  asnNameMatchers.set(group, new AhoCorasick(names));
}

function cacheAndReturn(ip, blocked, blockType, blockValue, customPage, asnOrgName) {
  const expiresAt = Date.now() + ipBlockCacheTTL;
  ipBlockCache.set(ip, { blocked, blockType, blockValue, customPage, asnOrgName, expiresAt });
  // Enforce maximum cache size
  if (ipBlockCacheMaxEntries > 0 && ipBlockCache.size > ipBlockCacheMaxEntries) {
    // Remove the oldest entry (first key in insertion order)
    const oldestKey = ipBlockCache.keys().next().value;
    ipBlockCache.delete(oldestKey);
  }
  return [blocked, blockType, blockValue, customPage, asnOrgName];
}

function isBlockedIPExtended(ip) {
  const now = Date.now();
  const entry = ipBlockCache.get(ip);
  if (entry) {
    if (entry.expiresAt > now) {
      // Refresh recency by re-inserting entry
      ipBlockCache.delete(ip);
      ipBlockCache.set(ip, entry);
      return [entry.blocked, entry.blockType, entry.blockValue, entry.customPage, entry.asnOrgName];
    } else {
      // Entry expired, remove it
      ipBlockCache.delete(ip);
    }
  }

  const countryReader = geoipCountryReader;
  const asnReader = geoipASNReader;

  if (!countryReader || !asnReader) {
    return [false, '', '', '', ''];
  }

  let countryInfo;
  try {
    countryInfo = countryReader.get(ip);
  } catch (e) {}
  if (countryInfo?.country && blockedCountryCodes.has(countryInfo.country.iso_code)) {
    const page = countryBlockPages[countryInfo.country.iso_code] || defaultBlockPage;
    return cacheAndReturn(ip, true, 'country', countryInfo.country.iso_code, page, '');
  }

  if (countryInfo?.continent && blockedContinentCodes.has(countryInfo.continent.code)) {
    const page = continentBlockPages[countryInfo.continent.code] || defaultBlockPage;
    return cacheAndReturn(ip, true, 'continent', countryInfo.continent.code, page, '');
  }

  let asnInfo;
  try {
    asnInfo = asnReader.get(ip);
  } catch (e) {}
  if (asnInfo?.autonomous_system_number) {
    const asn = asnInfo.autonomous_system_number;
    const orgName = asnInfo.autonomous_system_organization || '';

    for (const [group, arr] of Object.entries(blockedASNs)) {
      if (arr.includes(asn)) {
        const page = asnGroupBlockPages[group] || defaultBlockPage;
        return cacheAndReturn(ip, true, 'asn', group, page, orgName);
      }
    }

    for (const [group, matcher] of asnNameMatchers.entries()) {
      const matches = matcher.find(orgName);
      if (matches.length) {
        const page = asnGroupBlockPages[group] || defaultBlockPage;
        return cacheAndReturn(ip, true, 'asn', group, page, orgName);
      }
    }
  }

  return cacheAndReturn(ip, false, '', '', '', '');
}

function IPBlockMiddleware() {
  return async (request, server) => {
    const clientIP = getRealIP(request, server);
    logs.plugin('ipfilter', `Incoming request from IP: ${clientIP}`);
    const [blocked, blockType, blockValue, customPage, asnOrgName] = isBlockedIPExtended(clientIP);

    if (blocked) {
      recordEvent('ipfilter.block', {
        type: blockType,
        value: blockValue,
        asn_org: asnOrgName,
        ip: clientIP, // Include the IP address for stats
      });
      const url = new URL(request.url);

      if (url.pathname.startsWith('/api')) {
        return new Response(
          JSON.stringify({
            error: 'Access denied from your location or network.',
            reason: 'geoip',
            type: blockType,
            value: blockValue,
            asn_org: asnOrgName,
          }),
          {
            status: 403,
            headers: { 'Content-Type': 'application/json' },
          },
        );
      }

      // Normalize page paths by stripping leading slash
      const cleanCustomPage = customPage.replace(/^\/+/, '');
      const cleanDefaultPage = defaultBlockPage.replace(/^\/+/, '');

      let html = '';
      logs.plugin(
        'ipfilter',
        `Block pages: custom="${cleanCustomPage}", default="${cleanDefaultPage}"`,
      );
      logs.plugin('ipfilter', 'Searching for block page in the following locations:');
      const paths = [
        // allow absolute paths relative to project root first
        join(rootDir, cleanCustomPage),
      ];
      // Fallback to default block page if custom page isn't found
      if (customPage !== defaultBlockPage) {
        paths.push(
          // check default page at root directory
          join(rootDir, cleanDefaultPage),
        );
      }

      for (const p of paths) {
        logs.plugin('ipfilter', `Trying block page at: ${p}`);
        const content = await loadBlockPage(p);
        logs.plugin('ipfilter', `Load result for ${p}: ${content ? 'FOUND' : 'NOT FOUND'}`);
        if (content) {
          html = content;
          break;
        }
      }

      if (html) {
        const output = html.replace('{{.ASNName}}', asnOrgName || 'Blocked Network');
        return new Response(output, {
          status: 403,
          headers: { 'Content-Type': 'text/html; charset=utf-8' },
        });
      } else {
        return new Response('Access denied from your location or network.', {
          status: 403,
          headers: { 'Content-Type': 'text/plain' },
        });
      }
    }

    return undefined;
  };
}

if (enabled) {
  registerPlugin('ipfilter', IPBlockMiddleware());
} else {
  logs.plugin('ipfilter', 'IP filter plugin disabled via config');
}

export { checkAndUpdateDatabases, loadGeoDatabases };