From 4ef8c3b92a849d6a4f2d24438888967d6811dd36 Mon Sep 17 00:00:00 2001
From: Phillip Alexander <pal@jobindex.dk>
Date: Tue, 23 Jan 2024 17:19:15 +0100
Subject: [PATCH] Add PDF support (WIP)

---
 browser.go        |  8 ++++++++
 cmd/decap/main.go | 37 +++++++++++++++++++++++++------------
 query.go          | 19 +++++++++++++++++--
 3 files changed, 50 insertions(+), 14 deletions(-)

diff --git a/browser.go b/browser.go
index 73259e5..fb2d455 100644
--- a/browser.go
+++ b/browser.go
@@ -282,6 +282,14 @@ func outerHTML(out *[]string) chromedp.ActionFunc {
 	}
 }
 
+func printToPDF(buf *[]byte) chromedp.ActionFunc {
+	return func(ctx context.Context) error {
+		var err error
+		*buf, _, err = page.PrintToPDF().Do(ctx)
+		return err
+	}
+}
+
 func removeElements(sel string) chromedp.ActionFunc {
 	return func(ctx context.Context) error {
 		cmd := fmt.Sprintf("document.querySelectorAll('%s').forEach(e => e.remove());", sel)
diff --git a/cmd/decap/main.go b/cmd/decap/main.go
index cbea3bd..b95b982 100644
--- a/cmd/decap/main.go
+++ b/cmd/decap/main.go
@@ -105,13 +105,13 @@ func browseHandler(w http.ResponseWriter, req *http.Request) {
 
 	// execute query
 
+	err_status := http.StatusInternalServerError
 	var res *decap.Result
 	res, err = dec.Execute()
 	if err != nil {
 		// TODO: Propagate HTTP status properly
-		status := http.StatusInternalServerError
-		msg := fmt.Sprintf("%s: %s", http.StatusText(status), err)
-		http.Error(w, msg, status)
+		msg := fmt.Sprintf("%s: %s", http.StatusText(err_status), err)
+		http.Error(w, msg, err_status)
 		return
 	}
 
@@ -121,21 +121,34 @@ func browseHandler(w http.ResponseWriter, req *http.Request) {
 		w.Header().Set("Content-Type", "application/json")
 		err = json.NewEncoder(w).Encode(res)
 		if err != nil {
-			status := http.StatusInternalServerError
-			msg := fmt.Sprintf("%s: %s", http.StatusText(status), "Couldn't encode response")
-			http.Error(w, msg, status)
-			return
+			msg := fmt.Sprintf("%s: %s", http.StatusText(err_status), "Couldn't encode response")
+			http.Error(w, msg, err_status)
+		}
+		return
+	case "pdf":
+		w.Header().Set("Content-Type", "application/pdf")
+		_, err = w.Write(res.PDFBuffer())
+		if err != nil {
+			msg := fmt.Sprintf("%s: %s",
+				http.StatusText(err_status), "Couldn't write response bytes")
+			http.Error(w, msg, err_status)
 		}
+		return
 	case "png":
 		w.Header().Set("Content-Type", "image/png")
-		_, err := w.Write(res.ImgBuffer())
+		_, err = w.Write(res.ImgBuffer())
 		if err != nil {
-			status := http.StatusInternalServerError
 			msg := fmt.Sprintf("%s: %s",
-				http.StatusText(status), "Couldn't write response bytes")
-			http.Error(w, msg, status)
-			return
+				http.StatusText(err_status), "Couldn't write response bytes")
+			http.Error(w, msg, err_status)
 		}
+		return
+	default:
+		fmt.Fprintf(os.Stderr, `unknown results type "%s"`, res.Type())
+		msg := fmt.Sprintf(`%s: Unknown result type "%s"`,
+			http.StatusText(err_status), res.Type())
+		http.Error(w, msg, err_status)
+		 return
 	}
 }
 
diff --git a/query.go b/query.go
index e199c6e..77ba6a7 100644
--- a/query.go
+++ b/query.go
@@ -33,19 +33,28 @@ type Result struct {
 	TabID    string     `json:"tab_id"`
 	WindowID string     `json:"window_id"`
 	img      []byte
+	pdf []byte
 }
 
 func (res *Result) Type() string {
-	if len(res.img) != 0 {
+	switch {
+	case len(res.pdf) != 0:
+		return "pdf"
+	case len(res.img) != 0:
 		return "png"
+	default:
+		return "json"
 	}
-	return "json"
 }
 
 func (res *Result) ImgBuffer() []byte {
 	return res.img
 }
 
+func (res *Result) PDFBuffer() []byte {
+	return res.pdf
+}
+
 type QueryBlock struct {
 	Actions    []Action `json:"actions"`
 	Repeat     *int     `json:"repeat"`
@@ -416,6 +425,12 @@ func (r *Request) parseAction(xa Action) error {
 		}
 		r.appendActions(outerHTML(&r.res.Out[r.pos]))
 
+	case "print_to_pdf":
+		if err = xa.MustArgCount(0); err != nil {
+			return err
+		}
+		r.appendActions(printToPDF(&r.res.pdf))
+
 	case "remove":
 		if len(xa.Args()) == 0 {
 			return fmt.Errorf("remove: expected at least one argument")