added smart remote/branch selection

5 files changed, 30 insertions, 4 deletions

diff --git a/pkg/git/branch.go b/pkg/git/branch.go
index e691b2a..a700c30 100644
--- a/pkg/git/branch.go
+++ b/pkg/git/branch.go
@@ -14,7 +14,7 @@ type Branch struct {
 	Clean     bool
 }
 
-func (entity *RepoEntity) GetActiveBranch() (branch *Branch) {
+func (entity *RepoEntity) getActiveBranch() (branch *Branch) {
 	headRef, _ := entity.Repository.Head()
 	for _, lb := range entity.Branches {
 		if lb.Name == headRef.Name().Short() {
@@ -75,6 +75,12 @@ func (entity *RepoEntity) Checkout(branch *Branch) error {
 	entity.Commit = entity.Commits[0]
 	entity.Branch = branch
 	entity.Branch.Pushables, entity.Branch.Pullables = UpstreamDifferenceCount(entity.AbsPath)
+	// TODO: same code on 3 different occasion, maybe something wrong?
+	// make this conditional on global scale
+	if err = entity.Remote.switchRemoteBranch(entity.Remote.Name + "/" + entity.Branch.Name); err !=nil {
+		// probably couldn't find, but its ok.
+		return nil
+	}
 	return nil
 }
 
diff --git a/pkg/git/remote.go b/pkg/git/remote.go
index ba05e1a..6beef29 100644
--- a/pkg/git/remote.go
+++ b/pkg/git/remote.go
@@ -20,7 +20,10 @@ func (entity *RepoEntity) NextRemote() error {
 	} else {
 		entity.Remote = entity.Remotes[currentRemoteIndex+1]
 	}
-
+	// TODO: same code on 3 different occasion, maybe something wrong?
+	if err := entity.Remote.switchRemoteBranch(entity.Remote.Name + "/" + entity.Branch.Name); err !=nil {
+		// probably couldn't find, but its ok.
+	}
 	return nil
 }
 
diff --git a/pkg/git/remotebranch.go b/pkg/git/remotebranch.go
index 435a259..583c323 100644
--- a/pkg/git/remotebranch.go
+++ b/pkg/git/remotebranch.go
@@ -1,6 +1,7 @@
 package git
 
 import (
+	"errors"
 	"strings"
 
 	"gopkg.in/src-d/go-git.v4"
@@ -26,7 +27,6 @@ func (remote *Remote) NextRemoteBranch() error {
 	} else {
 		remote.Branch = remote.Branches[currentRemoteIndex+1]
 	}
-
 	return nil
 }
 
@@ -65,3 +65,13 @@ func remoteBranchesIter(s storer.ReferenceStorer) (storer.ReferenceIter, error)
 		return false
 	}, refs), nil
 }
+
+func (remote *Remote) switchRemoteBranch(remoteBranchName string) error {
+	for _, rb := range remote.Branches {
+		if rb.Name == remoteBranchName {
+			remote.Branch = rb
+			return nil
+		}
+	}
+	return errors.New("Remote branch not found.")
+}
diff --git a/pkg/git/repository.go b/pkg/git/repository.go
index d3ea1ec..8c6366c 100644
--- a/pkg/git/repository.go
+++ b/pkg/git/repository.go
@@ -60,10 +60,14 @@ func InitializeRepository(directory string) (entity *RepoEntity, err error) {
 		return entity, errors.New("There is no commit for this repository: " + directory)
 	}
 	entity.loadRemotes()
-	entity.Branch = entity.GetActiveBranch()
+	entity.Branch = entity.getActiveBranch()
 	if len(entity.Remotes) > 0 {
 		// TODO: tend to take origin/master as default
 		entity.Remote = entity.Remotes[0]
+		// TODO: same code on 3 different occasion, maybe something wrong?
+		if err = entity.Remote.switchRemoteBranch(entity.Remote.Name + "/" + entity.Branch.Name); err !=nil {
+		// probably couldn't find, but its ok.
+		}
 	} else {
 		return entity, errors.New("There is no remote for this repository: " + directory)
 	}
diff --git a/pkg/gui/branchview.go b/pkg/gui/branchview.go
index 0358e82..087f92a 100644
--- a/pkg/gui/branchview.go
+++ b/pkg/gui/branchview.go
@@ -56,6 +56,9 @@ func (gui *Gui) nextBranch(g *gocui.Gui, v *gocui.View) error {
 	if err = gui.updateCommits(g, entity); err != nil {
 		return err
 	}
+	if err = gui.updateRemoteBranches(g, entity); err != nil {
+		return err
+	}
 	if err = gui.refreshMain(g); err != nil {
 		return err
 	}