lazygit/cmd/copilot/main.go

package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/sanity-io/litter"
)

// ICopilotChat defines the interface for chat operations
type ICopilotChat interface {
	Authenticate() error
	IsAuthenticated() bool
	Chat(request Request) (string, error)
}

var _ ICopilotChat = &CopilotChat{}

type Role string

const (
	RoleUser      Role = "user"
	RoleAssistant Role = "assistant"
	RoleSystem    Role = "system"
)

type Model string

const (
	Gpt4o           Model = "gpt-4o-2024-05-13"
	Gpt4            Model = "gpt-4"
	Gpt3_5Turbo     Model = "gpt-3.5-turbo"
	O1Preview       Model = "o1-preview-2024-09-12"
	O1Mini          Model = "o1-mini-2024-09-12"
	Claude3_5Sonnet Model = "claude-3.5-sonnet"
)

const (
	COPILOT_CHAT_COMPLETION_URL = "https://api.githubcopilot.com/chat/completions"
	COPILOT_CHAT_AUTH_URL       = "https://api.github.com/copilot_internal/v2/token"
	EDITOR_VERSION              = "Lazygit/0.44.0"
	COPILOT_INTEGRATION_ID      = "vscode-chat"
)
const (
	CACHE_FILE_NAME = ".copilot_auth.json"
)
const (
	CHECK_INTERVAL = 30 * time.Second
	MAX_AUTH_TIME  = 5 * time.Minute
)
const (
	GITHUB_CLIENT_ID = "Iv1.b507a08c87ecfe98"
)

type ChatMessage struct {
	Role    Role   `json:"role"`
	Content string `json:"content"`
}

type Request struct {
	Intent      bool          `json:"intent"`
	N           int           `json:"n"`
	Stream      bool          `json:"stream"`
	Temperature float32       `json:"temperature"`
	Model       Model         `json:"model"`
	Messages    []ChatMessage `json:"messages"`
}

type ApiTokenResponse struct {
	Token     string `json:"token"`
	ExpiresAt int64  `json:"expires_at"`
}

type ApiToken struct {
	ApiKey    string
	ExpiresAt time.Time
}
type CacheData struct {
	OAuthToken string    `json:"oauth_token"`
	ApiKey     string    `json:"api_key"`
	ExpiresAt  time.Time `json:"expires_at"`
}
type DeviceCodeResponse struct {
	DeviceCode      string `json:"device_code"`
	UserCode        string `json:"user_code"`
	VerificationUri string `json:"verification_uri"`
	ExpiresIn       int    `json:"expires_in"`
	Interval        int    `json:"interval"`
}

type DeviceTokenResponse struct {
	AccessToken string `json:"access_token"`
	TokenType   string `json:"token_type"`
	Scope       string `json:"scope"`
	Error       string `json:"error,omitempty"`
}

type CopilotChat struct {
	OAuthToken string
	ApiToken   *ApiToken
	Client     *http.Client
	mu         sync.Mutex
}

func NewCopilotChat(client *http.Client) *CopilotChat {
	if client == nil {
		client = &http.Client{}
	}

	chat := &CopilotChat{
		Client: client,
	}

	if err := chat.loadFromCache(); err != nil {
		log.Printf("Warning: Failed to load from cache: %v", err)
	}

	return chat
}

func getCacheFilePath() (string, error) {
	execPath, err := os.Executable()
	if err != nil {
		return "", fmt.Errorf("failed to get executable path: %v", err)
	}
	return filepath.Join(filepath.Dir(execPath), CACHE_FILE_NAME), nil
}

func (self *CopilotChat) saveToCache() error {
	cacheFilePath, err := getCacheFilePath()
	if err != nil {
		return err
	}

	data := CacheData{
		OAuthToken: self.OAuthToken,
		ApiKey:     self.ApiToken.ApiKey,
		ExpiresAt:  self.ApiToken.ExpiresAt,
	}

	fileData, err := json.Marshal(data)
	if err != nil {
		return fmt.Errorf("failed to marshal cache data: %v", err)
	}

	if err := os.WriteFile(cacheFilePath, fileData, 0600); err != nil {
		return fmt.Errorf("failed to write cache file: %v", err)
	}

	return nil
}

