package main

import (
	"fmt"
	"io"
	"log"
	"math/rand"
	"net"
	"net/http"
	"sync"
	"time"

	"code.justin.tv/common/chitin"
	"code.justin.tv/common/chitin/cmd/grpc_demo/arithmetic"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

type adder struct {
	mu sync.Mutex
	cl arithmetic.AdderClient
}

func (a *adder) Add(ctx context.Context, req *arithmetic.AddRequest) (*arithmetic.AddReply, error) {
	a.mu.Lock()
	cl := a.cl
	a.mu.Unlock()
	if cl == nil {
		return nil, fmt.Errorf("no client available")
	}

	var reply arithmetic.AddReply

	if len(req.Number) == 0 {
		return &reply, nil
	}
	if len(req.Number) == 1 {
		reply.Number = req.Number[0]
		return &reply, nil
	}

	const jobs = 2

	stride := len(req.Number) / jobs
	var wg sync.WaitGroup
	results := make([]*arithmetic.AddReply, jobs+1)
	errs := make([]error, jobs+1)

	for i, start := 0, 0; start < len(req.Number); i, start = i+1, start+stride {
		wg.Add(1)
		go func(i, start int) {
			defer wg.Done()

			time.Sleep(time.Duration(rand.Int63n(100)) * time.Millisecond)

			end := start + stride
			if end > len(req.Number) {
				end = len(req.Number)
			}

			var subReq arithmetic.AddRequest
			subReq.Number = req.Number[start:end]
			subRep, err := cl.Add(ctx, &subReq)

			results[i], errs[i] = subRep, err
		}(i, start)
	}

	wg.Wait()

	for _, err := range errs {
		if err != nil {
			return nil, err
		}
	}

	for _, result := range results {
		if result == nil {
			continue
		}
		reply.Number += result.Number
	}

	time.Sleep(time.Duration(rand.Int63n(200)) * time.Millisecond)

	return &reply, nil
}

func (a *adder) RunningTotal(stream arithmetic.Adder_RunningTotalServer) error {
	a.mu.Lock()
	cl := a.cl
	a.mu.Unlock()
	if cl == nil {
		return fmt.Errorf("no client available")
	}

	ctx := stream.Context()

	var tally int64
	for {
		req, err := stream.Recv()
		if err != nil {
			log.Printf("server recv err=%q", err)
			if err == io.EOF {
				return nil
			}
			return err
		}

		var subReq arithmetic.AddRequest
		subReq.Number = make([]int64, len(req.Number)+1)
		subReq.Number[0] = tally
		copy(subReq.Number[1:], req.Number)

		rep, err := cl.Add(ctx, &subReq)
		if err != nil {
			log.Printf("server add err=%q", err)
			return err
		}
		tally = rep.Number

		err = stream.Send(rep)
		if err != nil {
			log.Printf("server send err=%q", err)
			return err
		}
	}
}

func (a *adder) setClient(cl arithmetic.AdderClient) {
	a.cl = cl
}

func main() {
	log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)
	rand.Seed(time.Now().UnixNano())

	s := grpc.NewServer()
	var srv adder

	arithmetic.RegisterAdderServer(s, &srv)

	go log.Printf("debug http err=%q", http.ListenAndServe("localhost:8080", nil))

	l, err := net.Listen("tcp", "localhost:0")
	if err != nil {
		log.Fatalf("listen err=%q", err)
	}

	err = chitin.ExperimentalTraceProcessOptIn()
	if err != nil {
		log.Fatalf("trace enable err=%q", err)
	}

	ctx := context.Background()
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	var wg sync.WaitGroup

	wg.Add(1)
	go func() {
		defer wg.Done()
		<-ctx.Done()
		err := l.Close()
		if err != nil {
			log.Printf("close err=%q", err)
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()

		err = s.Serve(l)
		if err != nil && ctx.Err() == nil {
			log.Fatalf("serve err=%q", err)
		}
	}()

	cc, err := grpc.Dial(l.Addr().String(), grpc.WithInsecure())
	if err != nil {
		log.Fatalf("dial err=%q", err)
	}

	client := arithmetic.NewAdderClient(cc)
	srv.setClient(client)

	// unary add request

	req := &arithmetic.AddRequest{
		Number: []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11},
	}
	reply, err := client.Add(ctx, req)
	if err != nil {
		log.Fatalf("reply err=%q", err)
	}

	log.Printf("sum=%d", reply.Number)

	// streaming add request

	var streamWg sync.WaitGroup
	stream, err := client.RunningTotal(ctx)
	if err != nil {
		log.Printf("client stream err=%q", err)
		return
	}
	streamWg.Add(1)
	go func() {
		defer streamWg.Done()
		for {
			rep, err := stream.Recv()
			if err != nil {
				log.Printf("client recv err=%q", err)
				break
			}
			log.Printf("client tally=%d", rep.Number)
		}
	}()
	for _, i := range []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11} {
		req := &arithmetic.AddRequest{
			Number: []int64{i},
		}
		err := stream.Send(req)
		if err != nil {
			log.Printf("client send err=%q", err)
			break
		}
	}
	err = stream.CloseSend()
	if err != nil {
		log.Printf("client close err=%q", err)
	}
	streamWg.Wait()

	// shut down server
	cancel()
	wg.Wait()
}
