'use strict';

const fs = require('fs');
const http = require('http');
const https = require('https');
const express = require('express');
const WebSocket = require('ws');

const LISTEN_PORT = Number(process.env.LISTEN_PORT || 8080);
const HOSTS_FILE = process.env.HOSTS_FILE || '/config/hosts.txt';
const CHECK_INTERVAL_SEC = Math.max(
  5,
  Number(process.env.CHECK_INTERVAL_SEC || 30) || 30
);
const PING_TIMEOUT_MS = Math.max(1000, Number(process.env.PING_TIMEOUT_MS || 8000));
const HEALTH_PATH = process.env.HEALTH_PATH || '/';
const REJECT_UNAUTHORIZED = process.env.PROXY_INSECURE_TLS !== '1';

const HOP_BY_HOP = new Set([
  'connection',
  'keep-alive',
  'proxy-authenticate',
  'proxy-authorization',
  'te',
  'trailers',
  'transfer-encoding',
  'upgrade',
]);

let currentTarget = null;
let lastStats = { checkedAt: null, results: [], selected: null };

function readHostsFile() {
  if (!fs.existsSync(HOSTS_FILE)) {
    console.error(`Hosts file not found: ${HOSTS_FILE}`);
    return [];
  }
  const raw = fs.readFileSync(HOSTS_FILE, 'utf8');
  const lines = raw.split(/\r?\n/);
  const urls = [];
  for (const line of lines) {
    const t = line.trim();
    if (!t || t.startsWith('#')) continue;
    try {
      const u = new URL(t);
      if (u.protocol !== 'http:' && u.protocol !== 'https:') continue;
      const base = `${u.protocol}//${u.host}`;
      urls.push(base);
    } catch {
      console.warn(`Skipping invalid URL line: ${t}`);
    }
  }
  return [...new Set(urls)];
}

function pingBase(base, path) {
  const HARD_CAP_MS = PING_TIMEOUT_MS + 2000;

  return new Promise((resolve) => {
    let url;
    try {
      const p = path.startsWith('/') ? path : `/${path}`;
      url = new URL(p, base.endsWith('/') ? base : `${base}/`);
    } catch {
      resolve({ base, ok: false, ms: null, error: 'bad_url' });
      return;
    }

    let settled = false;
    let clientReq = null;

    const finish = (result) => {
      if (settled) return;
      settled = true;
      clearTimeout(hardTimer);
      try {
        clientReq?.destroy();
      } catch (_) {}
      resolve(result);
    };

    const hardTimer = setTimeout(() => {
      try {
        clientReq?.destroy();
      } catch (_) {}
      finish({ base, ok: false, ms: null, error: 'hard_timeout' });
    }, HARD_CAP_MS);

    const lib = url.protocol === 'https:' ? https : http;
    const start = Date.now();
    clientReq = lib.request(
      {
        protocol: url.protocol,
        hostname: url.hostname,
        port: url.port || (url.protocol === 'https:' ? 443 : 80),
        path: url.pathname + url.search,
        method: 'GET',
        timeout: PING_TIMEOUT_MS,
        rejectUnauthorized: REJECT_UNAUTHORIZED,
        headers: {
          'User-Agent': 'xtream-proxy-health/1.0',
          Connection: 'close',
        },
      },
      (res) => {
        res.resume();
        res.on('end', () => {
          const ms = Date.now() - start;
          const ok = res.statusCode >= 200 && res.statusCode < 500;
          finish({ base, ok, ms, status: res.statusCode });
        });
        res.on('error', () => {
          finish({ base, ok: false, ms: null, error: 'response_error' });
        });
      }
    );
    clientReq.on('error', (e) => {
      finish({ base, ok: false, ms: null, error: e.code || e.message });
    });
    clientReq.on('timeout', () => {
      clientReq.destroy();
      finish({ base, ok: false, ms: null, error: 'timeout' });
    });
    clientReq.end();
  });
}

