package pipeserver

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"runtime/trace"
	"sync"
	"time"

	"code.justin.tv/rhys/nursery/cmd/multicp/conn"
	"code.justin.tv/rhys/nursery/cmd/multicp/join"
	"code.justin.tv/rhys/nursery/cmd/multicp/netpipe"
	"github.com/golang/protobuf/ptypes"
	"github.com/golang/protobuf/ptypes/empty"
)

const (
	maxKeepAlive    = 30 * time.Second
	maxMessageBytes = 1000
)

type pipe struct {
	closer   io.Closer
	pb       *netpipe.Pipe
	timeout  time.Duration
	timer    *time.Timer
	join     *join.Buffer
	toClient offsetBuffer

	mu       sync.Mutex
	readOnly bool // The pipe is closed, but still has data in the read buffer
}

type Server struct {
	NewConn func(ctx context.Context, target string, conn net.Conn) error
	dial    func(ctx context.Context, target string) (io.ReadWriteCloser, error)

	LocalAddr  func(ctx context.Context) net.Addr
	RemoteAddr func(ctx context.Context) net.Addr

	mu    sync.Mutex
	pipes map[string]*pipe
}

var _ netpipe.NetPipe = (*Server)(nil)

func (s *Server) expire(p *pipe) {}

func (s *Server) prepare(ctx context.Context) {
	defer trace.StartRegion(ctx, "prepare").End()
	trace.WithRegion(ctx, "Lock", s.mu.Lock)
	defer s.mu.Unlock()

	if s.pipes != nil {
		return
	}

	s.pipes = make(map[string]*pipe)
}

// CreatePipe opens a new bidirectional pipe.
func (s *Server) CreatePipe(ctx context.Context, req *netpipe.CreatePipeRequest) (*netpipe.Pipe, error) {
	ctx, task := trace.NewTask(ctx, "netpipe.CreatePipe")
	defer task.End()
	s.prepare(ctx)

	target := req.GetTarget()
	trace.Logf(ctx, "target", "%q", target)
	if target == "" {
		return nil, errors.New("did not specify target")
	}

	var timeout time.Duration
	if bid := req.GetKeepaliveDuration(); bid != nil {
		var err error
		timeout, err = ptypes.Duration(bid)
		if err != nil {
			return nil, errors.New("invalid keepalive")
		}
	}
	if timeout <= 0 || timeout > maxKeepAlive {
		timeout = maxKeepAlive
	}

	nameSeed := make([]byte, 32)
	_, err := io.ReadFull(rand.Reader, nameSeed)
	if err != nil {
		return nil, errors.New("could not name pipe")
	}
	name := base64.RawStdEncoding.EncodeToString(nameSeed)
	trace.Logf(ctx, "name", "%q", name)

	resp := &netpipe.Pipe{
		Name:              name,
		MaxMessageBytes:   maxMessageBytes,
		KeepaliveDuration: ptypes.DurationProto(timeout),
	}

	p := &pipe{pb: resp, timeout: timeout, join: join.NewBuffer()}

	// To read data sent by client
	reader := new(conn.TimeoutReader)
	go func() { reader.ReadFrom(p.join) }()
	// To write data to the client
	writer := new(conn.TimeoutWriter)
	go func() { writer.WriteTo(&p.toClient) }()

	local, remote := pipeAddr("server"), pipeAddr(name)
	if fn := s.LocalAddr; fn != nil {
		local = pipeAddr(fmt.Sprintf("%s", fn(ctx)))
	}
	if fn := s.RemoteAddr; fn != nil {
		remote = pipeAddr(fmt.Sprintf("%s:%s", fn(ctx), name))
	}

	c := &conn.ReadWriter{
		Reader: reader,
		Writer: writer,

		Closers: []io.Closer{&p.toClient},

		Local:  local,
		Remote: remote,
	}
	// After DeletePipe, Read calls on this side of the net.Conn will drain the
	// buffer and then return io.EOF
	p.closer = p.join

	err = s.NewConn(ctx, target, c)
	if err != nil {
		c.Close()
		return nil, errors.New("could not open pipe")
	}

	s.mu.Lock()
	defer s.mu.Unlock()
	s.pipes[name] = p

	return resp, nil
}

func (s *Server) DeletePipe(ctx context.Context, req *netpipe.DeletePipeRequest) (*empty.Empty, error) {
	ctx, task := trace.NewTask(ctx, "netpipe.DeletePipe")
	defer task.End()
	s.prepare(ctx)

	name := req.GetName()
	trace.Logf(ctx, "name", "%q", name)
	if name == "" {
		return nil, errors.New("did not specify name")
	}

	s.mu.Lock()
	pipe := s.pipes[name]
	delete(s.pipes, name) // TODO: Handle this in WriteAt when committing new read offset
	s.mu.Unlock()

	if pipe == nil {
		return nil, os.ErrNotExist
	}
	pipe.mu.Lock()
	pipe.readOnly = true
	pipe.mu.Unlock()

	pipe.closer.Close()

	return nil, io.EOF
}