func (self *CopilotChat) loadFromCache() error {
	cacheFilePath, err := getCacheFilePath()
	if err != nil {
		return err
	}

	fileData, err := os.ReadFile(cacheFilePath)
	if err != nil {
		if os.IsNotExist(err) {
			return nil // Cache doesn't exist yet, not an error
		}
		return fmt.Errorf("failed to read cache file: %v", err)
	}

	var data CacheData
	if err := json.Unmarshal(fileData, &data); err != nil {
		return fmt.Errorf("failed to unmarshal cache data: %v", err)
	}

	// Always load OAuth token if it exists
	if data.OAuthToken != "" {
		self.OAuthToken = data.OAuthToken
	}

	// If we have a valid API key, use it
	if data.ApiKey != "" && data.ExpiresAt.After(time.Now()) {
		self.ApiToken = &ApiToken{
			ApiKey:    data.ApiKey,
			ExpiresAt: data.ExpiresAt,
		}
		fmt.Println("Loaded valid API key from cache")
		return nil
	}

	// If we have OAuth token but no valid API key, fetch a new one
	if self.OAuthToken != "" {
		fmt.Println("OAuth token found, fetching new API key...")
		apiTokenReq, err := http.NewRequest(http.MethodGet, COPILOT_CHAT_AUTH_URL, nil)
		if err != nil {
			return fmt.Errorf("failed to create API token request: %v", err)
		}
		apiTokenReq.Header.Set("Authorization", fmt.Sprintf("token %s", self.OAuthToken))
		apiTokenReq.Header.Set("Accept", "application/json")
		apiTokenReq.Header.Set("Editor-Version", EDITOR_VERSION)
		apiTokenReq.Header.Set("Copilot-Integration-Id", COPILOT_INTEGRATION_ID)

		apiTokenResp, err := self.Client.Do(apiTokenReq)
		if err != nil {
			return fmt.Errorf("failed to get API token: %v", err)
		}
		defer apiTokenResp.Body.Close()

		var apiTokenResponse ApiTokenResponse
		if err := json.NewDecoder(apiTokenResp.Body).Decode(&apiTokenResponse); err != nil {
			return fmt.Errorf("failed to decode API token response: %v", err)
		}

		self.ApiToken = &ApiToken{
			ApiKey:    apiTokenResponse.Token,
			ExpiresAt: time.Unix(apiTokenResponse.ExpiresAt, 0),
		}

		// Save the new API key to cache
		if err := self.saveToCache(); err != nil {
			log.Printf("Warning: Failed to save new API key to cache: %v", err)
		}

		fmt.Println("Successfully fetched new API key")

		return nil
	}

	return nil
}

func (self *CopilotChat) Authenticate() error {
	// Try to load from cache first
	if err := self.loadFromCache(); err == nil {
		return nil
	}

	self.mu.Lock()
	defer self.mu.Unlock()

	// Step 1: Request device and user codes
	deviceCodeReq, err := http.NewRequest(
		http.MethodPost,
		"https://github.com/login/device/code",
		strings.NewReader(fmt.Sprintf(
			"client_id=%s&scope=copilot",
			GITHUB_CLIENT_ID,
		)),
	)
	if err != nil {
		return fmt.Errorf("failed to create device code request: %v", err)
	}
	deviceCodeReq.Header.Set("Accept", "application/json")
	deviceCodeReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := self.Client.Do(deviceCodeReq)
	if err != nil {
		return fmt.Errorf("failed to get device code: %v", err)
	}
	defer resp.Body.Close()

	var deviceCode DeviceCodeResponse
	if err := json.NewDecoder(resp.Body).Decode(&deviceCode); err != nil {
		return fmt.Errorf("failed to decode device code response: %v", err)
	}

	// Step 2: Display user code and verification URL
	fmt.Printf("\nPlease visit: %s\n", deviceCode.VerificationUri)
	fmt.Printf("And enter code: %s\n\n", deviceCode.UserCode)

	// Step 3: Poll for the access token with timeout
	startTime := time.Now()
	attempts := 0

	for {
		if time.Since(startTime) >= MAX_AUTH_TIME {
			return fmt.Errorf("authentication timed out after 5 minutes")
		}

		time.Sleep(CHECK_INTERVAL)
		attempts++
		fmt.Printf("Checking for authentication... attempt %d\n", attempts)

		tokenReq, err := http.NewRequest(http.MethodPost, "https://github.com/login/oauth/access_token",
			strings.NewReader(fmt.Sprintf(
				"client_id=%s&device_code=%s&grant_type=urn:ietf:params:oauth:grant-type:device_code", GITHUB_CLIENT_ID,
				deviceCode.DeviceCode)))
		if err != nil {
			return fmt.Errorf("failed to create token request: %v", err)
		}
		tokenReq.Header.Set("Accept", "application/json")
		tokenReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")

		tokenResp, err := self.Client.Do(tokenReq)
		if err != nil {
			return fmt.Errorf("failed to get access token: %v", err)
		}

		var tokenResponse DeviceTokenResponse
		if err := json.NewDecoder(tokenResp.Body).Decode(&tokenResponse); err != nil {
			tokenResp.Body.Close()
			return fmt.Errorf("failed to decode token response: %v", err)
		}
		tokenResp.Body.Close()

		if tokenResponse.Error == "authorization_pending" {
			fmt.Println("Login not detected. Please visit the URL and enter the code.")
			continue
		}
		if tokenResponse.Error != "" {
			if time.Since(startTime) >= MAX_AUTH_TIME {
				return fmt.Errorf("authentication timed out after 5 minutes")
			}
			continue
		}

		// Successfully got the access token
		self.OAuthToken = tokenResponse.AccessToken

		// Now get the Copilot API token
		apiTokenReq, err := http.NewRequest(http.MethodGet, COPILOT_CHAT_AUTH_URL, nil)
		if err != nil {
			return fmt.Errorf("failed to create API token request: %v", err)
		}
		apiTokenReq.Header.Set("Authorization", fmt.Sprintf("token %s", self.OAuthToken))
		apiTokenReq.Header.Set("Accept", "application/json")
		apiTokenReq.Header.Set("Editor-Version", EDITOR_VERSION)
		apiTokenReq.Header.Set("Copilot-Integration-Id", COPILOT_INTEGRATION_ID)

		apiTokenResp, err := self.Client.Do(apiTokenReq)
		if err != nil {
			return fmt.Errorf("failed to get API token: %v", err)
		}
		defer apiTokenResp.Body.Close()

		var apiTokenResponse ApiTokenResponse

		var buf bytes.Buffer
		tee := io.TeeReader(apiTokenResp.Body, &buf)

		// Debug response
		respBody, _ := io.ReadAll(tee)
		fmt.Printf("API Token Response: %s\n", string(respBody))

		// Decode using the buffer
		if err := json.Unmarshal(buf.Bytes(), &apiTokenResponse); err != nil {
			return fmt.Errorf("failed to decode API token response: %v", err)
		}

		// Re-create reader for JSON decoding
		apiTokenResp.Body = io.NopCloser(bytes.NewBuffer(respBody))

		self.ApiToken = &ApiToken{
			ApiKey:    apiTokenResponse.Token,
			ExpiresAt: time.Unix(apiTokenResponse.ExpiresAt, 0),
		}

		fmt.Println("Successfully authenticated!")
		// Save the new credentials to cache
		if err := self.saveToCache(); err != nil {
			log.Printf("Warning: Failed to save credentials to cache: %v", err)
		}
		return nil
	}
}

