import Koa from 'koa';
import Router from '@koa/router';
import sharp from 'sharp';
import dns from 'dns/promises';
import assert from 'assert';

const { DEST_SERVER, SERVER_PRETTY_NAME, LISTENING_PORT } = process.env;
assert(DEST_SERVER, 'missing DEST_SERVER env variable');
assert(SERVER_PRETTY_NAME, 'missing SERVER_PRETTY_NAME env variable');
const IS_PRODUCTION = process.env.NODE_ENV === 'production';
const FUCKED_UP_SERVERS = ['matrix.org'];
const defaultHeaders = {
  'User-Agent': 'why-is-synapse (https://git.sakamoto.pl/selfisekai/why-is-synapse)',
};

const existValueOrNoExistKey = (obj: Record<string, any>): Record<string, string> =>
  Object.fromEntries(
    Object.entries(obj).filter(([, val]) => typeof val === 'string') as [string, string][],
  );

const serverHostnameCache = new Map<string, string>();

const validateHostname = (serverName: string) => {
  assert(/^[^/?#]+$/.test(serverName), `invalid hostname: ${JSON.stringify(serverName)}`);
};

const getServerHostname = async (serverName: string) => {
  const cached = serverHostnameCache.get(serverName);
  if (cached) return cached;

  const jsonRes = await fetch(`https://${serverName}/.well-known/matrix/server`, {
    headers: defaultHeaders,
  });
  if (jsonRes.status === 200) {
    let res = (await jsonRes.json())['m.server'];
    if (typeof res === 'string') {
      validateHostname(res);
      if (!res.includes(':')) {
        res += ':8448'; // default port per spec
      }
      serverHostnameCache.set(serverName, res);
      return res;
    }
  }

  const dnsRes = await dns.resolveSrv(`_matrix._tcp.${serverName}`).catch(() => []);
  if (dnsRes.length > 0) {
    dnsRes.sort((a, b) => a.priority - b.priority);
    console.log(dnsRes);
    const res = `${dnsRes[0].name}:${dnsRes[0].port}`;
    validateHostname(res);
    serverHostnameCache.set(serverName, res);
    return res;
  }

  // if there are no SRV/well-known specs, assume the defaults
  return `${serverName}:8448`;
};

const koa = new Koa();
const router = new Router({
  prefix: '/_matrix',
});
const media = new Router();

media.get(
  ['/download/:serverName/:mediaId', '/download/:serverName/:mediaId/:filename'],
  async (ctx) => {
    const { serverName, mediaId, filename } = ctx.params;

    async function sendFile(f: Response) {
      ctx.res.statusCode = 200;
      (
        ['content-type', filename ? 'content-disposition' : null].filter(
          (header) => header,
        ) as string[]
      ).forEach((header) => {
        if (!f) return; // typescript is bork
        const val = f.headers.get(header);
        if (!val) return;
        ctx.res.setHeader(header, val);
      });
      if (filename) {
        ctx.res.setHeader('content-disposition', `attachment; filename="${filename}"`);
      }
      ctx.res.write(Buffer.from(await f.arrayBuffer()));
      ctx.res.end();
    }

    let file: Response | null = null;
    if (!FUCKED_UP_SERVERS.includes(serverName)) {
      file = await fetch(`${DEST_SERVER}/_matrix/media/r0/download/${serverName}/${mediaId}`, {
        headers: {
          ...defaultHeaders,
          ...existValueOrNoExistKey({
            Authorization: ctx.headers['authorization'],
          }),
        },
      });
      if (file && file.status === 200) {
        await sendFile(file);
        return;
      }
    }

    if (serverName === SERVER_PRETTY_NAME) {
      ctx.res.statusCode = 500;
      return;
    }

    const remoteHost = await getServerHostname(serverName);
    file = await fetch(`https://${remoteHost}/_matrix/media/r0/download/${serverName}/${mediaId}`, {
      headers: defaultHeaders,
    });
    if (file && file.status === 200) {
      await sendFile(file);
      return;
    }
  },
);

media.get('/thumbnail/:serverName/:mediaId', async (ctx) => {
  const { serverName, mediaId } = ctx.params;
  if (
    typeof ctx.query.width !== 'string' ||
    typeof ctx.query.height !== 'string' ||
    (typeof ctx.query.method !== 'string' && typeof ctx.query.method !== 'undefined')
  ) {
    return;
  }
  const { method } = ctx.query;
  if (method) {
    assert(['crop', 'scale'].includes(method));
  }
  const width = parseInt(ctx.query.width, 10);
  const height = parseInt(ctx.query.height, 10);

  let file: Response | null = null;
  if (!FUCKED_UP_SERVERS.includes(serverName)) {
    file = await fetch(`${DEST_SERVER}/_matrix/media/r0/download/${serverName}/${mediaId}`, {
      headers: defaultHeaders,
    });
  }

  if (!file || file.status !== 200) {
    if (serverName === SERVER_PRETTY_NAME) {
      ctx.res.statusCode = 500;
      return;
    }
    const remoteHost = await getServerHostname(serverName);
    file = await fetch(`https://${remoteHost}/_matrix/media/r0/download/${serverName}/${mediaId}`, {
      headers: {
        ...defaultHeaders,
        ...existValueOrNoExistKey({
          Authorization: ctx.headers['authorization'],
        }),
      },
    });
  }
  if (file && file.status === 200) {
    const resizer = sharp(Buffer.from(await file.arrayBuffer())).resize(width, height, {
      fit: method === 'crop' ? 'cover' : 'inside',
    });
    const img = await resizer.toBuffer({ resolveWithObject: true });
    ctx.res.statusCode = 200;
    ctx.res.setHeader('content-type', `image/${img.info.format}`);
    ctx.res.write(img.data);
    ctx.res.end();
    return;
  }
});

router.use('/media/r0', media.routes(), media.allowedMethods());
router.use('/media/v3', media.routes(), media.allowedMethods());

koa.use(router.routes()).use(router.allowedMethods());
koa.listen(parseInt(LISTENING_PORT || '8009', 10));