package main

import (
	"archive/zip"
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"path/filepath"
	"time"

	rpc "code.justin.tv/amzn/TwitchAutoprofLambdaTwirp"
	"code.justin.tv/video/autoprof"
	"github.com/google/pprof/profile"
	"golang.org/x/sync/errgroup"
)

var (
	errNoMeta    = errors.New("meta file not found")
	errNoBundles = errors.New("no bundles found")
)

func main() {
	client, err := newBundleCollectionClient()
	if err != nil {
		log.Fatalf("%v", err)
	}

	ctx := context.Background()
	ctx, cancel := context.WithTimeout(ctx, 15*time.Minute)
	defer cancel()

	if true {
		printGoVersions(ctx, client)
		return
	}
	printStats(ctx)
}

func withFile(path string, fn func(*os.File) error) error {
	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()

	return fn(file)
}

func printStats(ctx context.Context) {
	revStacks := make(map[string]map[string]struct{})

	sc := bufio.NewScanner(os.Stdin)
	for sc.Scan() {
		basePath := sc.Text()

		var meta autoprof.ArchiveMeta
		err := withFile(filepath.Join(basePath, "meta"), func(file *os.File) error {
			return json.NewDecoder(file).Decode(&meta)
		})
		if err != nil {
			log.Fatalf("%v", err)
		}

		var prof *profile.Profile
		err = withFile(filepath.Join(basePath, "pprof/goroutine"), func(file *os.File) error {
			var err error
			prof, err = profile.Parse(file)
			return err
		})
		if err != nil {
			log.Fatalf("%v", err)
		}

		for _, s := range prof.Sample {
			var names []string
			for _, loc := range s.Location {
				for _, line := range loc.Line {
					names = append(names, line.Function.Name)
				}
			}

			j, err := json.Marshal(names)
			if err != nil {
				log.Fatalf("%v", err)
			}

			// if s.Value[0] > 100 {
			// 	log.Printf("%s", j)
			// }

			// if bytes.Equal(j, []byte(`["runtime.gopark","runtime.netpollblock","internal/poll.runtime_pollWait","internal/poll.(*pollDesc).wait","internal/poll.(*pollDesc).waitRead","internal/poll.(*FD).Read","net.(*netFD).Read","net.(*conn).Read","crypto/tls.(*atLeastReader).Read","bytes.(*Buffer).ReadFrom","crypto/tls.(*Conn).readFromUntil","crypto/tls.(*Conn).readRecordOrCCS","crypto/tls.(*Conn).readRecord","crypto/tls.(*Conn).Read","io.(*LimitedReader).Read","bufio.(*Reader).fill","bufio.(*Reader).Peek","code.justin.tv/video/nydus/core/rawhttp.(*Server).serveConn"]`)) {
			// 	log.Printf("%s %d", meta.Revision, s.Value[0])
			// }

			stacks, ok := revStacks[meta.Revision]
			if !ok {
				stacks = make(map[string]struct{})
				revStacks[meta.Revision] = stacks
			}
			stacks[string(j)] = struct{}{}

			// for each bundle, be ready to say how many have some stack
			//
			// then look through all bundles to find the set of all stacks
			//
			// then for each stack, ask each bundle about its contribution
			//
			//
		}
	}

	for rev, stacks := range revStacks {
		log.Printf("rev=%q stacks=%d", rev, len(stacks))
	}
}

func printGoVersions(ctx context.Context, c *client) {
	end := time.Now()
	start := end.Add(-15 * time.Minute)

	apps, err := c.ListApps(ctx, start, end)
	if err != nil {
		log.Fatalf("%v", err)
	}

	eg, ctx := errgroup.WithContext(ctx)
	sem := make(chan struct{}, 10)
	for _, app := range apps {
		app := app
		eg.Go(func() error {
			sem <- struct{}{}
			defer func() { <-sem }()
			goVersion, err := findGoVersion(ctx, c, app, start, end)
			if err != nil {
				if errors.Is(err, errNoBundles) {
					return nil
				}
				return err
			}

			log.Printf("app=%q go=%q", app, goVersion)
			return nil
		})
	}
	err = eg.Wait()
	if err != nil {
		log.Fatalf("%v", err)
	}
}

func findGoVersion(ctx context.Context, c *client, app string, start, end time.Time) (string, error) {
	procs, err := c.ListProcesses(ctx, app, start, end)
	if err != nil {
		return "", fmt.Errorf("ListProcesses: %w", err)
	}

	var bundles []*rpc.Bundle
	for _, b := range procs {
		bundles = append(bundles, b...)
	}

	if len(bundles) == 0 {
		return "", errNoBundles
	}

	rng := rand.New(rand.NewSource(time.Now().UnixNano()))
	body, err := c.DownloadBundle(ctx, bundles[rng.Intn(len(bundles))].GetName())
	if err != nil {
		return "", fmt.Errorf("DownloadBundle: %w", err)
	}

	meta, err := bundleMeta(body)
	if err != nil {
		return "", fmt.Errorf("bundleMeta: %w", err)
	}

	return meta.GoVersion, nil
}

func bundleMeta(data []byte) (*autoprof.ArchiveMeta, error) {
	zr, err := zip.NewReader(bytes.NewReader(data), int64(len(data)))
	if err != nil {
		return nil, fmt.Errorf("zip.NewReader: %w", err)
	}

	for _, file := range zr.File {
		if file.Name != "meta" {
			continue
		}

		metaFile, err := file.Open()
		if err != nil {
			return nil, fmt.Errorf("Open: %w", err)
		}
		metaBuf, err := ioutil.ReadAll(metaFile)
		if err != nil {
			return nil, fmt.Errorf("ReadAll: %w", err)
		}
		metaFile.Close()

		var meta autoprof.ArchiveMeta
		err = json.Unmarshal(metaBuf, &meta)
		if err != nil {
			return nil, fmt.Errorf("json.Unmarshal: %w", err)
		}

		return &meta, nil
	}

	return nil, errNoMeta
}
