opencode-cursor-agent/internal/providers/geminiweb/provider.go

package geminiweb

import (
	"context"
	"cursor-api-proxy/internal/apitypes"
	"cursor-api-proxy/internal/config"
	"fmt"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/go-rod/rod"
)

type Provider struct {
	cfg  config.BridgeConfig
	pool *SessionPool
}

func NewProvider(cfg config.BridgeConfig) *Provider {
	return &Provider{cfg: cfg}
}

func (p *Provider) Name() string {
	return "gemini-web"
}

func (p *Provider) Close() error {
	return nil
}

func (p *Provider) initPool() error {
	if p.pool != nil {
		return nil
	}
	pool, err := NewSessionPool(p.cfg.GeminiAccountDir, p.cfg.GeminiMaxSessions)
	if err != nil {
		return fmt.Errorf("failed to init session pool: %w", err)
	}
	p.pool = pool
	return nil
}

func (p *Provider) Generate(ctx context.Context, model string, messages []apitypes.Message, tools []apitypes.Tool, cb func(apitypes.StreamChunk)) error {
	fmt.Printf("[GeminiWeb] Starting generation with model: %s\n", model)

	if err := p.initPool(); err != nil {
		return err
	}

	session := p.pool.GetAvailable()
	if session == nil {
		return fmt.Errorf("no available sessions - please run 'gemini-login <session-name>' first")
	}

	fmt.Printf("[GeminiWeb] Using session: %s\n", session.Name)

	p.pool.StartSession(session)
	defer p.pool.EndSession(session)

	browser, err := NewBrowser(p.cfg.GeminiBrowserVisible)
	if err != nil {
		return fmt.Errorf("failed to create browser: %w", err)
	}
	defer browser.Close()

	page, err := browser.NewPage()
	if err != nil {
		return fmt.Errorf("failed to create page: %w", err)
	}

	if session.CookieFile != "" {
		fmt.Printf("[GeminiWeb] Loading cookies from: %s\n", session.CookieFile)
		cookies, err := LoadCookiesFromFile(session.CookieFile)
		if err == nil {
			if err := SetCookiesOnPage(page, cookies); err != nil {
				return fmt.Errorf("failed to set cookies: %w", err)
			}
			fmt.Printf("[GeminiWeb] Loaded %d cookies\n", len(cookies))
		} else {
			fmt.Printf("[GeminiWeb] Warning: could not load cookies: %v\n", err)
		}
	}

	fmt.Printf("[GeminiWeb] Navigating to Gemini...\n")
	if err := NavigateToGemini(page); err != nil {
		return fmt.Errorf("failed to navigate: %w", err)
	}

	time.Sleep(2 * time.Second)

	fmt.Printf("[GeminiWeb] Checking login status...\n")
	if !IsLoggedIn(page) {
		return fmt.Errorf("session not logged in - please run 'gemini-login %s' first", session.Name)
	}
	fmt.Printf("[GeminiWeb] Logged in successfully\n")

	fmt.Printf("[GeminiWeb] Selecting model: %s\n", model)
	if err := SelectModel(page, model); err != nil {
		return fmt.Errorf("failed to select model: %w", err)
	}
	fmt.Printf("[GeminiWeb] Model selected\n")

	time.Sleep(500 * time.Millisecond)

	prompt := buildPromptFromMessages(messages)
	fmt.Printf("[GeminiWeb] Sending prompt (length: %d chars)\n", len(prompt))
	if err := SendPrompt(page, prompt); err != nil {
		return fmt.Errorf("failed to send prompt: %w", err)
	}
	fmt.Printf("[GeminiWeb] Prompt sent, waiting for response...\n")

	return WaitForResponse(page,
		func(text string) {
			cb(apitypes.StreamChunk{Type: apitypes.ChunkText, Text: text})
		},
		func(thinking string) {
			cb(apitypes.StreamChunk{Type: apitypes.ChunkThinking, Thinking: thinking})
		},
		func() {
			cb(apitypes.StreamChunk{Type: apitypes.ChunkDone, Done: true})
		},
	)
}

func buildPromptFromMessages(messages []apitypes.Message) string {
	var prompt string
	for _, m := range messages {
		switch m.Role {
		case "system":
			prompt += "System: " + m.Content + "\n\n"
		case "user":
			prompt += m.Content + "\n\n"
		case "assistant":
			prompt += "Assistant: " + m.Content + "\n\n"
		}
	}
	return prompt
}

func RunLogin(cfg config.BridgeConfig, sessionName string) error {
	if sessionName == "" {
		sessionName = fmt.Sprintf("session-%d", time.Now().Unix())
	}

	pool, err := NewSessionPool(cfg.GeminiAccountDir, cfg.GeminiMaxSessions)
	if err != nil {
		return fmt.Errorf("failed to init pool: %w", err)
	}

	session, err := pool.CreateSession(sessionName)
	if err != nil {
		return fmt.Errorf("failed to create session: %w", err)
	}

	fmt.Printf("Starting browser for login. Session: %s\n", sessionName)
	fmt.Println("Please log in to your Gemini account in the browser window.")
	fmt.Println("Press Ctrl+C when you have completed the login...")

	browser, err := NewBrowser(true)
	if err != nil {
		return fmt.Errorf("failed to create browser: %w", err)
	}
	defer browser.Close()

	page, err := browser.NewPage()
	if err != nil {
		return fmt.Errorf("failed to create page: %w", err)
	}

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

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)

	<-sigChan

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

	if err := SaveCookiesToFile(cookies, session.CookieFile); err != nil {
		return fmt.Errorf("failed to save cookies: %w", err)
	}

	fmt.Printf("Session saved successfully: %s\n", sessionName)
	return nil
}

func GetPageCookies(page *rod.Page) ([]Cookie, error) {
	cookies, err := page.Cookies([]string{})
	if err != nil {
		return nil, fmt.Errorf("failed to get cookies: %w", err)
	}

	var result []Cookie
	for _, c := range cookies {
		result = append(result, Cookie{
			Name:     c.Name,
			Value:    c.Value,
			Domain:   c.Domain,
			Path:     c.Path,
			HTTPOnly: c.HTTPOnly,
			Secure:   c.Secure,
		})
	}
	return result, nil
}