package main

import (
	"errors"
	"fmt"
	"io"
	"path/filepath"
	"sort"

	"git.unkin.net/unkin/agent-tools/internal/agent"

	"github.com/spf13/cobra"
)

// Verdicts a worktree can be classified into.
const (
	verdictKeep         = "keep"
	verdictRemove       = "remove"
	verdictRemoveBranch = "remove+branch"
)

// prLister is the slice of the Gitea client prune needs, so tests can drive
// classification without a live server.
type prLister interface {
	ListPRs(repoPath, state string) ([]agent.PullRequest, error)
}

type pruneResult struct {
	wt      managedWt
	verdict string
	reason  string
}

// repoCtx is the per-repo state classification is decided against.
type repoCtx struct {
	srcDir    string
	defBranch string
	prs       map[string]agent.PullRequest
	prsKnown  bool
}

func newPruneCmd() *cobra.Command {
	var apply, keepBranches bool
	cmd := &cobra.Command{
		Use:          "prune",
		Short:        "Classify managed worktrees and remove the ones whose work is safely upstream",
		Long:         "prune inspects every managed worktree, classifies it against git and its Gitea\npull request, and removes the ones whose work is provably upstream. It is a dry\nrun unless --yes is given.",
		SilenceUsage: true,
		RunE: func(cmd *cobra.Command, args []string) error {
			return runPrune(cmd.OutOrStdout(), pruneClient(), apply, keepBranches)
		},
	}
	f := cmd.Flags()
	f.BoolVar(&apply, "yes", false, "Actually remove worktrees (default is a dry run)")
	f.BoolVar(&keepBranches, "keep-branches", false, "Never delete a local branch, whatever the classification")
	return cmd
}

// pruneClient builds a Gitea client, falling back to anonymous access when no
// token can be minted; prune degrades to git-only signals if that fails too.
func pruneClient() prLister {
	tok, err := agent.GiteaToken()
	if err != nil {
		tok = ""
	}
	return agent.NewGiteaClient(tok)
}

func runPrune(out io.Writer, prs prLister, apply, keepBranches bool) error {
	managed, err := managedWorktrees()
	if err != nil {
		return err
	}
	if len(managed) == 0 {
		_, _ = fmt.Fprintln(out, "no managed worktrees")
		return nil
	}

	byRepo := map[string][]managedWt{}
	for _, w := range managed {
		byRepo[w.srcDir] = append(byRepo[w.srcDir], w)
	}
	srcDirs := make([]string, 0, len(byRepo))
	for dir := range byRepo {
		srcDirs = append(srcDirs, dir)
	}
	sort.Strings(srcDirs)

	var results []pruneResult
	for _, srcDir := range srcDirs {
		ctx, err := newRepoCtx(out, prs, srcDir)
		if err != nil {
			for _, w := range byRepo[srcDir] {
				results = append(results, pruneResult{wt: w, verdict: verdictKeep, reason: "repo state unknown: " + err.Error()})
			}
			continue
		}
		for _, w := range byRepo[srcDir] {
			res, err := classify(w, ctx)
			if err != nil {
				res = pruneResult{wt: w, verdict: verdictKeep, reason: "inspection failed: " + err.Error()}
			}
			results = append(results, res)
		}
	}

	for _, r := range results {
		_, _ = fmt.Fprintf(out, "%-44s %-34s %-14s %s\n", filepath.Base(r.wt.path), r.wt.branch, r.verdict, r.reason)
	}
	if !apply {
		_, _ = fmt.Fprintln(out, "dry run: nothing removed (pass --yes to apply)")
		return nil
	}

	var errs []error
	for _, r := range results {
		if r.verdict == verdictKeep {
			continue
		}
		deleteBranch := r.verdict == verdictRemoveBranch && !keepBranches
		if err := removeWorktree(out, r.wt, deleteBranch); err != nil {
			errs = append(errs, fmt.Errorf("%s: %w", r.wt.path, err))
		}
	}
	return errors.Join(errs...)
}

