package cli

import (
	"fmt"
	"os"
	"path/filepath"

	"github.com/spf13/cobra"
	"larc.wejust.rest/larc/internal/core"
	"larc.wejust.rest/larc/internal/repo"
	"larc.wejust.rest/larc/internal/status"
)

func BranchJumpCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:     "branchjump <branch>",
		Short:   "Switch to a different branch",
		Aliases: []string{"bj"},
		Args:    cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			wd, _ := os.Getwd()
			repoRoot, err := repo.FindRepoRoot(wd)
			if err != nil {
				return fmt.Errorf("not a larc repository")
			}

			r, err := repo.Open(repoRoot)
			if err != nil {
				return err
			}
			defer r.Close()

			branches, err := r.Meta.ListBranches()
			if err != nil {
				return err
			}

			if len(branches) == 0 {
				fmt.Println("No branches yet.")
				return nil
			}

			currentBranch, _ := r.CurrentBranch()

			branchName := args[0]
			var selectedBranch *core.Branch
			for _, b := range branches {
				if b.Name == branchName {
					selectedBranch = b
					break
				}
			}

			if selectedBranch == nil {
				return fmt.Errorf("branch '%s' not found", branchName)
			}

			if selectedBranch.Name == currentBranch {
				fmt.Printf("Already on branch %s.\n", selectedBranch.Name)
				return nil
			}

			fmt.Printf("Jumping to branch %s...\n", selectedBranch.Name)

			// Get the head revision of the selected branch
			rev, err := r.Meta.GetRevision(selectedBranch.HeadRev)
			if err != nil {
				return fmt.Errorf("failed to get revision for branch %s: %w", selectedBranch.Name, err)
			}

			// Get the tree for the new revision
			tree, err := r.GetTree(rev.TreeHash)
			if err != nil {
				return fmt.Errorf("failed to get tree for revision r%d: %w", rev.Number, err)
			}

			// Clean the working directory
			// NOTE(kroot): This is a simple implementation that just removes all files.
			// A more robust implementation would compare the files and only
			// remove/update the ones that have changed.

			scanner := status.NewScanner(r)
			entries, err := scanner.GetStagedEntries()
			if err != nil {
				return err
			}
			for _, entry := range entries {
				os.Remove(filepath.Join(repoRoot, entry.Path))
			}

			// Checkout the new tree
			for _, entry := range tree.Entries {
				filePath := filepath.Join(repoRoot, entry.Path)
				if err := os.MkdirAll(filepath.Dir(filePath), 0755); err != nil {
					return fmt.Errorf("failed to create directory %s: %w", filepath.Dir(filePath), err)
				}
				data, err := r.Blobs.Read(entry.BlobHash)
				if err != nil {
					return fmt.Errorf("failed to read blob %s: %w", entry.BlobHash, err)
				}
				if err := os.WriteFile(filePath, data, os.FileMode(entry.Mode)); err != nil {
					return fmt.Errorf("failed to write file %s: %w", filePath, err)
				}
			}

			// Update the current branch and revision
			if err := r.SetCurrentBranch(selectedBranch.Name); err != nil {
				return fmt.Errorf("failed to set current branch: %w", err)
			}
			if err := r.SetCurrentRevision(selectedBranch.HeadRev); err != nil {
				return fmt.Errorf("failed to set current revision: %w", err)
			}

			// Update the index
			scanner.ClearStaging()
			for _, entry := range tree.Entries {
				scanner.Stage(entry.Path, entry.BlobHash, entry.Size)
			}
			if err := scanner.SaveIndex(); err != nil {
				return fmt.Errorf("failed to save index: %w", err)
			}

			fmt.Printf("Successfully jumped to branch %s.\n", selectedBranch.Name)

			return nil
		},
	}
	return cmd
}