pm/scripts/gemini_review/main.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"strings"

	"code.gitea.io/sdk/gitea"
	"github.com/google/generative-ai-go/genai"
	"google.golang.org/api/option"
)

type Suggestion struct {
	LineNumber uint64 `json:"line_number"`
	Comment    string `json:"comment"`
	Snippet    string `json:"snippet,omitempty"`
}

func main() {
	apiKey := os.Getenv("GEMINI_API_KEY")
	token := os.Getenv("GITEA_TOKEN")
	baseURL := os.Getenv("GITEA_URL")
	repoFullName := os.Getenv("GITHUB_REPOSITORY")
	prNumberStr := os.Getenv("PR_NUMBER")

	if apiKey == "" || token == "" || repoFullName == "" || prNumberStr == "" {
		log.Fatal("Missing required environment variables: GEMINI_API_KEY, GITEA_TOKEN, GITHUB_REPOSITORY, PR_NUMBER")
	}

	if baseURL == "" {
		baseURL = "https://gitea.com"
	}

	prNumber, err := strconv.ParseInt(prNumberStr, 10, 64)
	if err != nil {
		log.Fatalf("Invalid PR_NUMBER: %v", err)
	}

	repoParts := strings.Split(repoFullName, "/")
	if len(repoParts) != 2 {
		log.Fatalf("Invalid GITHUB_REPOSITORY format: %s", repoFullName)
	}
	owner, repo := repoParts[0], repoParts[1]

	ctx := context.Background()

	// Initialize Gitea Client
	client, err := gitea.NewClient(baseURL, gitea.SetToken(token))
	if err != nil {
		log.Fatalf("Failed to create Gitea client: %v", err)
	}

	// Initialize Gemini Client
	geminiClient, err := genai.NewClient(ctx, option.WithAPIKey(apiKey))
	if err != nil {
		log.Fatalf("Failed to create Gemini client: %v", err)
	}
	defer geminiClient.Close()

	model := geminiClient.GenerativeModel("gemini-1.5-flash")

	// Get PR files
	files, _, err := client.ListPullRequestFiles(owner, repo, prNumber, gitea.ListPullRequestFilesOptions{})
	if err != nil {
		log.Fatalf("Failed to get PR files: %v", err)
	}

	var giteaComments []gitea.CreatePullReviewComment
	for _, file := range files {
		if !strings.HasSuffix(file.Filename, ".md") {
			continue
		}

		fmt.Printf("Reviewing file: %s\n", file.Filename)
		content, err := readFile(file.Filename)
		if err != nil {
			fmt.Printf("Error reading file %s: %v\n", file.Filename, err)
			continue
		}

		suggestions, err := getGeminiReview(ctx, model, content)
		if err != nil {
			fmt.Printf("Error getting review for %s: %v\n", file.Filename, err)
			continue
		}

		for _, s := range suggestions {
			body := s.Comment
			if s.Snippet != "" {
				body += fmt.Sprintf("\n\n```suggestion\n%s\n```", s.Snippet)
			}
			giteaComments = append(giteaComments, gitea.CreatePullReviewComment{
				Path:       file.Filename,
				Body:       body,
				NewLineNum: int64(s.LineNumber),
			})
		}
	}

	if len(giteaComments) > 0 {
		_, _, err = client.CreatePullReview(owner, repo, prNumber, gitea.CreatePullReviewOptions{
			State:    gitea.ReviewStateRequestChanges,
			Body:     "### 🤖 Gemini Writing Review\n\nI've found some areas for improvement in the documentation. Please see the inline comments below.",
			Comments: giteaComments,
		})
		if err != nil {
			log.Fatalf("Failed to create PR review: %v", err)
		}
		fmt.Printf("Successfully created PR review with %d inline comments.\n", len(giteaComments))
	} else {
		fmt.Println("No Markdown files to review or no suggestions found.")
	}
}

func readFile(path string) (string, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer f.Close()

	content, err := io.ReadAll(f)
	if err != nil {
		return "", err
	}
	return string(content), nil
}

func getGeminiReview(ctx context.Context, model *genai.GenerativeModel, content string) ([]Suggestion, error) {
	prompt := fmt.Sprintf(`
Review the following Markdown content for spelling errors, grammar mistakes, and style improvements.
Analyze the text line by line.

For each issue found, provide a suggestion in JSON format with:
- "line_number": The 1-indexed line number where the issue occurs.
- "comment": A brief explanation of the problem and the suggested fix.
- "snippet": The corrected text for that line (optional, but highly recommended for spell/grammar fixes).

Return ONLY a JSON array of suggestions. If no issues are found, return "[]".

Content with line numbers for reference:
%s
`, addLineNumbers(content))

	resp, err := model.GenerateContent(ctx, genai.Text(prompt))
	if err != nil {
		return nil, err
	}

	if len(resp.Candidates) == 0 || len(resp.Candidates[0].Content.Parts) == 0 {
		return nil, nil
	}

	var rawJS strings.Builder
	for _, part := range resp.Candidates[0].Content.Parts {
		if text, ok := part.(genai.Text); ok {
			rawJS.WriteString(string(text))
		}
	}

	// Clean up markdown code blocks if present
	js := strings.TrimSpace(rawJS.String())
	js = strings.TrimPrefix(js, "```json")
	js = strings.TrimSuffix(js, "```")
	js = strings.TrimSpace(js)

	var suggestions []Suggestion
	if err := json.Unmarshal([]byte(js), &suggestions); err != nil {
		return nil, fmt.Errorf("failed to unmarshal suggestions: %v (raw response: %s)", err, js)
	}

	return suggestions, nil
}

func addLineNumbers(text string) string {
	lines := strings.Split(text, "\n")
	for i, line := range lines {
		lines[i] = fmt.Sprintf("%d: %s", i+1, line)
	}
	return strings.Join(lines, "\n")
}