additional tools

14 files changed, 1361 insertions, 59 deletions

diff --git a/internal/tui/components/core/button.go b/internal/tui/components/core/button.go
new file mode 100644
index 000000000..102957e0e
--- /dev/null
+++ b/internal/tui/components/core/button.go
@@ -0,0 +1,287 @@
+package core
+
+import (
+	"github.com/charmbracelet/bubbles/key"
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/lipgloss"
+	"github.com/kujtimiihoxha/termai/internal/tui/styles"
+)
+
+// ButtonKeyMap defines key bindings for the button component
+type ButtonKeyMap struct {
+	Enter key.Binding
+}
+
+// DefaultButtonKeyMap returns default key bindings for the button
+func DefaultButtonKeyMap() ButtonKeyMap {
+	return ButtonKeyMap{
+		Enter: key.NewBinding(
+			key.WithKeys("enter"),
+			key.WithHelp("enter", "select"),
+		),
+	}
+}
+
+// ShortHelp returns keybinding help
+func (k ButtonKeyMap) ShortHelp() []key.Binding {
+	return []key.Binding{k.Enter}
+}
+
+// FullHelp returns full help info for keybindings
+func (k ButtonKeyMap) FullHelp() [][]key.Binding {
+	return [][]key.Binding{
+		{k.Enter},
+	}
+}
+
+// ButtonState represents the state of a button
+type ButtonState int
+
+const (
+	// ButtonNormal is the default state
+	ButtonNormal ButtonState = iota
+	// ButtonHovered is when the button is focused/hovered
+	ButtonHovered
+	// ButtonPressed is when the button is being pressed
+	ButtonPressed
+	// ButtonDisabled is when the button is disabled
+	ButtonDisabled
+)
+
+// ButtonVariant defines the visual style variant of a button
+type ButtonVariant int
+
+const (
+	// ButtonPrimary uses primary color styling
+	ButtonPrimary ButtonVariant = iota
+	// ButtonSecondary uses secondary color styling
+	ButtonSecondary
+	// ButtonDanger uses danger/error color styling
+	ButtonDanger
+	// ButtonWarning uses warning color styling
+	ButtonWarning
+	// ButtonNeutral uses neutral color styling
+	ButtonNeutral
+)
+
+// ButtonMsg is sent when a button is clicked
+type ButtonMsg struct {
+	ID      string
+	Payload interface{}
+}
+
+// ButtonCmp represents a clickable button component
+type ButtonCmp struct {
+	id         string
+	label      string
+	width      int
+	height     int
+	state      ButtonState
+	variant    ButtonVariant
+	keyMap     ButtonKeyMap
+	payload    interface{}
+	style      lipgloss.Style
+	hoverStyle lipgloss.Style
+}
+
+// NewButtonCmp creates a new button component
+func NewButtonCmp(id, label string) *ButtonCmp {
+	b := &ButtonCmp{
+		id:      id,
+		label:   label,
+		state:   ButtonNormal,
+		variant: ButtonPrimary,
+		keyMap:  DefaultButtonKeyMap(),
+		width:   len(label) + 4, // add some padding
+		height:  1,
+	}
+	b.updateStyles()
+	return b
+}
+
+// WithVariant sets the button variant
+func (b *ButtonCmp) WithVariant(variant ButtonVariant) *ButtonCmp {
+	b.variant = variant
+	b.updateStyles()
+	return b
+}
+
+// WithPayload sets the payload sent with button events
+func (b *ButtonCmp) WithPayload(payload interface{}) *ButtonCmp {
+	b.payload = payload
+	return b
+}
+
+// WithWidth sets a custom width
+func (b *ButtonCmp) WithWidth(width int) *ButtonCmp {
+	b.width = width
+	b.updateStyles()
+	return b
+}
+
+// updateStyles recalculates styles based on current state and variant
+func (b *ButtonCmp) updateStyles() {
+	// Base styles
+	b.style = styles.Regular.
+		Padding(0, 1).
+		Width(b.width).
+		Align(lipgloss.Center).
+		BorderStyle(lipgloss.RoundedBorder())
+
+	b.hoverStyle = b.style.
+		Bold(true)
+
+	// Variant-specific styling
+	switch b.variant {
+	case ButtonPrimary:
+		b.style = b.style.
+			Foreground(styles.Base).
+			Background(styles.Primary).
+			BorderForeground(styles.Primary)
+
+		b.hoverStyle = b.hoverStyle.
+			Foreground(styles.Base).
+			Background(styles.Blue).
+			BorderForeground(styles.Blue)
+
+	case ButtonSecondary:
+		b.style = b.style.
+			Foreground(styles.Base).
+			Background(styles.Secondary).
+			BorderForeground(styles.Secondary)
+
+		b.hoverStyle = b.hoverStyle.
+			Foreground(styles.Base).
+			Background(styles.Mauve).
+			BorderForeground(styles.Mauve)
+
+	case ButtonDanger:
+		b.style = b.style.
+			Foreground(styles.Base).
+			Background(styles.Error).
+			BorderForeground(styles.Error)
+
+		b.hoverStyle = b.hoverStyle.
+			Foreground(styles.Base).
+			Background(styles.Red).
+			BorderForeground(styles.Red)
+
+	case ButtonWarning:
+		b.style = b.style.
+			Foreground(styles.Text).
+			Background(styles.Warning).
+			BorderForeground(styles.Warning)
+
+		b.hoverStyle = b.hoverStyle.
+			Foreground(styles.Text).
+			Background(styles.Peach).
+			BorderForeground(styles.Peach)
+
+	case ButtonNeutral:
+		b.style = b.style.
+			Foreground(styles.Text).
+			Background(styles.Grey).
+			BorderForeground(styles.Grey)
+
+		b.hoverStyle = b.hoverStyle.
+			Foreground(styles.Text).
+			Background(styles.DarkGrey).
+			BorderForeground(styles.DarkGrey)
+	}
+
+	// Disabled style override
+	if b.state == ButtonDisabled {
+		b.style = b.style.
+			Foreground(styles.SubText0).
+			Background(styles.LightGrey).
+			BorderForeground(styles.LightGrey)
+	}
+}
+
+// SetSize sets the button size
+func (b *ButtonCmp) SetSize(width, height int) {
+	b.width = width
+	b.height = height
+	b.updateStyles()
+}
+
+// Focus sets the button to focused state
+func (b *ButtonCmp) Focus() tea.Cmd {
+	if b.state != ButtonDisabled {
+		b.state = ButtonHovered
+	}
+	return nil
+}
+
+// Blur sets the button to normal state
+func (b *ButtonCmp) Blur() tea.Cmd {
+	if b.state != ButtonDisabled {
+		b.state = ButtonNormal
+	}
+	return nil
+}
+
+// Disable sets the button to disabled state
+func (b *ButtonCmp) Disable() {
+	b.state = ButtonDisabled
+	b.updateStyles()
+}
+
+// Enable enables the button if disabled
+func (b *ButtonCmp) Enable() {
+	if b.state == ButtonDisabled {
+		b.state = ButtonNormal
+		b.updateStyles()
+	}
+}
+
+// IsDisabled returns whether the button is disabled
+func (b *ButtonCmp) IsDisabled() bool {
+	return b.state == ButtonDisabled
+}
+
+// IsFocused returns whether the button is focused
+func (b *ButtonCmp) IsFocused() bool {
+	return b.state == ButtonHovered
+}
+
+// Init initializes the button
+func (b *ButtonCmp) Init() tea.Cmd {
+	return nil
+}
+
+// Update handles messages and user input
+func (b *ButtonCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	// Skip updates if disabled
+	if b.state == ButtonDisabled {
+		return b, nil
+	}
+
+	switch msg := msg.(type) {
+	case tea.KeyMsg:
+		// Handle key presses when focused
+		if b.state == ButtonHovered {
+			switch {
+			case key.Matches(msg, b.keyMap.Enter):
+				b.state = ButtonPressed
+				return b, func() tea.Msg {
+					return ButtonMsg{
+						ID:      b.id,
+						Payload: b.payload,
+					}
+				}
+			}
+		}
+	}
+
+	return b, nil
+}
+
+// View renders the button
+func (b *ButtonCmp) View() string {
+	if b.state == ButtonHovered || b.state == ButtonPressed {
+		return b.hoverStyle.Render(b.label)
+	}
+	return b.style.Render(b.label)
+}
+
diff --git a/internal/tui/components/dialog/permission.go b/internal/tui/components/dialog/permission.go
new file mode 100644
index 000000000..f7581330a
--- /dev/null
+++ b/internal/tui/components/dialog/permission.go
@@ -0,0 +1,167 @@
+package dialog
+
+import (
+	"github.com/charmbracelet/bubbles/key"
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/lipgloss"
+	"github.com/kujtimiihoxha/termai/internal/permission"
+	"github.com/kujtimiihoxha/termai/internal/tui/components/core"
+	"github.com/kujtimiihoxha/termai/internal/tui/layout"
+	"github.com/kujtimiihoxha/termai/internal/tui/styles"
+	"github.com/kujtimiihoxha/termai/internal/tui/util"
+
+	"github.com/charmbracelet/huh"
+)
+
+type PermissionAction string
+
+// Permission responses
+const (
+	PermissionAllow           PermissionAction = "allow"
+	PermissionAllowForSession PermissionAction = "allow_session"
+	PermissionDeny            PermissionAction = "deny"
+)
+
+// PermissionResponseMsg represents the user's response to a permission request
+type PermissionResponseMsg struct {
+	Permission permission.PermissionRequest
+	Action     PermissionAction
+}
+
+// Width and height constants for the dialog
+var (
+	permissionWidth  = 60
+	permissionHeight = 10
+)
+
+// PermissionDialog interface for permission dialog component
+type PermissionDialog interface {
+	tea.Model
+	layout.Sizeable
+	layout.Bindings
+}
+
+// permissionDialogCmp is the implementation of PermissionDialog
+type permissionDialogCmp struct {
+	form       *huh.Form
+	content    string
+	width      int
+	height     int
+	permission permission.PermissionRequest
+}
+
+func (p *permissionDialogCmp) Init() tea.Cmd {
+	return nil
+}
+
+func (p *permissionDialogCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	var cmds []tea.Cmd
+
+	// Process the form
+	form, cmd := p.form.Update(msg)
+	if f, ok := form.(*huh.Form); ok {
+		p.form = f
+		cmds = append(cmds, cmd)
+	}
+
+	if p.form.State == huh.StateCompleted {
+		// Get the selected action
+		action := p.form.GetString("action")
+
+		// Close the dialog and return the response
+		return p, tea.Batch(
+			util.CmdHandler(core.DialogCloseMsg{}),
+			util.CmdHandler(PermissionResponseMsg{Action: PermissionAction(action), Permission: p.permission}),
+		)
+	}
+
+	return p, tea.Batch(cmds...)
+}
+
+func (p *permissionDialogCmp) View() string {
+	contentStyle := lipgloss.NewStyle().
+		Width(p.width).
+		Padding(1, 0).
+		Foreground(styles.Text).
+		Align(lipgloss.Center)
+
+	return lipgloss.JoinVertical(
+		lipgloss.Center,
+		contentStyle.Render(p.content),
+		p.form.View(),
+	)
+}
+
+func (p *permissionDialogCmp) GetSize() (int, int) {
+	return p.width, p.height
+}
+
+func (p *permissionDialogCmp) SetSize(width int, height int) {
+	p.width = width
+	p.height = height
+}
+
+func (p *permissionDialogCmp) BindingKeys() []key.Binding {
+	return p.form.KeyBinds()
+}
+
+func newPermissionDialogCmp(permission permission.PermissionRequest, content string) PermissionDialog {
+	// Create a note field for displaying the content
+
+	// Create select field for the permission options
+	selectOption := huh.NewSelect[string]().
+		Key("action").
+		Options(
+			huh.NewOption("Allow", string(PermissionAllow)),
+			huh.NewOption("Allow for this session", string(PermissionAllowForSession)),
+			huh.NewOption("Deny", string(PermissionDeny)),
+		).
+		Title("Permission Request")
+
+	// Apply theme
+	theme := styles.HuhTheme()
+
+	// Setup form width and height
+	form := huh.NewForm(huh.NewGroup(selectOption)).
+		WithWidth(permissionWidth - 2).
+		WithShowHelp(false).
+		WithTheme(theme).
+		WithShowErrors(false)
+
+	// Focus the form for immediate interaction
+	selectOption.Focus()
+
+	return &permissionDialogCmp{
+		permission: permission,
+		form:       form,
+		content:    content,
+		width:      permissionWidth,
+		height:     permissionHeight,
+	}
+}
+
+// NewPermissionDialogCmd creates a new permission dialog command
+func NewPermissionDialogCmd(permission permission.PermissionRequest, content string) tea.Cmd {
+	permDialog := newPermissionDialogCmp(permission, content)
+
+	// Create the dialog layout
+	dialogPane := layout.NewSinglePane(
+		permDialog.(*permissionDialogCmp),
+		layout.WithSignlePaneSize(permissionWidth+2, permissionHeight+2),
+		layout.WithSinglePaneBordered(true),
+		layout.WithSinglePaneFocusable(true),
+		layout.WithSinglePaneActiveColor(styles.Blue),
+		layout.WithSignlePaneBorderText(map[layout.BorderPosition]string{
+			layout.TopMiddleBorder: " Permission Required ",
+		}),
+	)
+
+	// Focus the dialog
+	dialogPane.Focus()
+
+	// Return the dialog command
+	return util.CmdHandler(core.DialogMsg{
+		Content: dialogPane,
+	})
+}
+
diff --git a/internal/tui/components/messages/message.go b/internal/tui/components/messages/message.go
new file mode 100644
index 000000000..619950850
--- /dev/null
+++ b/internal/tui/components/messages/message.go
@@ -0,0 +1,108 @@
+package messages
+
+import (
+	"fmt"
+
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/lipgloss"
+	"github.com/cloudwego/eino/schema"
+	"github.com/kujtimiihoxha/termai/internal/message"
+	"github.com/kujtimiihoxha/termai/internal/tui/layout"
+	"github.com/kujtimiihoxha/termai/internal/tui/styles"
+)
+
+const (
+	maxHeight = 10
+)
+
+type MessagesCmp interface {
+	tea.Model
+	layout.Focusable
+	layout.Bordered
+	layout.Sizeable
+}
+
+type messageCmp struct {
+	message  message.Message
+	width    int
+	height   int
+	focused  bool
+	expanded bool
+}
+
+func (m *messageCmp) Init() tea.Cmd {
+	return nil
+}
+
+func (m *messageCmp) Update(tea.Msg) (tea.Model, tea.Cmd) {
+	return m, nil
+}
+
+func (m *messageCmp) View() string {
+	wrapper := layout.NewSinglePane(
+		m,
+		layout.WithSinglePaneBordered(true),
+		layout.WithSinglePaneFocusable(true),
+		layout.WithSinglePanePadding(1),
+		layout.WithSinglePaneActiveColor(m.borderColor()),
+	)
+	if m.focused {
+		wrapper.Focus()
+	}
+	wrapper.SetSize(m.width, m.height)
+	return wrapper.View()
+}
+
+func (m *messageCmp) Blur() tea.Cmd {
+	m.focused = false
+	return nil
+}
+
+func (m *messageCmp) borderColor() lipgloss.TerminalColor {
+	switch m.message.MessageData.Role {
+	case schema.Assistant:
+		return styles.Mauve
+	case schema.User:
+		return styles.Flamingo
+	}
+	return styles.Blue
+}
+
+func (m *messageCmp) BorderText() map[layout.BorderPosition]string {
+	role := ""
+	icon := ""
+	switch m.message.MessageData.Role {
+	case schema.Assistant:
+		role = "Assistant"
+		icon = styles.BotIcon
+	case schema.User:
+		role = "User"
+		icon = styles.UserIcon
+	}
+	return map[layout.BorderPosition]string{
+		layout.TopLeftBorder: fmt.Sprintf("%s %s ", role, icon),
+	}
+}
+
+func (m *messageCmp) Focus() tea.Cmd {
+	m.focused = true
+	return nil
+}
+
+func (m *messageCmp) IsFocused() bool {
+	return m.focused
+}
+
+func (m *messageCmp) GetSize() (int, int) {
+	return m.width, 0
+}
+
+func (m *messageCmp) SetSize(width int, height int) {
+	m.width = width
+}
+
+func NewMessageCmp(msg message.Message) MessagesCmp {
+	return &messageCmp{
+		message: msg,
+	}
+}
diff --git a/internal/tui/components/repl/editor.go b/internal/tui/components/repl/editor.go
index 8d795eb14..d0af8d2c5 100644
--- a/internal/tui/components/repl/editor.go
+++ b/internal/tui/components/repl/editor.go
@@ -5,9 +5,11 @@ import (
 
 	"github.com/charmbracelet/bubbles/key"
 	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/lipgloss"
 	"github.com/cloudwego/eino/schema"
 	"github.com/kujtimiihoxha/termai/internal/app"
 	"github.com/kujtimiihoxha/termai/internal/tui/layout"
+	"github.com/kujtimiihoxha/termai/internal/tui/styles"
 	"github.com/kujtimiihoxha/vimtea"
 )
 
