package geoip

import (
	"archive/tar"
	"compress/gzip"
	"errors"
	"fmt"
	"io"
	"path"
	"strings"
)

// ErrTooLarge is a capped reader refusing to read past its cap (spec §10.3).
var ErrTooLarge = errors.New("geoip archive: larger than its cap")

type cappedReader struct {
	r    io.Reader
	left int64
}

func (c *cappedReader) Read(p []byte) (int, error) {
	if c.left < 0 {
		return 0, ErrTooLarge
	}
	if int64(len(p)) > c.left+1 {
		p = p[:c.left+1]
	}
	n, err := c.r.Read(p)
	c.left -= int64(n)
	if c.left < 0 {
		return n, ErrTooLarge
	}
	return n, err
}

// CappedReader reads r until more than max bytes have come through, then
// fails with ErrTooLarge. Exactly max bytes read cleanly.
func CappedReader(r io.Reader, max int64) io.Reader { return &cappedReader{r: r, left: max} }

// GunzipReader expands a gzip stream, refusing to expand past maxExpanded:
// a decompression bomb stops at the cap, never at the end of its expansion.
func GunzipReader(r io.Reader, maxExpanded int64) (io.Reader, error) {
	zr, err := gzip.NewReader(r)
	if err != nil {
		return nil, fmt.Errorf("geoip archive: not gzip: %w", err)
	}
	return CappedReader(zr, maxExpanded), nil
}

// TarGzMMDB returns the first regular *.mmdb member of a .tar.gz (MaxMind's
// download format). Every expanded byte counts toward maxExpanded,
// including the members walked past, so junk padding cannot fill a disk.
func TarGzMMDB(r io.Reader, maxExpanded int64) (io.Reader, error) {
	expanded, err := GunzipReader(r, maxExpanded)
	if err != nil {
		return nil, err
	}
	tr := tar.NewReader(expanded)
	for {
		h, err := tr.Next()
		if err == io.EOF {
			return nil, errors.New("geoip archive: no .mmdb member")
		}
		if err != nil {
			return nil, err
		}
		if h.Typeflag == tar.TypeReg && strings.HasSuffix(path.Base(h.Name), ".mmdb") {
			return tr, nil
		}
	}
}