func (self *CopilotChat) IsAuthenticated() bool {
	if self.ApiToken == nil {
		return false
	}
	return self.ApiToken.ExpiresAt.After(time.Now())
}

func (self *CopilotChat) Chat(request Request) (string, error) {
	fmt.Println("Chatting with Copilot...")

	if !self.IsAuthenticated() {
		fmt.Println("Not authenticated with Copilot. Authenticating...")
		if err := self.Authenticate(); err != nil {
			return "", fmt.Errorf("authentication failed: %v", err)
		}
	}

	apiKey := self.ApiToken.ApiKey
	fmt.Println("Authenticated with Copilot!")
	fmt.Println("API Key: ", apiKey)

	litter.Dump(self)

	requestBody, err := json.Marshal(request)
	if err != nil {
		return "", err
	}
	fmt.Println("Mounting request body: ", string(requestBody))

	self.mu.Lock()
	defer self.mu.Unlock()

	req, err := http.NewRequest(http.MethodPost, COPILOT_CHAT_COMPLETION_URL, strings.NewReader(string(requestBody)))
	if err != nil {
		return "", err
	}

	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Editor-Version", EDITOR_VERSION)
	req.Header.Set("Copilot-Integration-Id", COPILOT_INTEGRATION_ID)

	response, err := self.Client.Do(req)
	if err != nil {
		return "", err
	}
	defer response.Body.Close()

	if response.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(response.Body)
		return "", fmt.Errorf("failed to get completion: %s", string(body))
	}

	responseBody, err := io.ReadAll(response.Body)
	if err != nil {
		return "", err
	}

	return string(responseBody), nil
}

func main() {
	client := &http.Client{}
	fmt.Println("Starting...")
	copilotChat := NewCopilotChat(client)

	fmt.Println("Chatting...")

	err := copilotChat.Authenticate()
	if err != nil {
		if strings.Contains(err.Error(), "timed out") {
			log.Fatalf("Authentication process timed out. Please try again later.")
		}
		log.Fatalf("Error during authentication: %v", err)
	}

	fmt.Println("Authenticated!")

	messages := []ChatMessage{
		{
			Role:    RoleUser,
			Content: "Describe what is Lazygit in one sentence",
		},
	}

	request := Request{
		Intent:      true,
		N:           1,
		Stream:      false,
		Temperature: 0.1,
		Model:       Gpt4o,
		Messages:    messages,
	}

	response, err := copilotChat.Chat(request)
	if err != nil {
		log.Fatalf("Error during chat: %v", err)
	}

	fmt.Println(response)
}