#ifndef _DEFAULT_SOURCE
#define _DEFAULT_SOURCE
#endif

#include "cst_web_md.h"

#include "cst_util.h"
#include "sashidown.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

struct cst_web_md {
  sashidown_stream_t *st;
  cst_web_md_emit_fn emit;
  void *ud;
  cst_buf_t in;    /* HTML from sashidown, waiting for a complete tag */
  int failed;
};

/* --- sanitizer ------------------------------------------------------------ */

static const char *ALLOWED_TAGS[] = {
    "p",      "br",     "hr",     "h1",     "h2",      "h3",     "h4",
    "h5",     "h6",     "strong", "em",     "b",       "i",      "u",
    "s",      "del",    "code",   "pre",    "blockquote", "ul",  "ol",
    "li",     "table",  "thead",  "tbody",  "tr",      "th",     "td",
    "a",      "img",    "span",   "div",    NULL};

static int tag_allowed(const char *name, size_t n) {
  for (int i = 0; ALLOWED_TAGS[i]; i++) {
    if (strlen(ALLOWED_TAGS[i]) == n && strncmp(name, ALLOWED_TAGS[i], n) == 0)
      return 1;
  }
  return 0;
}

/* A URL we are willing to put in href: http(s), mailto, or a same-site
 * path. Anything else (javascript:, data:, file:, //evil) is dropped. */
static int href_ok(const char *v, size_t n) {
  if (n == 0)
    return 0;
  if (!strncmp(v, "http://", 7) || !strncmp(v, "https://", 8) ||
      !strncmp(v, "mailto:", 7))
    return 1;
  if (v[0] == '/' && !(n > 1 && (v[1] == '/' || v[1] == '\\')))
    return 1;
  if (v[0] == '#')
    return 1;
  return 0;
}

static int src_ok(const char *v, size_t n) {
  if (n > 11 && !strncmp(v, "data:image/", 11))
    return 1;
  if (n > 12 && !strncmp(v, "/api/attach/", 12))
    return 1;
  return 0;
}

static void attr_append(cst_buf_t *b, const char *name, size_t nlen,
                        const char *val, size_t vlen) {
  cst_buf_puts(b, " ");
  cst_buf_append(b, name, nlen);
  cst_buf_puts(b, "=\"");
  /* a single-quoted or bare value may hold '"': it must not end ours */
  for (size_t i = 0; i < vlen; i++) {
    if (val[i] == '"')
      cst_buf_puts(b, "&quot;");
    else if (val[i] == '<')
      cst_buf_puts(b, "&lt;");
    else if (val[i] == '>')
      cst_buf_puts(b, "&gt;");
    else
      cst_buf_append(b, val + i, 1);
  }
  cst_buf_puts(b, "\"");
}

/* Rewrite one tag (open, close or self-closing) into out; may emit nothing
 * when the tag is not allowed. Returns 0, or -1 on OOM. */
static int sanitize_tag(const char *tag, size_t len, cst_buf_t *out) {
  /* <name ...> or </name> or <!-- ... --> */
  size_t i = 0;
  int closing = 0;
  if (len < 3 || tag[0] != '<')
    return 0;
  i = 1;
  if (i < len && tag[i] == '/') {
    closing = 1;
    i++;
  }
  size_t name_start = i;
  while (i < len && ((tag[i] >= 'a' && tag[i] <= 'z') ||
                     (tag[i] >= 'A' && tag[i] <= 'Z') ||
                     (tag[i] >= '0' && tag[i] <= '9')))
    i++;
  size_t name_len = i - name_start;
  if (name_len == 0 || !tag_allowed(tag + name_start, name_len))
    return 0; /* comments, <script>, <iframe>: dropped entirely */
  int is_img = name_len == 3 && strncmp(tag + name_start, "img", 3) == 0;
  int is_a = name_len == 1 && tag[name_start] == 'a';
  int self_closing = len > 2 && tag[len - 2] == '/';
  size_t start = out->len;
  cst_buf_puts(out, "<");
  if (closing)
    cst_buf_puts(out, "/");
  cst_buf_append(out, tag + name_start, name_len);
  if (!closing) {
    /* copy the allowed attributes */
    int keep_target = 0; /* a link is only useful with a usable href */
    size_t p = name_start + name_len;
    while (p < len && tag[p] != '>' && !(tag[p] == '/' && p + 1 < len &&
                                         tag[p + 1] == '>')) {
      while (p < len && (tag[p] == ' ' || tag[p] == '\t' || tag[p] == '\n'))
        p++;
      size_t an = p;
      while (p < len && tag[p] != '=' && tag[p] != ' ' && tag[p] != '\t' &&
             tag[p] != '\n' && tag[p] != '>' && tag[p] != '/')
        p++;
      size_t an_len = p - an;
      if (an_len == 0) {
        p++;
        continue;
      }
      const char *aval = "";
      size_t av_len = 0;
      if (p < len && tag[p] == '=') {
        p++;
        char quote = 0;
        if (p < len && (tag[p] == '"' || tag[p] == '\''))
          quote = tag[p++];
        size_t as = p;
        while (p < len &&
               (quote ? tag[p] != quote : (tag[p] != ' ' && tag[p] != '>')))
          p++;
        aval = tag + as;
        av_len = p - as;
        if (quote && p < len)
          p++; /* closing quote */
      }
      int keep = 0;
      if (is_a && an_len == 4 && strncmp(tag + an, "href", 4) == 0) {
        keep = href_ok(aval, av_len);
        keep_target = keep;
      } else if (is_img && an_len == 3 && strncmp(tag + an, "src", 3) == 0) {
        keep = src_ok(aval, av_len);
        keep_target = keep;
      } else if (is_img && an_len == 3 && strncmp(tag + an, "alt", 3) == 0)
        keep = 1;
      else if (an_len == 5 && strncmp(tag + an, "class", 5) == 0 &&
               av_len > 9 && !strncmp(aval, "language-", 9))
        keep = 1;
      if (keep)
        attr_append(out, tag + an, an_len, aval, av_len);
    }
    if (is_a && keep_target) {
      cst_buf_puts(out, " target=\"_blank\" rel=\"noopener noreferrer\"");
    }
    if ((is_a && !keep_target) || (is_img && !keep_target)) {
      /* a link with an unusable href, or an image from nowhere: drop the
       * whole tag and keep only its text */
      out->len = start;
      if (out->buf)
        out->buf[start] = 0;
      return 0;
    }
    cst_buf_puts(out, self_closing ? " /" : "");
  }
  cst_buf_puts(out, ">");
  return 0;
}

