cremote/browser/browser.go

package browser

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/go-rod/rod"
	"github.com/go-rod/rod/lib/launcher"
	"github.com/go-rod/rod/lib/proto"
)

// Manager handles the connection to a Chrome browser instance
type Manager struct {
	browser      *rod.Browser
	tabs         map[string]*rod.Page
	mu           sync.Mutex
	isNewBrowser bool        // Tracks if we launched a new browser or connected to existing one
	storage      *TabStorage // Persistent storage for tab IDs
}

// NewManager creates a new browser manager
// If launchNew is true, it will launch a new browser instance
// Otherwise, it will try to connect to an existing browser instance
func NewManager(launchNew bool) (*Manager, error) {
	var browser *rod.Browser
	var isNewBrowser bool

	// Initialize tab storage
	storage, err := NewTabStorage()
	if err != nil {
		return nil, fmt.Errorf("failed to initialize tab storage: %w", err)
	}

	if launchNew {
		// Launch a new browser instance
		u := launcher.New().MustLaunch()
		browser = rod.New().ControlURL(u).MustConnect()
		isNewBrowser = true
	} else {
		// Connect to an existing browser instance
		// This assumes Chrome is running with --remote-debugging-port=9222
		u := launcher.MustResolveURL("")
		browser = rod.New().ControlURL(u)

		err := browser.Connect()
		if err != nil {
			return nil, fmt.Errorf("failed to connect to browser: %w\nMake sure Chrome is running with --remote-debugging-port=9222", err)
		}
		isNewBrowser = false
	}

	return &Manager{
		browser:      browser,
		tabs:         make(map[string]*rod.Page),
		isNewBrowser: isNewBrowser,
		storage:      storage,
	}, nil
}

// OpenTab opens a new tab and returns its ID
func (m *Manager) OpenTab() (string, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	page, err := m.browser.Page(proto.TargetCreateTarget{URL: "about:blank"})
	if err != nil {
		return "", fmt.Errorf("failed to create new tab: %w", err)
	}

	// Use the page ID as the tab ID
	tabID := string(page.TargetID)
	m.tabs[tabID] = page

	// Save the tab ID to persistent storage
	err = m.storage.SaveTab(tabID)
	if err != nil {
		return "", fmt.Errorf("failed to save tab ID: %w", err)
	}

	return tabID, nil
}

// GetTab returns a tab by its ID
func (m *Manager) GetTab(tabID string) (*rod.Page, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	// First check in-memory cache
	page, exists := m.tabs[tabID]
	if exists {
		return page, nil
	}

	// If not in memory, check persistent storage
	storedID, exists := m.storage.GetTab(tabID)
	if !exists {
		return nil, errors.New("tab not found")
	}

	// Try to get the page from the browser
	pages, err := m.browser.Pages()
	if err != nil {
		return nil, fmt.Errorf("failed to get browser pages: %w", err)
	}

	// Find the page with the matching ID
	for _, p := range pages {
		if string(p.TargetID) == storedID {
			// Cache it for future use
			m.tabs[tabID] = p
			return p, nil
		}
	}

	// If we get here, the tab no longer exists
	m.storage.RemoveTab(tabID)
	return nil, errors.New("tab not found or was closed")
}

// CloseTab closes a tab by its ID
func (m *Manager) CloseTab(tabID string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	// First check in-memory cache
	page, exists := m.tabs[tabID]
	if !exists {
		// If not in memory, check persistent storage
		storedID, exists := m.storage.GetTab(tabID)
		if !exists {
			return errors.New("tab not found")
		}

		// Try to get the page from the browser
		pages, err := m.browser.Pages()
		if err != nil {
			return fmt.Errorf("failed to get browser pages: %w", err)
		}

		// Find the page with the matching ID
		for _, p := range pages {
			if string(p.TargetID) == storedID {
				page = p
				exists = true
				break
			}
		}

		if !exists {
			// If we get here, the tab no longer exists, so just remove it from storage
			m.storage.RemoveTab(tabID)
			return errors.New("tab not found or was already closed")
		}
	}

	err := page.Close()
	if err != nil {
		return fmt.Errorf("failed to close tab: %w", err)
	}

	// Remove from in-memory cache and persistent storage
	delete(m.tabs, tabID)
	m.storage.RemoveTab(tabID)
	return nil
}

// Close closes the browser connection and all tabs
func (m *Manager) Close() error {
	m.mu.Lock()
	defer m.mu.Unlock()

	// Clear the tabs map
	m.tabs = make(map[string]*rod.Page)

	// Only close the browser if we launched it
	if m.isNewBrowser {
		return m.browser.Close()
	}

	// For existing browsers, just disconnect without closing
	return nil
}

// LoadURL loads a URL in a tab
func (m *Manager) LoadURL(tabID, url string) error {
	page, err := m.GetTab(tabID)
	if err != nil {
		return err
	}

	err = page.Navigate(url)
	if err != nil {
		return fmt.Errorf("failed to navigate to URL: %w", err)
	}

	// Wait for the page to be loaded
	err = page.WaitLoad()
	if err != nil {
		return fmt.Errorf("failed to wait for page load: %w", err)
	}

	return nil
}

// WaitNavigation waits for a navigation event to happen
func (m *Manager) WaitNavigation(tabID string, timeout int) error {
	page, err := m.GetTab(tabID)
	if err != nil {
		return err
	}

	// Set a timeout for the navigation wait
	page = page.Timeout(time.Duration(timeout) * time.Second)

	// Wait for navigation
	page.WaitNavigation(proto.PageLifecycleEventNameLoad)()

	// Wait for the page to be fully loaded
	err = page.WaitLoad()
	if err != nil {
		return fmt.Errorf("navigation wait failed: %w", err)
	}

	return nil
}

// GetPageSource returns the entire source code of a page
func (m *Manager) GetPageSource(tabID string) (string, error) {
	page, err := m.GetTab(tabID)
	if err != nil {
		return "", err
	}

	html, err := page.HTML()
	if err != nil {
		return "", fmt.Errorf("failed to get page HTML: %w", err)
	}

	return html, nil
}

// GetElementHTML returns the HTML of an element at the specified selector
func (m *Manager) GetElementHTML(tabID, selector string) (string, error) {
	page, err := m.GetTab(tabID)
	if err != nil {
		return "", err
	}

	// Find the element
	element, err := page.Element(selector)
	if err != nil {
		return "", fmt.Errorf("failed to find element: %w", err)
	}

	// Get the HTML of the element
	html, err := element.HTML()
	if err != nil {
		return "", fmt.Errorf("failed to get element HTML: %w", err)
	}

	return html, nil
}