package concurrency

import (
    "errors"
    "fmt"

    v3 "github.com/coreos/etcd/clientv3"
    "github.com/coreos/etcd/mvcc/mvccpb"
    "golang.org/x/net/context"
)

var (
    ErrElectionNotLeader = errors.New("election: not leader")
    ErrElectionNoLeader  = errors.New("election: no leader")
)

type Election struct {
    session *Session

    keyPrefix string

    leaderKey     string
    leaderRev     int64
    leaderSession *Session
}

// NewElection returns a new election on a given key prefix.
func NewElection(s *Session, pfx string) *Election {
    return &Election{session: s, keyPrefix: pfx + "/"}
}

// Campaign puts a value as eligible for the election. It blocks until
// it is elected, an error occurs, or the context is cancelled.
func (e *Election) Campaign(ctx context.Context, val string) error {
    s := e.session
    client := e.session.Client()

    k := fmt.Sprintf("%s%x", e.keyPrefix, s.Lease())
    txn := client.Txn(ctx).If(v3.Compare(v3.CreateRevision(k), "=", 0))
    txn = txn.Then(v3.OpPut(k, val, v3.WithLease(s.Lease())))
    txn = txn.Else(v3.OpGet(k))
    resp, err := txn.Commit()
    if err != nil {
        return err
    }
    e.leaderKey, e.leaderRev, e.leaderSession = k, resp.Header.Revision, s
    if !resp.Succeeded {
        kv := resp.Responses[0].GetResponseRange().Kvs[0]
        e.leaderRev = kv.CreateRevision
        if string(kv.Value) != val {
            if err = e.Proclaim(ctx, val); err != nil {
                e.Resign(ctx)
                return err
            }
        }
    }

    err = waitDeletes(ctx, client, e.keyPrefix, e.leaderRev-1)
    if err != nil {
        // clean up in case of context cancel
        select {
        case <-ctx.Done():
            e.Resign(client.Ctx())
        default:
            e.leaderSession = nil
        }
        return err
    }

    return nil
}

// Proclaim lets the leader announce a new value without another election.
func (e *Election) Proclaim(ctx context.Context, val string) error {
    if e.leaderSession == nil {
        return ErrElectionNotLeader
    }
    client := e.session.Client()
    cmp := v3.Compare(v3.CreateRevision(e.leaderKey), "=", e.leaderRev)
    txn := client.Txn(ctx).If(cmp)
    txn = txn.Then(v3.OpPut(e.leaderKey, val, v3.WithLease(e.leaderSession.Lease())))
    tresp, terr := txn.Commit()
    if terr != nil {
        return terr
    }
    if !tresp.Succeeded {
        e.leaderKey = ""
        return ErrElectionNotLeader
    }
    return nil
}

// Resign lets a leader start a new election.
func (e *Election) Resign(ctx context.Context) (err error) {
    if e.leaderSession == nil {
        return nil
    }
    client := e.session.Client()
    _, err = client.Delete(ctx, e.leaderKey)
    e.leaderKey = ""
    e.leaderSession = nil
    return err
}

// Leader returns the leader value for the current election.
func (e *Election) Leader(ctx context.Context) (string, error) {
    client := e.session.Client()
    resp, err := client.Get(ctx, e.keyPrefix, v3.WithFirstCreate()...)
    if err != nil {
        return "", err
    } else if len(resp.Kvs) == 0 {
        // no leader currently elected
        return "", ErrElectionNoLeader
    }
    return string(resp.Kvs[0].Value), nil
}

// Observe returns a channel that observes all leader proposal values as
// GetResponse values on the current leader key. The channel closes when
// the context is cancelled or the underlying watcher is otherwise disrupted.
func (e *Election) Observe(ctx context.Context) <-chan v3.GetResponse {
    retc := make(chan v3.GetResponse)
    go e.observe(ctx, retc)
    return retc
}

func (e *Election) observe(ctx context.Context, ch chan<- v3.GetResponse) {
    client := e.session.Client()

    defer close(ch)
    for {
        resp, err := client.Get(ctx, e.keyPrefix, v3.WithFirstCreate()...)
        if err != nil {
            return
        }

        var kv *mvccpb.KeyValue

        cctx, cancel := context.WithCancel(ctx)
        if len(resp.Kvs) == 0 {
            // wait for first key put on prefix
            opts := []v3.OpOption{v3.WithRev(resp.Header.Revision), v3.WithPrefix()}
            wch := client.Watch(cctx, e.keyPrefix, opts...)

            for kv == nil {
                wr, ok := <-wch
                if !ok || wr.Err() != nil {
                    cancel()
                    return
                }
                // only accept PUTs; a DELETE will make observe() spin
                for _, ev := range wr.Events {
                    if ev.Type == mvccpb.PUT {
                        kv = ev.Kv
                        break
                    }
                }
            }
        } else {
            kv = resp.Kvs[0]
        }

        wch := client.Watch(cctx, string(kv.Key), v3.WithRev(kv.ModRevision))
        keyDeleted := false
        for !keyDeleted {
            wr, ok := <-wch
            if !ok {
                return
            }
            for _, ev := range wr.Events {
                if ev.Type == mvccpb.DELETE {
                    keyDeleted = true
                    break
                }
                resp.Header = &wr.Header
                resp.Kvs = []*mvccpb.KeyValue{ev.Kv}
                select {
                case ch <- *resp:
                case <-cctx.Done():
                    return
                }
            }
        }
        cancel()
    }
}

// Key returns the leader key if elected, empty string otherwise.
func (e *Election) Key() string { return e.leaderKey }