// Copyright 2017 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package maintner

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/golang/protobuf/proto"
	"golang.org/x/build/maintner/maintpb"
	"golang.org/x/build/maintner/reclog"
)

// NewNetworkMutationSource returns a mutation source from a master server.
// The server argument should be a URL to the JSON logs index.
func NewNetworkMutationSource(server, cacheDir string) MutationSource {
	base, err := url.Parse(server)
	if err != nil {
		panic(fmt.Sprintf("invalid URL: %q", server))
	}
	return &netMutSource{
		server:   server,
		base:     base,
		cacheDir: cacheDir,
	}
}

type netMutSource struct {
	server   string
	base     *url.URL
	cacheDir string

	last []fileSeg

	// Hooks for testing. If nil, unused:
	testHookGetServerSegments      func(context.Context, int64) ([]LogSegmentJSON, error)
	testHookWaitAfterServerDupData func(context.Context) error
	testHookSyncSeg                func(context.Context, LogSegmentJSON) (fileSeg, error)
	testHookFilePrefixSum224       func(file string, n int64) string
}

func (ns *netMutSource) GetMutations(ctx context.Context) <-chan MutationStreamEvent {
	ch := make(chan MutationStreamEvent, 50)
	go func() {
		err := ns.sendMutations(ctx, ch)
		final := MutationStreamEvent{Err: err}
		if err == nil {
			final.End = true
		}
		select {
		case ch <- final:
		case <-ctx.Done():
		}
	}()
	return ch
}

// isNoInternetError reports whether the provided error is because there's no
// network connectivity.
func isNoInternetError(err error) bool {
	switch err := err.(type) {
	case *url.Error:
		return isNoInternetError(err.Err)
	case *net.OpError:
		return isNoInternetError(err.Err)
	case *net.DNSError:
		// Trashy:
		return err.Err == "no such host"
	default:
		log.Printf("Unknown error type %T: %#v", err, err)
		return false
	}
}

func (ns *netMutSource) locallyCachedSegments() (segs []fileSeg, err error) {
	defer func() {
		if err != nil {
			log.Printf("No network connection and failed to use local cache: %v", err)
		} else {
			log.Printf("No network connection; using %d locally cached segments.", len(segs))
		}
	}()
	fis, err := ioutil.ReadDir(ns.cacheDir)
	if err != nil {
		return nil, err
	}
	fiMap := map[string]os.FileInfo{}
	segHex := map[int]string{}
	segGrowing := map[int]bool{}
	for _, fi := range fis {
		name := fi.Name()
		if !strings.HasSuffix(name, ".mutlog") {
			continue
		}
		fiMap[name] = fi

		if len(name) == len("0000.6897fab4d3afcda332424b2a2a1a4469021074282bc7be5606aaa221.mutlog") {
			num, err := strconv.Atoi(name[:4])
			if err != nil {
				continue
			}
			segHex[num] = strings.TrimSuffix(name[5:], ".mutlog")
		} else if strings.HasSuffix(name, ".growing.mutlog") {
			num, err := strconv.Atoi(name[:4])
			if err != nil {
				continue
			}
			segGrowing[num] = true
		}
	}
	for num := 0; ; num++ {
		if hex, ok := segHex[num]; ok {
			name := fmt.Sprintf("%04d.%s.mutlog", num, hex)
			segs = append(segs, fileSeg{
				seg:    num,
				file:   filepath.Join(ns.cacheDir, name),
				size:   fiMap[name].Size(),
				sha224: hex,
			})
			continue
		}
		if segGrowing[num] {
			name := fmt.Sprintf("%04d.growing.mutlog", num)
			slurp, err := ioutil.ReadFile(filepath.Join(ns.cacheDir, name))
			if err != nil {
				return nil, err
			}
			segs = append(segs, fileSeg{
				seg:    num,
				file:   filepath.Join(ns.cacheDir, name),
				size:   int64(len(slurp)),
				sha224: fmt.Sprintf("%x", sha256.Sum224(slurp)),
			})
		}
		return segs, nil
	}
}

// waitSizeNot optionally specifies that the request should long-poll waiting for the server
// to have a sum of log segment sizes different than the value specified.
func (ns *netMutSource) getServerSegments(ctx context.Context, waitSizeNot int64) ([]LogSegmentJSON, error) {
	if fn := ns.testHookGetServerSegments; fn != nil {
		return fn(ctx, waitSizeNot)
	}

	logsURL := ns.server
	if waitSizeNot > 0 {
		logsURL += fmt.Sprintf("?waitsizenot=%d", waitSizeNot)
	}
	for {
		req, err := http.NewRequest("GET", logsURL, nil)
		if err != nil {
			return nil, err
		}
		req = req.WithContext(ctx)
		res, err := http.DefaultClient.Do(req)
		if err != nil {
			return nil, err
		}
		// If we're doing a long poll and the server replies
		// with a 304 response, that means the server is just
		// heart-beating us and trying to get a response back
		// within its various deadlines. But we should just
		// try again.
		if waitSizeNot > 0 && res.StatusCode == http.StatusNotModified {
			res.Body.Close()
			continue
		}
		defer res.Body.Close()
		if res.StatusCode != 200 {
			return nil, fmt.Errorf("%s: %v", ns.server, res.Status)
		}
		var segs []LogSegmentJSON
		err = json.NewDecoder(res.Body).Decode(&segs)
		if err != nil {
			return nil, fmt.Errorf("decoding %s JSON: %v", ns.server, err)
		}
		return segs, nil
	}
}

