package main

import (
	"archive/zip"
	"flag"
	"io"
	"log"
	"net/url"
	"os"

	"github.com/google/pprof/profile"
)

func main() {
	var (
		output   = flag.String("o", "/dev/null", "Output filename")
		profName = flag.String("profile", "profile", "Name of profile to merge")
	)
	flag.Parse()

	out, err := os.OpenFile(*output, os.O_CREATE|os.O_WRONLY, 0666)
	if err != nil {
		log.Fatalf("%v", err)
	}
	defer func() {
		err := out.Close()
		if err != nil {
			log.Fatalf("%v", err)
		}
	}()

	var profs []*profile.Profile
	for _, filename := range flag.Args() {
		prof, err := readFilebyName(filename, *profName)
		if err != nil {
			if os.IsNotExist(err) {
				continue
			}
			log.Fatalf("%v", err)
		}
		profs = append(profs, prof)
	}

	prof, err := profile.Merge(profs)
	if err != nil {
		log.Fatalf("%v", err)
	}
	err = prof.Write(out)
	if err != nil {
		log.Fatalf("%v", err)
	}
}

func readFilebyName(bundle string, profName string) (*profile.Profile, error) {
	f, err := os.Open(bundle)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	info, err := f.Stat()
	if err != nil {
		return nil, err
	}

	return readProfile(f, info.Size(), profName)
}

func readProfile(r io.ReaderAt, size int64, name string) (*profile.Profile, error) {
	zr, err := zip.NewReader(r, size)
	if err != nil {
		return nil, err
	}

	path := "pprof/" + url.PathEscape(name)

	for _, file := range zr.File {
		if file.Name == path {
			rc, err := file.Open()
			if err != nil {
				return nil, err
			}
			defer rc.Close()

			return profile.Parse(rc)
		}
	}

	return nil, os.ErrNotExist
}