async function pickFastestHost() {
  const hosts = readHostsFile();
  if (hosts.length === 0) {
    console.warn('No valid hosts in file; keeping previous target if any.');
    lastStats = {
      checkedAt: new Date().toISOString(),
      results: [],
      selected: currentTarget,
    };
    return;
  }

  const pings = await Promise.all(hosts.map((h) => pingBase(h, HEALTH_PATH)));
  const ok = pings.filter((p) => p.ok && typeof p.ms === 'number');
  ok.sort((a, b) => a.ms - b.ms);

  lastStats = {
    checkedAt: new Date().toISOString(),
    results: pings,
    selected: ok[0]?.base ?? currentTarget,
  };

  if (ok[0]) {
    const prev = currentTarget;
    currentTarget = ok[0].base;
    const others = ok.slice(1).map((x) => `${x.base}=${x.ms}ms`).join(', ') || 'none';
    if (prev !== currentTarget) {
      console.log(`Upstream switched → ${currentTarget} (${ok[0].ms} ms). Also: ${others}`);
    } else {
      console.log(`Health check: keeping ${currentTarget} (${ok[0].ms} ms). Others: ${others}`);
    }
  } else {
    console.warn(
      'No host responded successfully; keeping previous target:',
      currentTarget || '(none)'
    );
  }
}

let checkTimer = null;
let checkRunning = false;

function scheduleChecks() {
  const delayMs = CHECK_INTERVAL_SEC * 1000;
  const tick = async () => {
    if (checkRunning) {
      console.warn('Health check still running; skipping overlapping run (next in ' + CHECK_INTERVAL_SEC + 's)');
      checkTimer = setTimeout(tick, delayMs);
      return;
    }
    checkRunning = true;
    try {
      await pickFastestHost();
    } catch (e) {
      console.error('Health check error:', e);
    } finally {
      checkRunning = false;
      checkTimer = setTimeout(tick, delayMs);
    }
  };
  checkTimer = setTimeout(tick, delayMs);
}

function filterHopByHop(headers) {
  const out = Object.create(null);
  if (!headers) return out;
  for (const [k, v] of Object.entries(headers)) {
    if (v === undefined) continue;
    if (HOP_BY_HOP.has(k.toLowerCase())) continue;
    out[k] = v;
  }
  return out;
}

function buildRequestHeaders(req, targetUrl) {
  const out = filterHopByHop(req.headers);
  const existing = req.headers['x-forwarded-for'];
  const client = req.socket.remoteAddress || '';
  out['x-forwarded-for'] = existing ? `${existing}, ${client}` : client;
  const xfProto = req.headers['x-forwarded-proto'];
  out['x-forwarded-proto'] = typeof xfProto === 'string' && xfProto ? xfProto : 'http';
  const xfHost = req.headers['x-forwarded-host'] || req.headers.host;
  if (xfHost) out['x-forwarded-host'] = xfHost;
  out.host = targetUrl.host;
  return out;
}

function proxyHttp(req, res) {
  const base = currentTarget.replace(/\/$/, '');
  let targetUrl;
  try {
    targetUrl = new URL(req.url, `${base}/`);
  } catch {
    res.status(400).type('text/plain').send('Bad request URL');
    return;
  }

  const lib = targetUrl.protocol === 'https:' ? https : http;
  const opts = {
    protocol: targetUrl.protocol,
    hostname: targetUrl.hostname,
    port: targetUrl.port || (targetUrl.protocol === 'https:' ? 443 : 80),
    path: targetUrl.pathname + targetUrl.search,
    method: req.method,
    headers: buildRequestHeaders(req, targetUrl),
    rejectUnauthorized: REJECT_UNAUTHORIZED,
  };

  const pReq = lib.request(opts, (pRes) => {
    res.writeHead(pRes.statusCode, filterHopByHop(pRes.headers));
    pRes.pipe(res);
  });

  pReq.on('error', (err) => {
    if (!res.headersSent) {
      res.writeHead(502, { 'Content-Type': 'text/plain' });
    }
    res.end(`Upstream error: ${err.message}`);
  });

  req.pipe(pReq);
}

