diff --git a/main.c b/main.c
index 21f9755..04eda12 100644
--- a/main.c
+++ b/main.c
@@ -19,10 +19,9 @@ enum {
 };
 
 enum {
-    CODE_OPEN       = 1 << 0,
-    CODE_PRINT      = 1 << 1,
-    CODE_LINES      = 1 << 2,
-    CODE_HIGHLIGHT  = 1 << 3,
+    CODE_LINES      = (1 << 0),
+    CODE_HIGHLIGHT  = (1 << 1),
+    CODE_BUFFERED   = (CODE_LINES | CODE_HIGHLIGHT),
 };
 
 struct ctx {
@@ -30,6 +29,11 @@ struct ctx {
         int flags;
         int from;
     } code;
+    struct {
+        char *str;
+        size_t len;
+        FILE *f;
+    } buf;
 };
 
 static int enter_block(MD_BLOCKTYPE type, void *details, void *ctx_)
@@ -71,48 +75,55 @@ static int enter_block(MD_BLOCKTYPE type, void *details, void *ctx_)
                 struct ctx *ctx = ctx_;
                 MD_BLOCK_CODE_DETAIL *d = details;
 
-                ctx->code.flags = CODE_OPEN;
-                if (!d->info.text) {
-                    ctx->code.flags |= CODE_HIGHLIGHT;
+                if (ctx->code.flags)
+                    return ERR_INVALID;
+
+                if (!d->info.text || !strncmp(d->info.text, "pre\n", 4)) {
+                    printf("<pre>");
                 } else {
                     const char *t = d->info.text;
                     size_t l = d->info.size;
-                    if (strncmp(t, "pre", 3) || (t[4] != ' ' && t[4] != '\n')) {
-                        ctx->code.flags |= CODE_LINES;
+
+                    if (strncmp(t, "pre ", 4)) {
+                        ctx->code.flags = CODE_LINES;
                         ctx->code.from = 1;
+                    }
 
-                        const char *s;
-                        s = memchr(t, ' ', l);
-                        if (!s) break;
-
-                        ++s;
-                        l -= s - t;
-                        t = s;
-                        while (l > 0) {
-                            if (l > 5 && !strncmp(t, "from=", 5)) {
-                                char *end = NULL;
-                                ctx->code.from = strtol(t + 5, &end, 10);
-                                if (!end || (*end != ' ' && *end != '\n'))
-                                    return ERR_INVALID;
-                                l -= end - t;
-                                t = end;
-                            } else if (l >= 2 && !strncmp(t, "hl", 2)
-                                    && (s[2] == ' ' || s[2] == '\n')) {
-                                ctx->code.flags |= CODE_HIGHLIGHT;
-                                l -= 2;
-                                t += 2;
-                            } else {
-                                s = memchr(t, ' ', l);
-                                if (!s) break;
-                                l -= s - t;
-                                t = s;
-                            }
-                            while (*t == ' ' && l > 0) {
-                                ++t;
-                                --l;
-                            }
+                    const char *s;
+                    s = memchr(t, ' ', l);
+                    if (!s) break;
+
+                    ++s;
+                    l -= s - t;
+                    t = s;
+                    while (l > 0) {
+                        if ((ctx->code.flags & CODE_LINES)
+                                && l > 5 && !strncmp(t, "from=", 5)) {
+                            char *end = NULL;
+                            ctx->code.from = strtol(t + 5, &end, 10);
+                            if (!end || (*end != ' ' && *end != '\n'))
+                                return ERR_INVALID;
+                            l -= end - t;
+                            t = end;
+                        } else if (l >= 2 && !strncmp(t, "hl", 2)
+                                && (s[2] == ' ' || s[2] == '\n')) {
+                            ctx->code.flags |= CODE_HIGHLIGHT;
+                            l -= 2;
+                            t += 2;
+                        } else {
+                            s = memchr(t, ' ', l);
+                            if (!s) break;
+                            l -= s - t;
+                            t = s;
+                        }
+                        while (*t == ' ' && l > 0) {
+                            ++t;
+                            --l;
                         }
                     }
+
+                    if (!(ctx->code.flags & CODE_BUFFERED))
+                        printf("<pre>");
                 }
             }
             break;
@@ -173,15 +184,27 @@ static int leave_block(MD_BLOCKTYPE type, void *details, void *ctx_)
             {
                 struct ctx *ctx = ctx_;
 
-                if (!(ctx->code.flags & CODE_OPEN))
-                    return ERR_INVALID;
-                if (ctx->code.flags & CODE_PRINT) {
+                if (!(ctx->code.flags & CODE_BUFFERED)) {
+                    printf("</pre>");
+                } else {
+                    if (fclose(ctx->buf.f)) return ERR_MEM;
+                    ctx->buf.f = NULL;
+
+                    if (ctx->code.flags & CODE_LINES) {
+                        printf("<pre class=\"lineno\">%d...</pre>", ctx->code.from);
+                        printf("<code>");
+                    } else {
+                        printf("<pre>");
+                    }
+
+                    fwrite(ctx->buf.str, 1, ctx->buf.len, stdout);
+
                     if (ctx->code.flags & CODE_LINES)
                         printf("</code>");
                     else
                         printf("</pre>");
+                    ctx->code.flags = 0;
                 }
-                ctx->code.flags = 0;
             }
             break;
 
@@ -361,6 +384,8 @@ static int escape_text(const char *text, size_t size)
 
 static int highlight_code_text(const char *text, size_t size, void *ctx_)
 {
+    struct ctx *ctx = ctx_;
+    FILE *f = (ctx->code.flags & CODE_BUFFERED) ? ctx->buf.f : stdout;
     const char *end = text + size;
 
     for ( ; text < end; ) {
@@ -369,7 +394,7 @@ static int highlight_code_text(const char *text, size_t size, void *ctx_)
         if (!open) break;
         close = memmem(open + 4, end - open - 4, "</hl>", 5);
         if (!close) break;
-        printf("%.*s%s%.*s%s",
+        fprintf(f, "%.*s%s%.*s%s",
                (int) (open - text), text,
                "<span class=\"highlighted\">",
                (int) (close - open - 4), open + 4,
@@ -377,7 +402,7 @@ static int highlight_code_text(const char *text, size_t size, void *ctx_)
         text = close + 5;
         size = end - text;
     }
-    printf("%.*s", (int) size, text);
+    fprintf(f, "%.*s", (int) size, text);
     return 0;
 }
 
@@ -406,26 +431,20 @@ static int text(MD_TEXTTYPE type, const MD_CHAR *text, MD_SIZE size, void *ctx_)
             {
                 struct ctx *ctx = ctx_;
 
-                if (!(ctx->code.flags & CODE_OPEN)) {
-                    /* inline code */
+                if (!(ctx->code.flags & CODE_BUFFERED)) {
                     printf("%.*s", size, text);
                     break;
                 }
 
-                if (!(ctx->code.flags & CODE_PRINT)) {
-                    if (ctx->code.flags & CODE_LINES) {
-                        printf("<pre class=\"lineno\">%d...</pre>", ctx->code.from);
-                        printf("<code>");
-                    } else {
-                        printf("<pre>");
-                    }
-                    ctx->code.flags |= CODE_PRINT;
+                if (!ctx->buf.f) {
+                    ctx->buf.f = open_memstream(&ctx->buf.str, &ctx->buf.len);
+                    if (!ctx->buf.f) return ERR_MEM;
                 }
 
                 if (ctx->code.flags & CODE_HIGHLIGHT)
                     return highlight_code_text(text, size, ctx_);
                 else
-                    printf("%.*s", size, text);
+                    fprintf(ctx->buf.f, "%.*s", size, text);
                 break;
             }