func (ns *netMutSource) getNewSegments(ctx context.Context) ([]fileSeg, error) {
	for {
		sumLast := sumSegSize(ns.last)

		segs, err := ns.getServerSegments(ctx, sumLast)
		if isNoInternetError(err) {
			if sumLast == 0 {
				return ns.locallyCachedSegments()
			}
			log.Printf("No internet; blocking.")
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case <-time.After(15 * time.Second):
				continue
			}
		}
		if err != nil {
			return nil, err
		}
		// TODO: optimization: if already on GCE, skip sync to disk part and just
		// read from network. fast & free network inside.

		var fileSegs []fileSeg
		for _, seg := range segs {
			fileSeg, err := ns.syncSeg(ctx, seg)
			if err != nil {
				return nil, fmt.Errorf("syncing segment %d: %v", seg.Number, err)
			}
			fileSegs = append(fileSegs, fileSeg)
		}
		sumCommon := ns.sumCommonPrefixSize(fileSegs, ns.last)
		if sumLast != sumCommon {
			return nil, ErrSplit
		}
		sumCur := sumSegSize(fileSegs)
		if sumCommon == sumCur {
			// Nothing new. This shouldn't happen once the
			// server is updated to respect the
			// "?waitsizenot=NNN" long polling parameter.
			// But keep this brief pause as a backup to
			// prevent spinning and because clients &
			// servers won't be updated simultaneously.
			if ns.testHookGetServerSegments == nil {
				log.Printf("maintner.netsource: server returned unchanged log segments; old server?")
			}
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case <-time.After(1 * time.Second):
			}
			continue
		}
		ns.last = fileSegs

		newSegs := trimLeadingSegBytes(fileSegs, sumCommon)
		return newSegs, nil
	}
}

func trimLeadingSegBytes(in []fileSeg, trim int64) []fileSeg {
	// First trim off whole segments, sharing the same underlying memory.
	for len(in) > 0 && trim >= in[0].size {
		trim -= in[0].size
		in = in[1:]
	}
	if len(in) == 0 {
		return nil
	}
	// Now copy, since we'll be modifying the first element.
	out := append([]fileSeg(nil), in...)
	out[0].skip = trim
	return out
}

// filePrefixSum224 returns the lowercase hex SHA-224 of the first n bytes of file.
func (ns *netMutSource) filePrefixSum224(file string, n int64) string {
	if fn := ns.testHookFilePrefixSum224; fn != nil {
		return fn(file, n)
	}
	f, err := os.Open(file)
	if err != nil {
		if !os.IsNotExist(err) {
			log.Print(err)
		}
		return ""
	}
	defer f.Close()
	h := sha256.New224()
	_, err = io.CopyN(h, f, n)
	if err != nil {
		log.Print(err)
		return ""
	}
	return fmt.Sprintf("%x", h.Sum(nil))
}

func sumSegSize(segs []fileSeg) (sum int64) {
	for _, seg := range segs {
		sum += seg.size
	}
	return
}

func (ns *netMutSource) sumCommonPrefixSize(a, b []fileSeg) (sum int64) {
	for len(a) > 0 && len(b) > 0 {
		sa, sb := a[0], b[0]
		if sa.sha224 == sb.sha224 {
			// Whole chunk in common.
			sum += sa.size
			a, b = a[1:], b[1:]
			continue
		}
		if sa.size == sb.size {
			// If they're the same size but different
			// sums, it must've forked.
			return
		}
		// See if one chunk is a prefix of the other.
		// Make sa be the smaller one.
		if sb.size < sa.size {
			sa, sb = sb, sa
		}
		// Hash the beginning of the bigger size.
		bPrefixSum := ns.filePrefixSum224(sb.file, sa.size)
		if bPrefixSum == sa.sha224 {
			sum += sa.size
		}
		break
	}
	return
}