@@ -105,8 +107,12 @@ func (m *editorCmp) Blur() tea.Cmd {
 }
 
 func (m *editorCmp) BorderText() map[layout.BorderPosition]string {
+	title := "New Message"
+	if m.focused {
+		title = lipgloss.NewStyle().Foreground(styles.Primary).Render(title)
+	}
 	return map[layout.BorderPosition]string{
-		layout.TopLeftBorder: "New Message",
+		layout.TopLeftBorder: title,
 	}
 }
 
@@ -148,7 +154,9 @@ func (m *editorCmp) BindingKeys() []key.Binding {
 
 func NewEditorCmp(app *app.App) EditorCmp {
 	return &editorCmp{
-		app:    app,
-		editor: vimtea.NewEditor(),
+		app: app,
+		editor: vimtea.NewEditor(
+			vimtea.WithFileName("message.md"),
+		),
 	}
 }
diff --git a/internal/tui/components/repl/messages.go b/internal/tui/components/repl/messages.go
index feddf7bf6..9b3c5bde8 100644
--- a/internal/tui/components/repl/messages.go
+++ b/internal/tui/components/repl/messages.go
@@ -1,15 +1,22 @@
 package repl
 
 import (
+	"fmt"
+	"slices"
+	"strings"
+
 	"github.com/charmbracelet/bubbles/key"
 	"github.com/charmbracelet/bubbles/viewport"
 	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/glamour"
 	"github.com/charmbracelet/lipgloss"
+	"github.com/cloudwego/eino/schema"
 	"github.com/kujtimiihoxha/termai/internal/app"
 	"github.com/kujtimiihoxha/termai/internal/message"
 	"github.com/kujtimiihoxha/termai/internal/pubsub"
 	"github.com/kujtimiihoxha/termai/internal/session"
 	"github.com/kujtimiihoxha/termai/internal/tui/layout"
+	"github.com/kujtimiihoxha/termai/internal/tui/styles"
 )
 
 type MessagesCmp interface {
@@ -21,13 +28,15 @@ type MessagesCmp interface {
 }
 
 type messagesCmp struct {
-	app      *app.App
-	messages []message.Message
-	session  session.Session
-	viewport viewport.Model
-	width    int
-	height   int
-	focused  bool
+	app        *app.App
+	messages   []message.Message
+	session    session.Session
+	viewport   viewport.Model
+	mdRenderer *glamour.TermRenderer
+	width      int
+	height     int
+	focused    bool
+	cachedView string
 }
 
 func (m *messagesCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
@@ -35,6 +44,8 @@ func (m *messagesCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	case pubsub.Event[message.Message]:
 		if msg.Type == pubsub.CreatedEvent {
 			m.messages = append(m.messages, msg.Payload)
+			m.renderView()
+			m.viewport.GotoBottom()
 		}
 	case pubsub.Event[session.Session]:
 		if msg.Type == pubsub.UpdatedEvent {
@@ -45,60 +56,182 @@ func (m *messagesCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	case SelectedSessionMsg:
 		m.session, _ = m.app.Sessions.Get(msg.SessionID)
 		m.messages, _ = m.app.Messages.List(m.session.ID)
+		m.renderView()
+		m.viewport.GotoBottom()
+	}
+	if m.focused {
+		u, cmd := m.viewport.Update(msg)
+		m.viewport = u
+		return m, cmd
 	}
 	return m, nil
 }
 
-func (i *messagesCmp) View() string {
-	stringMessages := make([]string, len(i.messages))
-	for idx, msg := range i.messages {
-		stringMessages[idx] = msg.MessageData.Content
+func borderColor(role schema.RoleType) lipgloss.TerminalColor {
+	switch role {
+	case schema.Assistant:
+		return styles.Mauve
+	case schema.User:
+		return styles.Rosewater
+	case schema.Tool:
+		return styles.Peach
 	}
-	return lipgloss.JoinVertical(lipgloss.Top, stringMessages...)
+	return styles.Blue
+}
+
+func borderText(msgRole schema.RoleType, currentMessage int) map[layout.BorderPosition]string {
+	role := ""
+	icon := ""
+	switch msgRole {
+	case schema.Assistant:
+		role = "Assistant"
+		icon = styles.BotIcon
+	case schema.User:
+		role = "User"
+		icon = styles.UserIcon
+	}
+	return map[layout.BorderPosition]string{
+		layout.TopLeftBorder: lipgloss.NewStyle().
+			Padding(0, 1).
+			Bold(true).
+			Foreground(styles.Crust).
+			Background(borderColor(msgRole)).
+			Render(fmt.Sprintf("%s %s ", role, icon)),
+		layout.TopRightBorder: lipgloss.NewStyle().
+			Padding(0, 1).
+			Bold(true).
+			Foreground(styles.Crust).
+			Background(borderColor(msgRole)).
+			Render(fmt.Sprintf("#%d ", currentMessage)),
+	}
+}
+
+func (m *messagesCmp) renderView() {
+	stringMessages := make([]string, 0)
+	r, _ := glamour.NewTermRenderer(
+		glamour.WithStyles(styles.CatppuccinMarkdownStyle()),
+		glamour.WithWordWrap(m.width-10),
+		glamour.WithEmoji(),
+	)
+	textStyle := lipgloss.NewStyle().Width(m.width - 4)
+	currentMessage := 1
+	for _, msg := range m.messages {
+		if msg.MessageData.Role == schema.Tool {
+			continue
+		}
+		content := msg.MessageData.Content
+		if content != "" {
+			content, _ = r.Render(msg.MessageData.Content)
+			stringMessages = append(stringMessages, layout.Borderize(
+				textStyle.Render(content),
+				layout.BorderOptions{
+					InactiveBorder: lipgloss.DoubleBorder(),
+					ActiveBorder:   lipgloss.DoubleBorder(),
+					ActiveColor:    borderColor(msg.MessageData.Role),
+					InactiveColor:  borderColor(msg.MessageData.Role),
+					EmbeddedText:   borderText(msg.MessageData.Role, currentMessage),
+				},
+			))
+			currentMessage++
+		}
+		for _, toolCall := range msg.MessageData.ToolCalls {
+			resultInx := slices.IndexFunc(m.messages, func(m message.Message) bool {
+				return m.MessageData.ToolCallID == toolCall.ID
+			})
+			content := fmt.Sprintf("**Arguments**\n```json\n%s\n```\n", toolCall.Function.Arguments)
+			if resultInx == -1 {
+				content += "Running..."
+			} else {
+				result := m.messages[resultInx].MessageData.Content
+				if result != "" {
+					lines := strings.Split(result, "\n")
+					if len(lines) > 15 {
+						result = strings.Join(lines[:15], "\n")
+					}
+					content += fmt.Sprintf("**Result**\n```\n%s\n```\n", result)
+					if len(lines) > 15 {
+						content += fmt.Sprintf("\n\n *...%d lines are truncated* ", len(lines)-15)
+					}
+				}
+			}
+			content, _ = r.Render(content)
+			stringMessages = append(stringMessages, layout.Borderize(
+				textStyle.Render(content),
+				layout.BorderOptions{
+					InactiveBorder: lipgloss.DoubleBorder(),
+					ActiveBorder:   lipgloss.DoubleBorder(),
+					ActiveColor:    borderColor(schema.Tool),
+					InactiveColor:  borderColor(schema.Tool),
+					EmbeddedText: map[layout.BorderPosition]string{
+						layout.TopLeftBorder: lipgloss.NewStyle().
+							Padding(0, 1).
+							Bold(true).
+							Foreground(styles.Crust).
+							Background(borderColor(schema.Tool)).
+							Render(
+								fmt.Sprintf("Tool [%s] %s ", toolCall.Function.Name, styles.ToolIcon),
+							),
+						layout.TopRightBorder: lipgloss.NewStyle().
+							Padding(0, 1).
+							Bold(true).
+							Foreground(styles.Crust).
+							Background(borderColor(schema.Tool)).
+							Render(fmt.Sprintf("#%d ", currentMessage)),
+					},
+				},
+			))
+			currentMessage++
+		}
+	}
+	m.viewport.SetContent(lipgloss.JoinVertical(lipgloss.Top, stringMessages...))
+}
+
+func (m *messagesCmp) View() string {
+	return lipgloss.NewStyle().Padding(1).Render(m.viewport.View())
 }
 
-// BindingKeys implements MessagesCmp.
 func (m *messagesCmp) BindingKeys() []key.Binding {
-	return []key.Binding{}
+	return layout.KeyMapToSlice(m.viewport.KeyMap)
 }
 
-// Blur implements MessagesCmp.
 func (m *messagesCmp) Blur() tea.Cmd {
 	m.focused = false
 	return nil
 }
 
-// BorderText implements MessagesCmp.
 func (m *messagesCmp) BorderText() map[layout.BorderPosition]string {
 	title := m.session.Title
-	if len(title) > 20 {
-		title = title[:20] + "..."
+	titleWidth := m.width / 2
+	if len(title) > titleWidth {
+		title = title[:titleWidth] + "..."
+	}
+	if m.focused {
+		title = lipgloss.NewStyle().Foreground(styles.Primary).Render(title)
 	}
 	return map[layout.BorderPosition]string{
-		layout.TopLeftBorder: title,
+		layout.TopLeftBorder:     title,
+		layout.BottomRightBorder: formatTokensAndCost(m.session.CompletionTokens+m.session.PromptTokens, m.session.Cost),
 	}
 }
 
-// Focus implements MessagesCmp.
 func (m *messagesCmp) Focus() tea.Cmd {
 	m.focused = true
 	return nil
 }
 
-// GetSize implements MessagesCmp.
 func (m *messagesCmp) GetSize() (int, int) {
 	return m.width, m.height
 }
 
-// IsFocused implements MessagesCmp.
 func (m *messagesCmp) IsFocused() bool {
 	return m.focused
 }
 
-// SetSize implements MessagesCmp.
 func (m *messagesCmp) SetSize(width int, height int) {
 	m.width = width
 	m.height = height
+	m.viewport.Width = width - 2   // padding
+	m.viewport.Height = height - 2 // padding
 }
 
 func (m *messagesCmp) Init() tea.Cmd {
@@ -109,5 +242,6 @@ func NewMessagesCmp(app *app.App) MessagesCmp {
 	return &messagesCmp{
 		app:      app,
 		messages: []message.Message{},
+		viewport: viewport.New(0, 0),
 	}
 }
diff --git a/internal/tui/components/repl/sessions.go b/internal/tui/components/repl/sessions.go
index 5d2411fb6..0f208ced9 100644
--- a/internal/tui/components/repl/sessions.go
+++ b/internal/tui/components/repl/sessions.go
@@ -7,6 +7,7 @@ import (
 	"github.com/charmbracelet/bubbles/key"
 	"github.com/charmbracelet/bubbles/list"
 	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/lipgloss"
 	"github.com/kujtimiihoxha/termai/internal/app"
 	"github.com/kujtimiihoxha/termai/internal/pubsub"
 	"github.com/kujtimiihoxha/termai/internal/session"
@@ -160,8 +161,13 @@ func (i *sessionsCmp) BorderText() map[layout.BorderPosition]string {
 		current,
 		totalCount,
 	)
+
+	title := "Sessions"
+	if i.focused {
+		title = lipgloss.NewStyle().Foreground(styles.Primary).Render(title)
+	}
 	return map[layout.BorderPosition]string{
-		layout.TopMiddleBorder:    "Sessions",
+		layout.TopMiddleBorder:    title,
 		layout.BottomMiddleBorder: pageInfo,
 	}
 }
diff --git a/internal/tui/layout/border.go b/internal/tui/layout/border.go
index a3c80396f..8fe5c430c 100644
--- a/internal/tui/layout/border.go
+++ b/internal/tui/layout/border.go
@@ -24,24 +24,43 @@ var (
 	InactivePreviewBorder = styles.Grey
 )
 
-func Borderize(content string, active bool, embeddedText map[BorderPosition]string, activeColor lipgloss.TerminalColor) string {
-	if embeddedText == nil {
-		embeddedText = make(map[BorderPosition]string)
+type BorderOptions struct {
+	Active         bool
+	EmbeddedText   map[BorderPosition]string
+	ActiveColor    lipgloss.TerminalColor
+	InactiveColor  lipgloss.TerminalColor
+	ActiveBorder   lipgloss.Border
+	InactiveBorder lipgloss.Border
+}
+
+func Borderize(content string, opts BorderOptions) string {
+	if opts.EmbeddedText == nil {
+		opts.EmbeddedText = make(map[BorderPosition]string)
+	}
+	if opts.ActiveColor == nil {
+		opts.ActiveColor = ActiveBorder
 	}
-	if activeColor == nil {
-		activeColor = ActiveBorder
+	if opts.InactiveColor == nil {
+		opts.InactiveColor = InactivePreviewBorder
 	}
+	if opts.ActiveBorder == (lipgloss.Border{}) {
+		opts.ActiveBorder = lipgloss.ThickBorder()
+	}
+	if opts.InactiveBorder == (lipgloss.Border{}) {
+		opts.InactiveBorder = lipgloss.NormalBorder()
+	}
+
 	var (
 		thickness = map[bool]lipgloss.Border{
-			true:  lipgloss.Border(lipgloss.ThickBorder()),
-			false: lipgloss.Border(lipgloss.NormalBorder()),
+			true:  opts.ActiveBorder,
+			false: opts.InactiveBorder,
 		}
 		color = map[bool]lipgloss.TerminalColor{
-			true:  activeColor,
-			false: InactivePreviewBorder,
+			true:  opts.ActiveColor,
+			false: opts.InactiveColor,
 		}
-		border = thickness[active]
-		style  = lipgloss.NewStyle().Foreground(color[active])
+		border = thickness[opts.Active]
+		style  = lipgloss.NewStyle().Foreground(color[opts.Active])
 		width  = lipgloss.Width(content)
 	)
 
@@ -80,20 +99,20 @@ func Borderize(content string, active bool, embeddedText map[BorderPosition]stri
 	// Stack top border, content and horizontal borders, and bottom border.
 	return strings.Join([]string{
 		buildHorizontalBorder(
-			embeddedText[TopLeftBorder],
-			embeddedText[TopMiddleBorder],
-			embeddedText[TopRightBorder],
+			opts.EmbeddedText[TopLeftBorder],
+			opts.EmbeddedText[TopMiddleBorder],
+			opts.EmbeddedText[TopRightBorder],
 			border.TopLeft,
 			border.Top,
 			border.TopRight,
 		),
 		lipgloss.NewStyle().
-			BorderForeground(color[active]).
+			BorderForeground(color[opts.Active]).
 			Border(border, false, true, false, true).Render(content),
 		buildHorizontalBorder(
-			embeddedText[BottomLeftBorder],
-			embeddedText[BottomMiddleBorder],
-			embeddedText[BottomRightBorder],
+			opts.EmbeddedText[BottomLeftBorder],
+			opts.EmbeddedText[BottomMiddleBorder],
+			opts.EmbeddedText[BottomRightBorder],
 			border.BottomLeft,
 			border.Bottom,
 			border.BottomRight,
diff --git a/internal/tui/layout/grid.go b/internal/tui/layout/grid.go
new file mode 100644
index 000000000..d6f0b4ab9
--- /dev/null
+++ b/internal/tui/layout/grid.go
@@ -0,0 +1,254 @@
+package layout
+
+import (
+	"github.com/charmbracelet/bubbles/key"
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/lipgloss"
+)
+
+type GridLayout interface {
+	tea.Model
+	Sizeable
+	Bindings
+	Panes() [][]tea.Model
+}
+
+type gridLayout struct {
+	width  int
+	height int
+
+	rows    int
+	columns int
+
+	panes [][]tea.Model
+
+	gap       int
+	bordered  bool
+	focusable bool
+
+	currentRow    int
+	currentColumn int
+
+	activeColor lipgloss.TerminalColor
+}
+
+type GridOption func(*gridLayout)
+
+func (g *gridLayout) Init() tea.Cmd {
+	var cmds []tea.Cmd
+	for i := range g.panes {
+		for j := range g.panes[i] {
+			if g.panes[i][j] != nil {
+				cmds = append(cmds, g.panes[i][j].Init())
+			}
+		}
+	}
+	return tea.Batch(cmds...)
+}
+
+func (g *gridLayout) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	var cmds []tea.Cmd
+
+	switch msg := msg.(type) {
+	case tea.WindowSizeMsg:
+		g.SetSize(msg.Width, msg.Height)
+		return g, nil
+	case tea.KeyMsg:
+		if key.Matches(msg, g.nextPaneBinding()) {
+			return g.focusNextPane()
+		}
+	}
+
+	// Update all panes
+	for i := range g.panes {
+		for j := range g.panes[i] {
+			if g.panes[i][j] != nil {
+				var cmd tea.Cmd
+				g.panes[i][j], cmd = g.panes[i][j].Update(msg)
+				if cmd != nil {
+					cmds = append(cmds, cmd)
+				}
+			}
+		}
+	}
+
+	return g, tea.Batch(cmds...)
+}
+
+func (g *gridLayout) focusNextPane() (tea.Model, tea.Cmd) {
+	if !g.focusable {
+		return g, nil
+	}
+
+	var cmds []tea.Cmd
+
+	// Blur current pane
+	if g.currentRow < len(g.panes) && g.currentColumn < len(g.panes[g.currentRow]) {
+		if currentPane, ok := g.panes[g.currentRow][g.currentColumn].(Focusable); ok {
+			cmds = append(cmds, currentPane.Blur())
+		}
+	}
+
+	// Find next valid pane
+	g.currentColumn++
+	if g.currentColumn >= len(g.panes[g.currentRow]) {
+		g.currentColumn = 0
+		g.currentRow++
+		if g.currentRow >= len(g.panes) {
+			g.currentRow = 0
+		}
+	}
+
+	// Focus next pane
+	if g.currentRow < len(g.panes) && g.currentColumn < len(g.panes[g.currentRow]) {
+		if nextPane, ok := g.panes[g.currentRow][g.currentColumn].(Focusable); ok {
+			cmds = append(cmds, nextPane.Focus())
+		}
+	}
+
+	return g, tea.Batch(cmds...)
+}
+
+func (g *gridLayout) nextPaneBinding() key.Binding {
+	return key.NewBinding(
+		key.WithKeys("tab"),
+		key.WithHelp("tab", "next pane"),
+	)
+}
+
+func (g *gridLayout) View() string {
+	if len(g.panes) == 0 {
+		return ""
+	}
+
+	// Calculate dimensions for each cell
+	cellWidth := (g.width - (g.columns-1)*g.gap) / g.columns
+	cellHeight := (g.height - (g.rows-1)*g.gap) / g.rows
+
+	// Render each row
+	rows := make([]string, g.rows)
+	for i := 0; i < g.rows; i++ {
+		// Render each column in this row
+		cols := make([]string, len(g.panes[i]))
+		for j := 0; j < len(g.panes[i]); j++ {
+			if g.panes[i][j] == nil {
+				cols[j] = ""
+				continue
+			}
+
+			// Set size for each pane
+			if sizable, ok := g.panes[i][j].(Sizeable); ok {
+				effectiveWidth, effectiveHeight := cellWidth, cellHeight
+				if g.bordered {
+					effectiveWidth -= 2
+					effectiveHeight -= 2
+				}
+				sizable.SetSize(effectiveWidth, effectiveHeight)
+			}
+
+			// Render the pane
+			content := g.panes[i][j].View()
+			
+			// Apply border if needed
+			if g.bordered {
+				isFocused := false
+				if focusable, ok := g.panes[i][j].(Focusable); ok {
+					isFocused = focusable.IsFocused()
+				}
+				
+				borderText := map[BorderPosition]string{}
+				if bordered, ok := g.panes[i][j].(Bordered); ok {
+					borderText = bordered.BorderText()
+				}
+				
+				content = Borderize(content, BorderOptions{
+					Active:       isFocused,
+					EmbeddedText: borderText,
+				})
+			}
+			
+			cols[j] = content
+		}
+		
+		// Join columns with gap
+		rows[i] = lipgloss.JoinHorizontal(lipgloss.Top, cols...)
+	}
+	
+	// Join rows with gap
+	return lipgloss.JoinVertical(lipgloss.Left, rows...)
+}
+
+func (g *gridLayout) SetSize(width, height int) {
+	g.width = width
+	g.height = height
+}
+
+func (g *gridLayout) GetSize() (int, int) {
+	return g.width, g.height
+}
+
+func (g *gridLayout) BindingKeys() []key.Binding {
+	var bindings []key.Binding
+	bindings = append(bindings, g.nextPaneBinding())
+	
+	// Collect bindings from all panes
+	for i := range g.panes {
+		for j := range g.panes[i] {
+			if g.panes[i][j] != nil {
+				if bindable, ok := g.panes[i][j].(Bindings); ok {
+					bindings = append(bindings, bindable.BindingKeys()...)
+				}
+			}
+		}
+	}
+	
+	return bindings
+}
+
+func (g *gridLayout) Panes() [][]tea.Model {
+	return g.panes
+}
+
+// NewGridLayout creates a new grid layout with the given number of rows and columns
+func NewGridLayout(rows, cols int, panes [][]tea.Model, opts ...GridOption) GridLayout {
+	grid := &gridLayout{
+		rows:    rows,
+		columns: cols,
+		panes:   panes,
+		gap:     1,
+	}
+	
+	for _, opt := range opts {
+		opt(grid)
+	}
+	
+	return grid
+}
+
+// WithGridGap sets the gap between cells
+func WithGridGap(gap int) GridOption {
+	return func(g *gridLayout) {
+		g.gap = gap
+	}
+}
+
+// WithGridBordered sets whether cells should have borders
+func WithGridBordered(bordered bool) GridOption {
+	return func(g *gridLayout) {
+		g.bordered = bordered
+	}
+}
+
+// WithGridFocusable sets whether the grid supports focus navigation
+func WithGridFocusable(focusable bool) GridOption {
+	return func(g *gridLayout) {
+		g.focusable = focusable
+	}
+}
+
+// WithGridActiveColor sets the active border color
+func WithGridActiveColor(color lipgloss.TerminalColor) GridOption {
+	return func(g *gridLayout) {
+		g.activeColor = color
+	}
+}
diff --git a/internal/tui/layout/single.go b/internal/tui/layout/single.go
index 1e4d0881c..e5c9a61c4 100644
--- a/internal/tui/layout/single.go
+++ b/internal/tui/layout/single.go
@@ -64,7 +64,10 @@ func (s *singlePaneLayout) View() string {
 		if bordered, ok := s.content.(Bordered); ok {
 			s.borderText = bordered.BorderText()
 		}
-		return Borderize(content, s.focused, s.borderText, s.activeColor)
+		return Borderize(content, BorderOptions{
+			Active:       s.focused,
+			EmbeddedText: s.borderText,
+		})
 	}
 	return content
 }
diff --git a/internal/tui/page/init.go b/internal/tui/page/init.go
index 884aaf9f9..a010d14da 100644
--- a/internal/tui/page/init.go
+++ b/internal/tui/page/init.go
@@ -1,36 +1,307 @@
 package page
 
 import (
+	"fmt"
+	"os"
+	"path/filepath"
+	"strconv"
+
+	"github.com/charmbracelet/bubbles/key"
 	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/huh"
+	"github.com/charmbracelet/lipgloss"
+	"github.com/kujtimiihoxha/termai/internal/llm/models"
 	"github.com/kujtimiihoxha/termai/internal/tui/layout"
+	"github.com/kujtimiihoxha/termai/internal/tui/styles"
+	"github.com/kujtimiihoxha/termai/internal/tui/util"
+	"github.com/spf13/viper"
 )
 
 var InitPage PageID = "init"
 
+type configSaved struct{}
+
 type initPage struct {
-	layout layout.SinglePaneLayout
+	form         *huh.Form
+	width        int
+	height       int
+	saved        bool
+	errorMsg     string
+	statusMsg    string
+	modelOpts    []huh.Option[string]
+	bigModel     string
+	smallModel   string
+	openAIKey    string
+	anthropicKey string
+	groqKey      string
+	maxTokens    string
+	dataDir      string
+	agent        string
+}
+
+func (i *initPage) Init() tea.Cmd {
+	return i.form.Init()
+}
+
+func (i *initPage) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	var cmds []tea.Cmd
+
+	switch msg := msg.(type) {
+	case tea.WindowSizeMsg:
+		i.width = msg.Width - 4 // Account for border
+		i.height = msg.Height - 4
+		i.form = i.form.WithWidth(i.width).WithHeight(i.height)
+		return i, nil
+
+	case configSaved:
+		i.saved = true
+		i.statusMsg = "Configuration saved successfully. Press any key to continue."
+		return i, nil
+	}
+
+	if i.saved {
+		switch msg.(type) {
+		case tea.KeyMsg:
+			return i, util.CmdHandler(PageChangeMsg{ID: ReplPage})
+		}
+		return i, nil
+	}
+
+	// Process the form
+	form, cmd := i.form.Update(msg)
+	if f, ok := form.(*huh.Form); ok {
+		i.form = f
+		cmds = append(cmds, cmd)
+	}
+
+	if i.form.State == huh.StateCompleted {
+		// Save configuration to file
+		configPath := filepath.Join(os.Getenv("HOME"), ".termai.yaml")
+		maxTokens, _ := strconv.Atoi(i.maxTokens)
+		config := map[string]interface{}{
+			"models": map[string]string{
+				"big":   i.bigModel,
+				"small": i.smallModel,
+			},
+			"providers": map[string]interface{}{
+				"openai": map[string]string{
+					"key": i.openAIKey,
+				},
+				"anthropic": map[string]string{
+					"key": i.anthropicKey,
+				},
+				"groq": map[string]string{
+					"key": i.groqKey,
+				},
+				"common": map[string]int{
+					"max_tokens": maxTokens,
+				},
+			},
+			"data": map[string]string{
+				"dir": i.dataDir,
+			},
+			"agents": map[string]string{
+				"default": i.agent,
+			},
+			"log": map[string]string{
+				"level": "info",
+			},
+		}
+
+		// Write config to viper
+		for k, v := range config {
+			viper.Set(k, v)
+		}
+
+		// Save configuration
+		err := viper.WriteConfigAs(configPath)
+		if err != nil {
+			i.errorMsg = fmt.Sprintf("Failed to save configuration: %s", err)
+			return i, nil
+		}
+
+		// Return to main page
+		return i, util.CmdHandler(configSaved{})
+	}
+
+	return i, tea.Batch(cmds...)
 }
 
-func (i initPage) Init() tea.Cmd {
-	return nil
+func (i *initPage) View() string {
+	if i.saved {
+		return lipgloss.NewStyle().
+			Width(i.width).
+			Height(i.height).
+			Align(lipgloss.Center, lipgloss.Center).
+			Render(lipgloss.JoinVertical(
+				lipgloss.Center,
+				lipgloss.NewStyle().Foreground(styles.Green).Render("✓ Configuration Saved"),
+				"",
+				lipgloss.NewStyle().Foreground(styles.Blue).Render(i.statusMsg),
+			))
+	}
+
+	view := i.form.View()
+	if i.errorMsg != "" {
+		errorBox := lipgloss.NewStyle().
+			Padding(1).
+			Border(lipgloss.RoundedBorder()).
+			BorderForeground(styles.Red).
+			Width(i.width - 4).
+			Render(i.errorMsg)
+		view = lipgloss.JoinVertical(lipgloss.Left, errorBox, view)
+	}
+	return view
 }
 
-func (i initPage) Update(_ tea.Msg) (tea.Model, tea.Cmd) {
-	return i, nil
+func (i *initPage) GetSize() (int, int) {
+	return i.width, i.height
 }
 
-func (i initPage) View() string {
-	return "Initializing..."
+func (i *initPage) SetSize(width int, height int) {
+	i.width = width
+	i.height = height
+	i.form = i.form.WithWidth(width).WithHeight(height)
+}
+
+func (i *initPage) BindingKeys() []key.Binding {
+	if i.saved {
+		return []key.Binding{
+			key.NewBinding(
+				key.WithKeys("enter", "space", "esc"),
+				key.WithHelp("any key", "continue"),
+			),
+		}
+	}
+	return i.form.KeyBinds()
 }
 
 func NewInitPage() tea.Model {
+	// Create model options
+	var modelOpts []huh.Option[string]
+	for id, model := range models.SupportedModels {
+		modelOpts = append(modelOpts, huh.NewOption(model.Name, string(id)))
+	}
+
+	// Create agent options
+	agentOpts := []huh.Option[string]{
+		huh.NewOption("Coder", "coder"),
+		huh.NewOption("Assistant", "assistant"),
+	}
+
+	// Init page with form
+	initModel := &initPage{
+		modelOpts:  modelOpts,
+		bigModel:   string(models.DefaultBigModel),
+		smallModel: string(models.DefaultLittleModel),
+		maxTokens:  "4000",
+		dataDir:    ".termai",
+		agent:      "coder",
+	}
+
+	// API Keys group
+	apiKeysGroup := huh.NewGroup(
+		huh.NewNote().
+			Title("API Keys").
+			Description("You need to provide at least one API key to use termai"),
+
+		huh.NewInput().
+			Title("OpenAI API Key").
+			Placeholder("sk-...").
+			Key("openai_key").
+			Value(&initModel.openAIKey),
+
+		huh.NewInput().
+			Title("Anthropic API Key").
+			Placeholder("sk-ant-...").
+			Key("anthropic_key").
+			Value(&initModel.anthropicKey),
+
+		huh.NewInput().
+			Title("Groq API Key").
+			Placeholder("gsk_...").
+			Key("groq_key").
+			Value(&initModel.groqKey),
+	)
+
+	// Model configuration group
+	modelsGroup := huh.NewGroup(
+		huh.NewNote().
+			Title("Model Configuration").
+			Description("Select which models to use"),
+
+		huh.NewSelect[string]().
+			Title("Big Model").
+			Options(modelOpts...).
+			Key("big_model").
+			Value(&initModel.bigModel),
+
+		huh.NewSelect[string]().
+			Title("Small Model").
+			Options(modelOpts...).
+			Key("small_model").
+			Value(&initModel.smallModel),
+
+		huh.NewInput().
+			Title("Max Tokens").
+			Placeholder("4000").
+			Key("max_tokens").
+			CharLimit(5).
+			Validate(func(s string) error {
+				var n int
+				_, err := fmt.Sscanf(s, "%d", &n)
+				if err != nil || n <= 0 {
+					return fmt.Errorf("must be a positive number")
+				}
+				initModel.maxTokens = s
+				return nil
+			}).
+			Value(&initModel.maxTokens),
+	)
+
+	// General settings group
+	generalGroup := huh.NewGroup(
+		huh.NewNote().
+			Title("General Settings").
+			Description("Configure general termai settings"),
+
+		huh.NewInput().
+			Title("Data Directory").
+			Placeholder(".termai").
+			Key("data_dir").
+			Value(&initModel.dataDir),
+
+		huh.NewSelect[string]().
+			Title("Default Agent").
+			Options(agentOpts...).
+			Key("agent").
+			Value(&initModel.agent),
+
+		huh.NewConfirm().
+			Title("Save Configuration").
+			Affirmative("Save").
+			Negative("Cancel"),
+	)
+
+	// Create form with theme
+	form := huh.NewForm(
+		apiKeysGroup,
+		modelsGroup,
+		generalGroup,
+	).WithTheme(styles.HuhTheme()).
+		WithShowHelp(true).
+		WithShowErrors(true)
+
+	// Set the form in the model
+	initModel.form = form
+
 	return layout.NewSinglePane(
-		&initPage{},
+		initModel,
 		layout.WithSinglePaneFocusable(true),
 		layout.WithSinglePaneBordered(true),
 		layout.WithSignlePaneBorderText(
 			map[layout.BorderPosition]string{
-				layout.TopMiddleBorder: "Welcome to termai",
+				layout.TopMiddleBorder: "Welcome to termai - Initial Setup",
 			},
 		),
 	)
diff --git a/internal/tui/page/page.go b/internal/tui/page/page.go
index 578e0aa9a..482df5fd7 100644
--- a/internal/tui/page/page.go
+++ b/internal/tui/page/page.go
@@ -1,3 +1,8 @@
 package page
 
 type PageID string
+
+// PageChangeMsg is used to change the current page
+type PageChangeMsg struct {
+	ID PageID
+}
diff --git a/internal/tui/styles/markdown.go b/internal/tui/styles/markdown.go
index cadb9a2e3..8c8000715 100644
--- a/internal/tui/styles/markdown.go
+++ b/internal/tui/styles/markdown.go
@@ -146,8 +146,8 @@ var catppuccinDark = ansi.StyleConfig{
 	Code: ansi.StyleBlock{
 		StylePrimitive: ansi.StylePrimitive{
 			Color:  stringPtr(dark.Green().Hex),
-			Prefix: " ",
-			Suffix: " ",
+			Prefix: "",
+			Suffix: "",
 		},
 	},
 	CodeBlock: ansi.StyleCodeBlock{
diff --git a/internal/tui/styles/styles.go b/internal/tui/styles/styles.go
index 807521d6f..307079c32 100644
--- a/internal/tui/styles/styles.go
+++ b/internal/tui/styles/styles.go
@@ -20,6 +20,7 @@ var (
 	DoubleBorder = Regular.Border(lipgloss.DoubleBorder())
 
 	// Colors
+	White = lipgloss.Color("#ffffff")
 
 	Surface0 = lipgloss.AdaptiveColor{
 		Dark:  dark.Surface0().Hex,
diff --git a/internal/tui/tui.go b/internal/tui/tui.go
index fcc03657c..0d1d15011 100644
--- a/internal/tui/tui.go
+++ b/internal/tui/tui.go
@@ -1,13 +1,17 @@
 package tui
 
 import (
+	"fmt"
 	"log"
+	"os"
+	"path/filepath"
 
 	"github.com/charmbracelet/bubbles/key"
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/lipgloss"
 	"github.com/kujtimiihoxha/termai/internal/app"
 	"github.com/kujtimiihoxha/termai/internal/llm"
+	"github.com/kujtimiihoxha/termai/internal/permission"
 	"github.com/kujtimiihoxha/termai/internal/pubsub"
 	"github.com/kujtimiihoxha/termai/internal/tui/components/core"
 	"github.com/kujtimiihoxha/termai/internal/tui/components/dialog"
@@ -48,6 +52,11 @@ var keys = keyMap{
 	),
 }
 
+var editorKeyMap = key.NewBinding(
+	key.WithKeys("i"),
+	key.WithHelp("i", "insert mode"),
+)
+
 type appModel struct {
 	width, height int
 	currentPage   page.PageID
@@ -71,8 +80,27 @@ func (a appModel) Init() tea.Cmd {
 func (a appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	switch msg := msg.(type) {
 	case pubsub.Event[llm.AgentEvent]:
-		log.Println("Event received")
+		log.Println("AgentEvent")
 		log.Println(msg)
+	case pubsub.Event[permission.PermissionRequest]:
+		return a, dialog.NewPermissionDialogCmd(
+			msg.Payload,
+			fmt.Sprintf(
+				"Tool: %s\nAction: %s\nParams: %v",
+				msg.Payload.ToolName,
+				msg.Payload.Action,
+				msg.Payload.Params,
+			),
+		)
+	case dialog.PermissionResponseMsg:
+		switch msg.Action {
+		case dialog.PermissionAllow:
+			permission.Default.Grant(msg.Permission)
+		case dialog.PermissionAllowForSession:
+			permission.Default.GrantPersistant(msg.Permission)
+		case dialog.PermissionDeny:
+			permission.Default.Deny(msg.Permission)
+		}
 	case vimtea.EditorModeMsg:
 		a.editorMode = msg.Mode
 	case tea.WindowSizeMsg:
@@ -97,6 +125,8 @@ func (a appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		a.dialog = d.(core.DialogCmp)
 		a.dialogVisible = false
 		return a, cmd
+	case page.PageChangeMsg:
+		return a, a.moveToPage(msg.ID)
 	case util.InfoMsg:
 		a.status, _ = a.status.Update(msg)
 	case util.ErrorMsg:
@@ -201,8 +231,17 @@ func (a appModel) View() string {
 }
 
 func New(app *app.App) tea.Model {
+	// Check if config file exists, if not, start with init page
+	homedir, _ := os.UserHomeDir()
+	configPath := filepath.Join(homedir, ".termai.yaml")
+
+	startPage := page.ReplPage
+	if _, err := os.Stat(configPath); os.IsNotExist(err) {
+		startPage = page.InitPage
+	}
+
 	return &appModel{
-		currentPage: page.ReplPage,
+		currentPage: startPage,
 		loadedPages: make(map[page.PageID]bool),
 		status:      core.NewStatusCmp(),
 		help:        core.NewHelpCmp(),