// newRepoCtx refreshes a source repo and collects the signals prune classifies
// against. A failed fetch or an unreachable Gitea is reported and tolerated:
// the git-only signals still work offline.
func newRepoCtx(out io.Writer, prs prLister, srcDir string) (repoCtx, error) {
	ctx := repoCtx{srcDir: srcDir, prs: map[string]agent.PullRequest{}}
	repo := filepath.Base(srcDir)
	if err := agent.GitFetchPrune(srcDir, "origin", credentialHelperArgs()...); err != nil {
		_, _ = fmt.Fprintf(out, "warn: fetch %s: %v (using local refs)\n", repo, err)
	}
	def, err := agent.GitRemoteDefaultBranch(srcDir, "origin")
	if err != nil {
		return repoCtx{}, err
	}
	ctx.defBranch = def

	if prs == nil {
		return ctx, nil
	}
	list, err := prs.ListPRs(repoPath(srcDir, repo), "all")
	if err != nil {
		_, _ = fmt.Fprintf(out, "warn: list PRs for %s: %v (git signals only)\n", repo, err)
		return ctx, nil
	}
	ctx.prs = prsByBranch(list)
	ctx.prsKnown = true
	return ctx, nil
}

// repoPath is the Gitea "owner/repo" for a checkout, read from origin's URL
// because not every managed repo lives under AGENTWS_OWNER.
func repoPath(srcDir, repo string) string {
	url, err := agent.GitRemoteURL(srcDir, "origin")
	if err == nil && agent.RemoteHost(url) == giteaHost() {
		if path, err := agent.RepoPathFromRemoteURL(url); err == nil {
			return path
		}
	}
	return owner() + "/" + repo
}

// prsByBranch indexes PRs by head branch, preferring an open PR and otherwise
// the most recent one when a branch has been used more than once.
func prsByBranch(list []agent.PullRequest) map[string]agent.PullRequest {
	out := map[string]agent.PullRequest{}
	for _, pr := range list {
		branch := agent.PRHeadBranch(pr)
		if branch == "" {
			continue
		}
		if cur, ok := out[branch]; ok && !supersedes(pr, cur) {
			continue
		}
		out[branch] = pr
	}
	return out
}

func supersedes(a, b agent.PullRequest) bool {
	if a.IsOpen() != b.IsOpen() {
		return a.IsOpen()
	}
	if a.Merged != b.Merged {
		return a.Merged
	}
	return a.Number > b.Number
}

// classify applies the prune precedence: dirty and open-PR worktrees are kept,
// provably-upstream work loses its branch too, and anything unproven keeps its
// branch so no commits become unreachable.
func classify(wt managedWt, ctx repoCtx) (pruneResult, error) {
	res := pruneResult{wt: wt}
	dirty, err := agent.GitIsDirty(wt.path)
	if err != nil {
		return res, err
	}
	if dirty {
		res.verdict, res.reason = verdictKeep, "dirty"
		return res, nil
	}

	pr, hasPR := ctx.prs[wt.branch]
	if hasPR && pr.IsOpen() {
		res.verdict, res.reason = verdictKeep, fmt.Sprintf("PR open #%d", pr.Number)
		return res, nil
	}

	upstream := "origin/" + ctx.defBranch
	contained, err := agent.GitIsAncestor(wt.path, "HEAD", upstream)
	if err != nil {
		return res, err
	}
	if contained {
		res.verdict, res.reason = verdictRemoveBranch, "contained in "+upstream
		return res, nil
	}
	unmerged, err := agent.GitUnmergedCommits(wt.path, upstream, "HEAD")
	if err != nil {
		return res, err
	}
	if unmerged == 0 {
		res.verdict, res.reason = verdictRemoveBranch, "cherry-clean against "+upstream
		return res, nil
	}
	if hasPR && pr.Merged {
		res.verdict, res.reason = verdictRemoveBranch, fmt.Sprintf("PR merged #%d", pr.Number)
		return res, nil
	}
	if hasPR {
		if agent.GitRemoteBranchExists(ctx.srcDir, "origin", wt.branch) {
			res.verdict, res.reason = verdictRemoveBranch, fmt.Sprintf("PR closed #%d, branch on origin", pr.Number)
			return res, nil
		}
		res.verdict, res.reason = verdictRemove, fmt.Sprintf("PR closed #%d, branch gone", pr.Number)
		return res, nil
	}
	res.verdict = verdictRemove
	res.reason = "no PR"
	if !ctx.prsKnown {
		res.reason = "PR state unknown"
	}
	return res, nil
}