func (ns *netMutSource) sendMutations(ctx context.Context, ch chan<- MutationStreamEvent) error {
	newSegs, err := ns.getNewSegments(ctx)
	if err != nil {
		return err
	}
	return foreachFileSeg(newSegs, func(seg fileSeg) error {
		f, err := os.Open(seg.file)
		if err != nil {
			return err
		}
		defer f.Close()
		if seg.skip > 0 {
			if _, err := f.Seek(seg.skip, io.SeekStart); err != nil {
				return err
			}
		}
		return reclog.ForeachRecord(io.LimitReader(f, seg.size-seg.skip), seg.skip, func(off int64, hdr, rec []byte) error {
			m := new(maintpb.Mutation)
			if err := proto.Unmarshal(rec, m); err != nil {
				return err
			}
			select {
			case ch <- MutationStreamEvent{Mutation: m}:
				return nil
			case <-ctx.Done():
				return ctx.Err()
			}
		})
	})
}

func foreachFileSeg(segs []fileSeg, fn func(seg fileSeg) error) error {
	for _, seg := range segs {
		if err := fn(seg); err != nil {
			return err
		}
	}
	return nil
}

// TODO: add a constructor for this? or simplify it. make it Size +
// File + embedded LogSegmentJSON?
type fileSeg struct {
	seg    int
	file   string // full path
	sha224 string
	skip   int64
	size   int64
}

func (ns *netMutSource) syncSeg(ctx context.Context, seg LogSegmentJSON) (fileSeg, error) {
	if fn := ns.testHookSyncSeg; fn != nil {
		return fn(ctx, seg)
	}

	isFinalSeg := !strings.HasPrefix(seg.URL, "https://storage.googleapis.com/")
	relURL, err := url.Parse(seg.URL)
	if err != nil {
		return fileSeg{}, err
	}
	segURL := ns.base.ResolveReference(relURL)

	frozen := filepath.Join(ns.cacheDir, fmt.Sprintf("%04d.%s.mutlog", seg.Number, seg.SHA224))

	// Do we already have it? Files named in their final form with the sha224 are considered
	// complete and immutable.
	if fi, err := os.Stat(frozen); err == nil && fi.Size() == seg.Size {
		return fileSeg{seg: seg.Number, file: frozen, size: fi.Size(), sha224: seg.SHA224}, nil
	}

	// See how much data we already have in the partial growing file.
	partial := filepath.Join(ns.cacheDir, fmt.Sprintf("%04d.growing.mutlog", seg.Number))
	have, _ := ioutil.ReadFile(partial)
	if int64(len(have)) == seg.Size {
		got224 := fmt.Sprintf("%x", sha256.Sum224(have))
		if got224 == seg.SHA224 {
			if !isFinalSeg {
				// This was growing for us, but the server started a new growing segment.
				if err := os.Rename(partial, frozen); err != nil {
					return fileSeg{}, err
				}
				return fileSeg{seg: seg.Number, file: frozen, sha224: seg.SHA224, size: seg.Size}, nil
			}
			return fileSeg{seg: seg.Number, file: partial, sha224: seg.SHA224, size: seg.Size}, nil
		}
	}

	// Otherwise, download.
	req, err := http.NewRequest("GET", segURL.String(), nil)
	if err != nil {
		return fileSeg{}, err
	}
	req = req.WithContext(ctx)
	req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", len(have), seg.Size-1))

	log.Printf("Downloading %d bytes of %s ...", seg.Size-int64(len(have)), segURL)
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return fileSeg{}, err
	}
	if res.StatusCode != 200 && res.StatusCode != 206 {
		return fileSeg{}, fmt.Errorf("%s: %s", segURL.String(), res.Status)
	}
	slurp, err := ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		return fileSeg{}, err
	}

	var newContents []byte
	if int64(len(slurp)) == seg.Size {
		newContents = slurp
	} else if int64(len(have)+len(slurp)) == seg.Size {
		newContents = append(have, slurp...)
	}
	got224 := fmt.Sprintf("%x", sha256.Sum224(newContents))
	if got224 != seg.SHA224 {
		if len(have) == 0 {
			return fileSeg{}, errors.New("corrupt download")
		}
		// Try again
		os.Remove(partial)
		return ns.syncSeg(ctx, seg)
	}
	tf, err := ioutil.TempFile(ns.cacheDir, "tempseg")
	if err != nil {
		return fileSeg{}, err
	}
	if _, err := tf.Write(newContents); err != nil {
		return fileSeg{}, err
	}
	if err := tf.Close(); err != nil {
		return fileSeg{}, err
	}
	finalName := partial
	if !isFinalSeg {
		finalName = frozen
	}
	if err := os.Rename(tf.Name(), finalName); err != nil {
		return fileSeg{}, err
	}
	log.Printf("wrote %v", finalName)
	return fileSeg{seg: seg.Number, file: finalName, size: seg.Size, sha224: seg.SHA224}, nil
}

type LogSegmentJSON struct {
	Number int    `json:"number"`
	Size   int64  `json:"size"`
	SHA224 string `json:"sha224"`
	URL    string `json:"url"`
}