func (s *Server) Read(ctx context.Context, req *netpipe.ReadRequest) (*netpipe.ReadResponse, error) {
	ctx, task := trace.NewTask(ctx, "netpipe.Read")
	defer task.End()
	s.prepare(ctx)

	name := req.GetName()
	trace.Logf(ctx, "name", "%q", name)
	if name == "" {
		return nil, errors.New("did not specify name")
	}

	s.mu.Lock()
	pipe := s.pipes[name]
	s.mu.Unlock()
	if pipe == nil {
		return nil, errors.New("pipe not found")
	}

	timeout, err := ptypes.Duration(req.GetTimeoutDuration())
	if err != nil || timeout < 0 {
		return nil, errors.New("invalid timeout")
	}

	buf := make([]byte, 1<<10)
	if max := int(req.GetMaxReadBytes()); max < len(buf) {
		buf = buf[:max]
	}

	resp := &netpipe.ReadResponse{}

	var retry, done <-chan time.Time
	if timeout > 0 {
		const tick = 5 * time.Millisecond
		ticker := time.NewTicker(tick)
		defer ticker.Stop()
		retry = ticker.C

		timeoutDone := time.NewTimer(timeout)
		defer timeoutDone.Stop()
		done = timeoutDone.C
	}

	var readOnly bool
	for i := 0; ; i++ {
		pipe.toClient.Contents(func(b []byte, off int64, err error) {
			// Pull data from the start of the buffer.
			// TODO: schedule these to allow more than one buffer per RTT

			if err != nil {
				readOnly = true
			}

			n := copy(buf, b)
			resp.ReadData = buf[:n]
			resp.ReadOffset = off

			trace.Logf(ctx, "Read",
				"commitOffset=%d readOffset=%d available=%d returned=%d err=%q",
				off, resp.ReadOffset, len(b), len(resp.ReadData), err)
		})
		if len(resp.ReadData) == 0 && retry != nil {
			// TODO: jitter the actual sleep time
			//
			// TODO: subscribe to new data, rather than polling
			//
			// TODO: invert the control here so the buffer is able to pick which
			// connection to use for particular data, especially since there may
			// be several requests waiting for the same data.
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case <-retry:
				continue
			case <-done:
			}
		}
		break
	}

	pipe.mu.Lock()
	if readOnly || pipe.readOnly {
		readOnly, pipe.readOnly = true, true
	}
	pipe.mu.Unlock()

	if len(resp.GetReadData()) == 0 && readOnly {
		// This is the end of the data stream; send EOF. We'll delete when the data is acknowledged.
		resp.EndOfFile = true
	}

	return resp, nil
}

func (s *Server) WriteAt(ctx context.Context, req *netpipe.WriteAtRequest) (*netpipe.WriteAtResponse, error) {
	ctx, task := trace.NewTask(ctx, "netpipe.WriteAt")
	defer task.End()
	s.prepare(ctx)

	name := req.GetName()
	trace.Logf(ctx, "name", "%q", name)
	if name == "" {
		return nil, errors.New("did not specify name")
	}

	readOffset := req.GetCommittedReadOffset()
	if readOffset < 0 {
		return nil, errors.New("invalid committed read offset")
	}
	// TODO: check if pipe is read-only, and if so check whether its read buffer is now empty

	writeOffset := req.GetWriteOffset()
	if writeOffset < 0 {
		return nil, errors.New("invalid write offset")
	}

	s.mu.Lock()
	pipe := s.pipes[name]
	s.mu.Unlock()

	if pipe == nil {
		return nil, errors.New("pipe not found")
	}

	trace.Logf(ctx, "WriteAt",
		"length=%d offset=%d readOffset=%d",
		len(req.GetWriteData()), writeOffset, readOffset)

	// pipe.timer.Reset(pipe.timeout)

	pipe.toClient.Commit(readOffset)

	pipe.mu.Lock()
	readOnly := pipe.readOnly
	pipe.mu.Unlock()
	if readOnly {
		var ackEOF bool
		pipe.toClient.Contents(func(b []byte, off int64, err error) {
			if len(b) == 0 {
				// Client has ACKed all data through to EOF.
				ackEOF = true
			}
		})
		if ackEOF {
			// TODO: delete the pipe.
		}
	}

	// TODO: check that the pipe is still active

	writeData := req.GetWriteData()
	if l := len(writeData); l > int(pipe.pb.GetMaxMessageBytes()) {
		return nil, errors.New("write buffer is too long")
	}

	// TODO: check end offset against joiner's window
	_, err := pipe.join.WriteAt(writeData, writeOffset)
	if err != nil {
		return nil, errors.New("could not write data")
	}

	resp := &netpipe.WriteAtResponse{
		CommittedWriteOffset: pipe.join.ContiguousBytes(),
	}

	return resp, nil
}

type offsetBuffer struct {
	mu       sync.Mutex
	pending  bytes.Buffer
	consumed int64
	err      error
}

func (ob *offsetBuffer) Close() error {
	ob.mu.Lock()
	defer ob.mu.Unlock()

	ob.err = os.ErrClosed
	return nil
}

func (ob *offsetBuffer) Write(b []byte) (int, error) {
	ob.mu.Lock()
	defer ob.mu.Unlock()

	n, err := ob.pending.Write(b)
	return n, err
}

func (ob *offsetBuffer) Commit(offset int64) {
	ob.mu.Lock()
	defer ob.mu.Unlock()

	var (
		start   = ob.consumed
		pending = int64(ob.pending.Len())
		delta   = offset - start
	)

	if offset < 0 {
		// It's meaningless to commit negative byte offsets. Ignore.
		return
	}

	if delta < 0 {
		// Old news, nothing to do
		return
	}
	if delta > pending {
		// Trying to commit future data .. that's not right, but let's ignore
		return
	}

	_ = ob.pending.Next(int(delta))
	ob.consumed = offset
}

func (ob *offsetBuffer) Contents(fn func(b []byte, off int64, err error)) {
	ob.mu.Lock()
	defer ob.mu.Unlock()

	b := ob.pending.Bytes()
	off := ob.consumed
	err := ob.err

	fn(b, off, err)
}

type pipeAddr string

func (pa pipeAddr) Network() string { return "netpipe" }
func (pa pipeAddr) String() string  { return string(pa) }
