package backup

import (
	"archive/tar"
	"compress/gzip"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"
)

// ErrChecksum is returned when an archive member does not match the checksum
// the manifest recorded for it.
var ErrChecksum = errors.New("archive member failed its checksum")

// Writer streams a tar.gz. Members are written in the order they are produced
// and the manifest goes last, because its per-collection checksums are only
// known once every collection has been written.
type Writer struct {
	gz  *gzip.Writer
	tar *tar.Writer
}

// NewWriter starts an archive on out. out may be a file or stdout; nothing here
// seeks.
func NewWriter(out io.Writer) *Writer {
	gz := gzip.NewWriter(out)
	return &Writer{gz: gz, tar: tar.NewWriter(gz)}
}

func (w *Writer) writeMember(name string, body []byte) error {
	h := &tar.Header{
		Name:     name,
		Mode:     0o600,
		Size:     int64(len(body)),
		ModTime:  time.Now().UTC(),
		Typeflag: tar.TypeReg,
	}
	if err := w.tar.WriteHeader(h); err != nil {
		return fmt.Errorf("write header %s: %w", name, err)
	}
	if _, err := w.tar.Write(body); err != nil {
		return fmt.Errorf("write %s: %w", name, err)
	}
	return nil
}

// WriteCollection writes the concatenated raw BSON of one collection and
// returns the manifest entry describing it.
func (w *Writer) WriteCollection(name string, docs [][]byte) (CollectionEntry, error) {
	var body []byte
	for _, d := range docs {
		body = append(body, d...)
	}
	sum := sha256.Sum256(body)
	entry := CollectionEntry{
		Name:      name,
		Documents: int64(len(docs)),
		Bytes:     int64(len(body)),
		SHA256:    hex.EncodeToString(sum[:]),
	}
	if err := w.writeMember(collectionMember(name), body); err != nil {
		return CollectionEntry{}, err
	}
	return entry, nil
}

// WriteIndexes writes a collection's index specifications verbatim.
func (w *Writer) WriteIndexes(name string, specsJSON []byte) error {
	return w.writeMember(indexMember(name), specsJSON)
}

// Close writes the manifest and finishes the archive.
func (w *Writer) Close(m Manifest) error {
	raw, err := json.MarshalIndent(m, "", "  ")
	if err != nil {
		return fmt.Errorf("marshal manifest: %w", err)
	}
	if err := w.writeMember(ManifestName, raw); err != nil {
		return err
	}
	if err := w.tar.Close(); err != nil {
		return err
	}
	return w.gz.Close()
}

func collectionMember(name string) string { return "collections/" + name + ".bson" }
func indexMember(name string) string      { return "indexes/" + name + ".json" }

// Reader is an opened archive.
//
// Open extracts to a temporary directory rather than streaming, because gzip
// offers no random access and the manifest — which carries the checksums every
// other member is judged against — is written last. Verifying before writing a
// single document to the target is worth one pass over local disk. This is why
// the container image needs a /tmp.
type Reader struct {
	dir      string
	manifest Manifest
}

// Open extracts, verifies and returns the archive at path. The caller must
// Close it.
func Open(archivePath string) (*Reader, error) {
	dir, err := os.MkdirTemp("", "vantage-restore-*")
	if err != nil {
		return nil, fmt.Errorf("temp dir: %w", err)
	}
	r := &Reader{dir: dir}

	if err := r.extract(archivePath); err != nil {
		r.Close()
		return nil, err
	}
	if err := r.loadManifest(); err != nil {
		r.Close()
		return nil, err
	}
	if err := r.verifyMembers(); err != nil {
		r.Close()
		return nil, err
	}
	return r, nil
}

func (r *Reader) extract(archivePath string) error {
	f, err := os.Open(archivePath)
	if err != nil {
		return fmt.Errorf("open archive: %w", err)
	}
	defer f.Close()

	gz, err := gzip.NewReader(f)
	if err != nil {
		return fmt.Errorf("archive is not gzip: %w", err)
	}
	defer gz.Close()

	tr := tar.NewReader(gz)
	for {
		h, err := tr.Next()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return fmt.Errorf("read archive: %w", err)
		}
		if h.Typeflag != tar.TypeReg {
			continue
		}
		dest, err := safeJoin(r.dir, h.Name)
		if err != nil {
			return err
		}
		if err := os.MkdirAll(filepath.Dir(dest), 0o700); err != nil {
			return fmt.Errorf("mkdir: %w", err)
		}
		out, err := os.OpenFile(dest, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
		if err != nil {
			return fmt.Errorf("create %s: %w", h.Name, err)
		}
		if _, err := io.Copy(out, tr); err != nil {
			out.Close()
			return fmt.Errorf("extract %s: %w", h.Name, err)
		}
		if err := out.Close(); err != nil {
			return err
		}
	}
}

// safeJoin refuses a member name that escapes the extraction directory. An
// archive is operator-supplied input and may not be one we wrote.
func safeJoin(dir, name string) (string, error) {
	clean := path.Clean("/" + name)
	dest := filepath.Join(dir, filepath.FromSlash(strings.TrimPrefix(clean, "/")))
	if !strings.HasPrefix(dest, filepath.Clean(dir)+string(os.PathSeparator)) {
		return "", fmt.Errorf("archive member %q escapes the extraction directory", name)
	}
	return dest, nil
}

func (r *Reader) loadManifest() error {
	raw, err := os.ReadFile(filepath.Join(r.dir, ManifestName))
	if err != nil {
		return fmt.Errorf("archive has no %s: %w", ManifestName, err)
	}
	if err := json.Unmarshal(raw, &r.manifest); err != nil {
		return fmt.Errorf("parse %s: %w", ManifestName, err)
	}
	return r.manifest.Check()
}

func (r *Reader) verifyMembers() error {
	for _, c := range r.manifest.Collections {
		f, err := os.Open(filepath.Join(r.dir, collectionMember(c.Name)))
		if err != nil {
			return fmt.Errorf("%w: %s is named in the manifest but absent from the archive",
				ErrChecksum, c.Name)
		}
		h := sha256.New()
		n, err := io.Copy(h, f)
		f.Close()
		if err != nil {
			return fmt.Errorf("read %s: %w", c.Name, err)
		}
		if n != c.Bytes {
			return fmt.Errorf("%w: %s is %d bytes, manifest says %d", ErrChecksum, c.Name, n, c.Bytes)
		}
		if got := hex.EncodeToString(h.Sum(nil)); got != c.SHA256 {
			return fmt.Errorf("%w: %s checksum %s, manifest says %s", ErrChecksum, c.Name, got, c.SHA256)
		}
	}
	return nil
}

// Manifest returns the verified manifest.
func (r *Reader) Manifest() Manifest { return r.manifest }

// OpenCollection returns the raw BSON stream for one collection.
func (r *Reader) OpenCollection(name string) (io.ReadCloser, error) {
	return os.Open(filepath.Join(r.dir, collectionMember(name)))
}

// IndexesJSON returns a collection's index specifications, or nil when the
// archive holds none. A collection with no indexes beyond _id_ is ordinary and
// is not an error.
func (r *Reader) IndexesJSON(name string) ([]byte, error) {
	raw, err := os.ReadFile(filepath.Join(r.dir, indexMember(name)))
	if os.IsNotExist(err) {
		return nil, nil
	}
	return raw, err
}

// Close removes the extraction directory.
func (r *Reader) Close() error { return os.RemoveAll(r.dir) }