From c4eb32a8fa47b9b10a510120968b7011bed4d7da Mon Sep 17 00:00:00 2001
From: Raphael Vigee <raphvigee@gmail.com>
Date: Sat, 20 Apr 2024 11:55:36 +0100
Subject: [PATCH] better handle force ctrlc

---
 utils/xcontext/context.go | 10 +++++++---
 utils/xtea/program.go     |  2 ++
 utils/xtea/singletui.go   | 21 +++++++++++++++++++++
 3 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/utils/xcontext/context.go b/utils/xcontext/context.go
index 06f6bc04..ffde2b94 100644
--- a/utils/xcontext/context.go
+++ b/utils/xcontext/context.go
@@ -5,6 +5,7 @@ import (
 	"github.com/hephbuild/heph/log/log"
 	"github.com/hephbuild/heph/utils/ads"
 	"github.com/hephbuild/heph/utils/xsync"
+	"github.com/hephbuild/heph/utils/xtea"
 	"os"
 	"os/signal"
 	"sync"
@@ -138,6 +139,8 @@ func Cancel(ctx context.Context) {
 	cancel()
 }
 
+const stuckTimeout = 5 * time.Second
+
 func BootstrapSoftCancel() (context.Context, context.CancelFunc) {
 	ctx, cancel := context.WithCancel(context.Background())
 
@@ -168,12 +171,12 @@ func BootstrapSoftCancel() (context.Context, context.CancelFunc) {
 			// Wait for soft cancel to all be unregistered, should be fast, unless something is stuck
 			case <-sc.wait():
 				// Wait for graceful exit
-				<-time.After(2 * time.Second)
-			case <-time.After(2 * time.Second):
+				<-time.After(stuckTimeout)
+			case <-time.After(stuckTimeout):
 				// All soft cancel did not unregister, something is stuck...
 			}
 		} else {
-			<-time.After(2 * time.Second)
+			<-time.After(stuckTimeout)
 		}
 
 		log.Error("Something seems to be stuck, ctrl+c one more time to forcefully exit")
@@ -182,6 +185,7 @@ func BootstrapSoftCancel() (context.Context, context.CancelFunc) {
 		if sig, ok := sig.(syscall.Signal); ok {
 			sigN = int(sig)
 		}
+		xtea.ResetTerminal()
 		os.Exit(128 + sigN)
 	}()
 
diff --git a/utils/xtea/program.go b/utils/xtea/program.go
index 2adf1307..114c3982 100644
--- a/utils/xtea/program.go
+++ b/utils/xtea/program.go
@@ -13,9 +13,11 @@ func RunModel(model tea.Model, opts ...tea.ProgramOption) error {
 
 func Run(p *tea.Program) error {
 	defer func() {
+		SetResetTerminal(nil)
 		log.SetDiversion(nil)
 	}()
 
+	SetResetTerminal(p)
 	_, err := p.Run()
 	return err
 }
diff --git a/utils/xtea/singletui.go b/utils/xtea/singletui.go
index a3ad2faa..d81d1858 100644
--- a/utils/xtea/singletui.go
+++ b/utils/xtea/singletui.go
@@ -2,6 +2,7 @@ package xtea
 
 import (
 	"fmt"
+	tea "github.com/charmbracelet/bubbletea"
 	"runtime/debug"
 	"sync"
 )
@@ -33,3 +34,23 @@ func SingleflightDone() {
 	}
 	tuim.Unlock()
 }
+
+var resetTermFunc func()
+
+func SetResetTerminal(p *tea.Program) {
+	if p == nil {
+		resetTermFunc = nil
+		return
+	}
+
+	resetTermFunc = func() {
+		_ = p.ReleaseTerminal()
+		_ = p.RestoreTerminal()
+	}
+}
+
+func ResetTerminal() {
+	if f := resetTermFunc; f != nil {
+		f()
+	}
+}