function toWsUrl(httpBase, reqPath) {
  const base = httpBase.replace(/\/$/, '');
  const u = new URL(reqPath, `${base}/`);
  const proto = u.protocol === 'https:' ? 'wss:' : 'ws:';
  return `${proto}//${u.host}${u.pathname}${u.search}`;
}

function buildWsClientHeaders(req, targetWsUrl) {
  const u = new URL(targetWsUrl);
  const out = filterHopByHop(req.headers);
  for (const key of Object.keys(out)) {
    if (key.toLowerCase().startsWith('sec-websocket')) {
      delete out[key];
    }
  }
  delete out.host;
  out.host = u.host;
  return out;
}

const app = express();

app.get('/health', (_req, res) => {
  res.json({
    ok: Boolean(currentTarget),
    upstream: currentTarget,
    intervalSec: CHECK_INTERVAL_SEC,
    lastCheck: lastStats,
  });
});

app.use((req, res) => {
  if (!currentTarget) {
    res.status(503).type('text/plain').send('No upstream available yet (hosts file empty or all pings failed).');
    return;
  }
  proxyHttp(req, res);
});

const server = http.createServer(app);
const wss = new WebSocket.Server({ noServer: true });

server.on('upgrade', (req, socket, head) => {
  if (!currentTarget) {
    socket.write('HTTP/1.1 503 Service Unavailable\r\nConnection: close\r\n\r\n');
    socket.destroy();
    return;
  }

  let upstreamUrl;
  try {
    upstreamUrl = toWsUrl(currentTarget, req.url);
  } catch {
    socket.destroy();
    return;
  }

  wss.handleUpgrade(req, socket, head, (clientWs) => {
    const upstream = new WebSocket(upstreamUrl, {
      rejectUnauthorized: REJECT_UNAUTHORIZED,
      headers: buildWsClientHeaders(req, upstreamUrl),
    });

    const pending = [];
    let upstreamReady = false;

    clientWs.on('message', (data, isBinary) => {
      if (upstreamReady && upstream.readyState === WebSocket.OPEN) {
        upstream.send(data, { binary: isBinary });
      } else {
        pending.push({ data, isBinary });
      }
    });

    upstream.on('open', () => {
      upstreamReady = true;
      for (const m of pending) {
        if (upstream.readyState === WebSocket.OPEN) upstream.send(m.data, { binary: m.isBinary });
      }
      pending.length = 0;

      upstream.on('message', (data, isBinary) => {
        if (clientWs.readyState === WebSocket.OPEN) clientWs.send(data, { binary: isBinary });
      });
      clientWs.on('ping', (buf) => upstream.ping(buf));
      upstream.on('ping', (buf) => clientWs.ping(buf));
      clientWs.on('pong', (buf) => upstream.pong(buf));
      upstream.on('pong', (buf) => clientWs.pong(buf));
    });

    const shutdown = () => {
      try {
        clientWs.terminate();
      } catch (_) {}
      try {
        upstream.terminate();
      } catch (_) {}
    };

    upstream.on('error', shutdown);
    clientWs.on('error', shutdown);
    upstream.on('close', shutdown);
    clientWs.on('close', shutdown);
  });
});

async function main() {
  await pickFastestHost();
  server.listen(LISTEN_PORT, () => {
    console.log(`Xtream proxy listening on :${LISTEN_PORT}`);
    console.log(
      `Hosts file: ${HOSTS_FILE}, interval: ${CHECK_INTERVAL_SEC}s (from end of each run), health path: ${HEALTH_PATH}`
    );
    scheduleChecks();
  });
}

function shutdown() {
  if (checkTimer) {
    clearTimeout(checkTimer);
    checkTimer = null;
  }
  server.close(() => process.exit(0));
}

process.on('SIGTERM', shutdown);
process.on('SIGINT', shutdown);

main().catch((e) => {
  console.error(e);
  process.exit(1);
});