package main

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strings"
)

var preferredRemotes = []string{"origin", "github", "work"}

type GitRepository struct {
	WorkingDirectory string
	RemoteName       string
	RemoteURL        string
	PrimaryBranch    string
}

func NewGitRepository(workingDir string) (*GitRepository, error) {
	repo := &GitRepository{
		WorkingDirectory: workingDir,
	}

	if err := repo.isGitRepository(); err != nil {
		return nil, err
	}

	if err := repo.workdirIsClean(); err != nil {
		return nil, err
	}

	if err := repo.guessRemote(); err != nil {
		return nil, err
	}

	if err := repo.guessPrimaryBranch(); err != nil {
		return nil, err
	}

	return repo, nil
}

func (r *GitRepository) runGitCommand(args ...string) ([]byte, error) {
	cmd := exec.Command("git", append([]string{"-C", r.WorkingDirectory}, args...)...)
	return cmd.Output()
}

func (r *GitRepository) guessPrimaryBranch() error {
	output, err := r.runGitCommand("ls-remote", "--symref", r.RemoteName, "HEAD")
	if err != nil {
		return fmt.Errorf("failed to run git ls-remote: %w", err)
	}

	re := regexp.MustCompile(`ref: refs/heads/(\S+)\s+HEAD`)
	matches := re.FindStringSubmatch(string(output))

	if len(matches) < 2 {
		return fmt.Errorf("could not determine primary branch")
	}

	r.PrimaryBranch = matches[1]
	return nil
}

func (r *GitRepository) guessRemote() error {
	output, err := r.runGitCommand("config", "--get-regexp", `remote\.[a-z0-9]+\.url`)
	if err != nil {
		return fmt.Errorf("failed to run git config: %w", err)
	}

	remotes := make(map[string]string)
	scanner := bufio.NewScanner(strings.NewReader(string(output)))

	for scanner.Scan() {
		line := scanner.Text()
		parts := strings.Fields(line)
		if len(parts) < 2 {
			continue
		}

		configParts := strings.Split(parts[0], ".")
		if len(configParts) < 2 {
			continue
		}

		remoteName := configParts[1]
		remoteURL := parts[1]
		remotes[remoteName] = remoteURL
	}

	for _, preferred := range preferredRemotes {
		if url, exists := remotes[preferred]; exists {
			r.RemoteName = preferred
			r.RemoteURL = url
			return nil
		}
	}

	return fmt.Errorf("no preferred remote configured")
}

func (r *GitRepository) isGitRepository() error {
	_, err := r.runGitCommand("rev-parse", "--show-toplevel")
	if err != nil {
		return fmt.Errorf("%s is not a git repository", r.WorkingDirectory)
	}
	return nil
}

func (r *GitRepository) workdirIsClean() error {
	_, err := r.runGitCommand("diff", "--quiet", "--exit-code")
	if err != nil {
		return fmt.Errorf("%s has unstaged changes, commit or stash them first", r.WorkingDirectory)
	}
	return nil
}

// getMergedBranches returns a list of local branches that have been merged into the primary branch
func (r *GitRepository) getMergedBranches() ([]string, error) {
	output, err := r.runGitCommand("branch", "--merged", r.PrimaryBranch)
	if err != nil {
		return nil, fmt.Errorf("failed to get merged branches: %w", err)
	}

	var branches []string
	scanner := bufio.NewScanner(strings.NewReader(string(output)))

	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		// Skip current branch indicator and the primary branch itself
		if strings.HasPrefix(line, "*") {
			line = strings.TrimSpace(line[1:])
		}

		// Don't delete the primary branch
		if line != r.PrimaryBranch && line != "" {
			branches = append(branches, line)
		}
	}

	return branches, scanner.Err()
}

// deleteBranch deletes a local branch
func (r *GitRepository) deleteBranch(branch string) error {
	_, err := r.runGitCommand("branch", "-d", branch)
	if err != nil {
		// Try force delete if normal delete fails
		_, forceErr := r.runGitCommand("branch", "-D", branch)
		if forceErr != nil {
			return fmt.Errorf("failed to delete branch %s: %w", branch, err)
		}
		log.Printf("Force deleted branch: %s", branch)
		return nil
	}
	log.Printf("Deleted branch: %s", branch)
	return nil
}

// cleanupMergedBranches deletes all local branches that have been merged into the primary branch
func (r *GitRepository) cleanupMergedBranches(dryRun bool) error {
	mergedBranches, err := r.getMergedBranches()
	if err != nil {
		return err
	}

	if len(mergedBranches) == 0 {
		log.Println("No merged branches to clean up")
		return nil
	}

	log.Printf("Found %d merged branches:", len(mergedBranches))
	for _, branch := range mergedBranches {
		if dryRun {
			log.Printf("Would delete: %s", branch)
		} else {
			if err := r.deleteBranch(branch); err != nil {
				log.Printf("Error deleting branch %s: %v", branch, err)
			}
		}
	}

	return nil
}

func main() {
	wd, err := os.Getwd()
	if err != nil {
		log.Fatalf("Failed to get working directory: %v", err)
	}

	repo, err := NewGitRepository(wd)
	if err != nil {
		log.Fatalf("Error: %v", err)
	}

	log.Printf("Working in: %s", repo.WorkingDirectory)
	log.Printf("Main branch is: %s", repo.PrimaryBranch)
	log.Printf("Remote is named %s and is at %s", repo.RemoteName, repo.RemoteURL)

	dryRun := len(os.Args) > 1 && os.Args[1] == "--dry-run"

	if dryRun {
		log.Println("Running in dry-run mode")
	}

	if err := repo.cleanupMergedBranches(dryRun); err != nil {
		log.Fatalf("Failed to cleanup merged branches: %v", err)
	}
}