/* Emit the text before every '<' and every tag through the sanitizer. */
static void scan_chunk(cst_web_md_t *s, int final) {
  if (!s->emit || s->failed)
    return;
  size_t i = 0, emitted = 0;
  while (i < s->in.len) {
    if (s->in.buf[i] != '<') {
      i++;
      continue;
    }
    size_t close = i;
    while (close < s->in.len && s->in.buf[close] != '>')
      close++;
    if (close >= s->in.len && !final)
      break; /* the tag may still be growing: keep it buffered */
    if (i > emitted && s->emit(s->in.buf + emitted, i - emitted, s->ud) != 0)
      s->failed = 1;
    size_t tag_len = close >= s->in.len ? s->in.len - i : close - i + 1;
    cst_buf_t tag = {0};
    if (sanitize_tag(s->in.buf + i, tag_len, &tag) != 0) {
      cst_buf_free(&tag);
      s->failed = 1;
      return;
    }
    if (tag.len && s->emit(tag.buf, tag.len, s->ud) != 0)
      s->failed = 1;
    cst_buf_free(&tag);
    i = i + tag_len;
    emitted = i;
    if (s->failed)
      break;
  }
  if (emitted > 0) {
    memmove(s->in.buf, s->in.buf + emitted, s->in.len - emitted);
    s->in.len -= emitted;
    s->in.buf[s->in.len] = 0;
  }
}

/* Collector used by the one-shot sanitizer. */
typedef struct {
  cst_buf_t *b;
  int failed;
} md_collect_t;

static int collect_emit(const char *html, size_t len, void *ud) {
  md_collect_t *c = ud;
  if (cst_buf_append(c->b, html, len) != 0) {
    c->failed = 1;
    return 1;
  }
  return 0;
}

char *cst_web_md_sanitize(const char *html, size_t len) {
  cst_web_md_t tmp;
  memset(&tmp, 0, sizeof(tmp));
  cst_buf_t out = {0};
  md_collect_t col = {&out, 0};
  if (cst_buf_append(&tmp.in, html ? html : "", len) != 0)
    return NULL;
  tmp.emit = collect_emit;
  tmp.ud = &col;
  scan_chunk(&tmp, 1);
  cst_buf_free(&tmp.in);
  if (col.failed) {
    cst_buf_free(&out);
    return NULL;
  }
  return cst_buf_steal(&out);
}

/* --- sashidown glue ------------------------------------------------------- */

static int on_sashidown_html(const char *html, size_t len, void *user) {
  cst_web_md_t *s = user;
  if (cst_buf_append(&s->in, html, len) != 0) {
    s->failed = 1;
    return 1;
  }
  scan_chunk(s, 0);
  return s->failed ? 1 : 0;
}

cst_web_md_t *cst_web_md_new(cst_web_md_emit_fn emit, void *ud) {
  cst_web_md_t *s = calloc(1, sizeof(*s));
  if (!s)
    return NULL;
  sashidown_stream_opts_t opts;
  memset(&opts, 0, sizeof(opts));
  opts.on_html = on_sashidown_html;
  opts.user = s;
  s->st = sashidown_stream_create(&opts);
  if (!s->st) {
    free(s);
    return NULL;
  }
  s->emit = emit;
  s->ud = ud;
  return s;
}

int cst_web_md_feed(cst_web_md_t *s, const char *data, size_t len) {
  if (!s || !s->st)
    return -1;
  if (len && sashidown_stream_feed(s->st, data, len) != 0) {
    s->failed = 1;
    return -1;
  }
  return s->failed ? -1 : 0;
}

int cst_web_md_finish(cst_web_md_t *s) {
  if (!s || !s->st)
    return -1;
  if (sashidown_stream_finish(s->st) != 0) {
    s->failed = 1;
    return -1;
  }
  scan_chunk(s, 1);
  return s->failed ? -1 : 0;
}

void cst_web_md_free(cst_web_md_t *s) {
  if (!s)
    return;
  if (s->st)
    sashidown_stream_destroy(s->st);
  cst_buf_free(&s->in);
  free(s);
}

char *cst_web_md_html(const char *md, size_t len) {
  if (!md)
    return NULL;
  size_t cap = 65536;
  char *out = NULL;
  for (int attempt = 0; attempt < 8; attempt++) {
    char *buf = malloc(cap);
    if (!buf)
      return NULL;
    size_t out_len = 0;
    if (sashidown_convert(md, len, buf, cap, &out_len, SASHIDOWN_FMT_HTML) ==
        0) {
      out = cst_web_md_sanitize(buf, out_len);
      free(buf);
      return out;
    }
    free(buf);
    if (cap >= (size_t)8 * 1024 * 1024)
      break;
    cap *= 2;
  }
  return NULL;
}