package gmi

import (
	"bytes"
	"context"
	"io"
	"sync"

	"git.sr.ht/~adnano/go-gemini"
	"github.com/mediocregopher/blog.mediocregopher.com/srv/cache"
)

type cacheRW struct {
	gemini.ResponseWriter
	status    gemini.Status
	mediaType string
	buf       []byte
}

func (c *cacheRW) SetMediaType(mediaType string) {
	c.mediaType = mediaType
	c.ResponseWriter.SetMediaType(mediaType)
}

func (c *cacheRW) WriteHeader(status gemini.Status, meta string) {
	c.status = status
	c.ResponseWriter.WriteHeader(status, meta)
}

func (c *cacheRW) Write(b []byte) (int, error) {
	c.buf = append(c.buf, b...)
	return c.ResponseWriter.Write(b)
}

func cacheMiddleware(cache cache.Cache) func(h gemini.Handler) gemini.Handler {

	type entry struct {
		mediaType string
		body      []byte
	}

	pool := sync.Pool{
		New: func() interface{} { return new(bytes.Reader) },
	}

	return func(h gemini.Handler) gemini.Handler {
		return gemini.HandlerFunc(func(
			ctx context.Context,
			rw gemini.ResponseWriter,
			r *gemini.Request,
		) {

			id := r.URL.String()

			if value := cache.Get(id); value != nil {

				entry := value.(entry)

				if entry.mediaType != "" {
					rw.SetMediaType(entry.mediaType)
				}

				reader := pool.Get().(*bytes.Reader)
				defer pool.Put(reader)

				reader.Reset(entry.body)

				io.Copy(rw, reader)
				return
			}

			cacheRW := &cacheRW{
				ResponseWriter: rw,
				status:         gemini.StatusSuccess,
			}

			h.ServeGemini(ctx, cacheRW, r)

			if cacheRW.status == gemini.StatusSuccess {
				cache.Set(id, entry{
					mediaType: cacheRW.mediaType,
					body:      cacheRW.buf,
				})
			}
		})
	}
}