// Package fcgi implements a minimal FastCGI client.
//
// Supports sending a single request over a pre-dialed net.Conn and streaming
// the CGI response back to an http.ResponseWriter.
package fcgi

import (
	"bufio"
	"encoding/binary"
	"io"
	"net/http"
	"net/textproto"
	"strings"
)

// Reference: https://fastcgi-archives.github.io/FastCGI_Specification.html

const (
	fcgiVersion       = 1
	fcgiBeginRequest  = 1
	fcgiParams        = 4
	fcgiStdin         = 5
	fcgiStdout        = 6
	fcgiEndRequest    = 3
	fcgiRoleResponder = 1
	fcgiRequestID     = 1
)

// Do sends params and the request body over conn as a FastCGI request, then
// parses the response and writes it to w. conn is closed by the caller.
func Do(w http.ResponseWriter, r *http.Request, conn io.ReadWriter, params map[string]string) error {
	bw := bufio.NewWriter(conn)

	if err := writeRecord(bw, fcgiBeginRequest, fcgiRequestID,
		[]byte{0, fcgiRoleResponder, 0, 0, 0, 0, 0, 0}); err != nil {
		return err
	}

	if err := writeRecord(bw, fcgiParams, fcgiRequestID, encodeParams(params)); err != nil {
		return err
	}
	if err := writeRecord(bw, fcgiParams, fcgiRequestID, nil); err != nil {
		return err
	}

	if r.Body != nil {
		buf := make([]byte, 4096)
		for {
			n, err := r.Body.Read(buf)
			if n > 0 {
				if werr := writeRecord(bw, fcgiStdin, fcgiRequestID, buf[:n]); werr != nil {
					return werr
				}
			}
			if err == io.EOF {
				break
			}
			if err != nil {
				return err
			}
		}
	}
	if err := writeRecord(bw, fcgiStdin, fcgiRequestID, nil); err != nil {
		return err
	}
	if err := bw.Flush(); err != nil {
		return err
	}

	// Read response
	var stdout strings.Builder
	br := bufio.NewReader(conn)

	for {
		hdr := make([]byte, 8)
		if _, err := io.ReadFull(br, hdr); err != nil {
			return err
		}
		recType := hdr[1]
		contentLen := int(binary.BigEndian.Uint16(hdr[4:6]))
		paddingLen := int(hdr[6])

		content := make([]byte, contentLen)
		if _, err := io.ReadFull(br, content); err != nil {
			return err
		}
		if paddingLen > 0 {
			if _, err := io.ReadFull(br, make([]byte, paddingLen)); err != nil {
				return err
			}
		}

		switch recType {
		case fcgiStdout:
			stdout.Write(content)
		case fcgiEndRequest:
			return forwardResponse(w, stdout.String())
		}
		// fcgiStderr is silently discarded
	}
}

func forwardResponse(w http.ResponseWriter, raw string) error {
	sep := "\r\n\r\n"
	idx := strings.Index(raw, sep)
	advance := 4
	if idx == -1 {
		sep = "\n\n"
		idx = strings.Index(raw, sep)
		advance = 2
	}
	if idx == -1 {
		w.Write([]byte(raw))
		return nil
	}

	headerPart := raw[:idx]
	body := raw[idx+advance:]

	tp := textproto.NewReader(bufio.NewReader(strings.NewReader(headerPart + "\r\n\r\n")))
	mime, _ := tp.ReadMIMEHeader()

	status := 200
	if s := mime.Get("Status"); s != "" {
		code, _, _ := strings.Cut(s, " ")
		if n := parseIntFast(code); n > 0 {
			status = n
		}
		mime.Del("Status")
	}

	for k, vs := range mime {
		for _, v := range vs {
			w.Header().Add(k, v)
		}
	}
	w.WriteHeader(status)
	w.Write([]byte(body))
	return nil
}

func writeRecord(w io.Writer, recType uint8, reqID uint16, content []byte) error {
	length := len(content)
	padding := (8 - (length % 8)) % 8
	hdr := []byte{
		fcgiVersion, recType,
		byte(reqID >> 8), byte(reqID),
		byte(length >> 8), byte(length),
		byte(padding), 0,
	}
	if _, err := w.Write(hdr); err != nil {
		return err
	}
	if len(content) > 0 {
		if _, err := w.Write(content); err != nil {
			return err
		}
	}
	if padding > 0 {
		if _, err := w.Write(make([]byte, padding)); err != nil {
			return err
		}
	}
	return nil
}

func encodeParams(params map[string]string) []byte {
	var buf []byte
	for k, v := range params {
		buf = appendLen(buf, len(k))
		buf = appendLen(buf, len(v))
		buf = append(buf, k...)
		buf = append(buf, v...)
	}
	return buf
}

func appendLen(buf []byte, n int) []byte {
	if n <= 127 {
		return append(buf, byte(n))
	}
	return append(buf, byte((n>>24)|0x80), byte(n>>16), byte(n>>8), byte(n))
}

func parseIntFast(s string) int {
	n := 0
	for i := 0; i < len(s); i++ {
		c := s[i]
		if c < '0' || c > '9' {
			return 0
		}
		n = n*10 + int(c-'0')
	}
	return n
}