package main

import (
	"bytes"
	"io"

	"git.sr.ht/~adnano/go-gemini"
)

// ResponseWriter implements a gemini.ResponseWriter that constructs
// a Gemini response.
type ResponseWriter struct {
	status      gemini.Status
	meta        string
	body        bytes.Buffer
	mediatype   string
	wroteHeader bool
	bodyAllowed bool
}

// SetMediaType implements the gemini.ResponseWriter interface.
func (w *ResponseWriter) SetMediaType(mediatype string) {
	w.mediatype = mediatype
}

// Write implements the gemini.ResponseWriter interface.
func (w *ResponseWriter) Write(b []byte) (int, error) {
	if !w.wroteHeader {
		meta := w.mediatype
		if meta == "" {
			// Use default media type
			meta = "text/gemini; charset=utf-8"
		}
		w.WriteHeader(gemini.StatusSuccess, meta)
	}
	if !w.bodyAllowed {
		return 0, gemini.ErrBodyNotAllowed
	}
	return w.body.Write(b)
}

// WriteHeader implements the gemini.ResponseWriter interface.
func (w *ResponseWriter) WriteHeader(status gemini.Status, meta string) {
	if w.wroteHeader {
		return
	}

	if status.Class() == gemini.StatusSuccess {
		w.bodyAllowed = true
	}

	w.status = status
	w.meta = meta
	w.wroteHeader = true
}

// Response returns the constructed response.
func (w *ResponseWriter) Response() *gemini.Response {
	return &gemini.Response{
		Status: w.status,
		Meta:   w.meta,
		Body:   io.NopCloser(&w.body),
	}
}

func (w *ResponseWriter) Flush() error {
	return nil
}