categoryNode.ts

/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *  Licensed under the MIT License. See License.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/

import * as vscode from 'vscode';
import { RemoteInfo } from '../../../common/types';
import { AuthenticationError, AuthProvider } from '../../common/authentication';
import { DEV_MODE, PR_SETTINGS_NAMESPACE } from '../../common/settingKeys';
import { ITelemetry } from '../../common/telemetry';
import { toQueryUri } from '../../common/uri';
import { formatError } from '../../common/utils';
import { isCopilotQuery } from '../../github/copilotPrWatcher';
import { FolderRepositoryManager, ItemsResponseResult } from '../../github/folderRepositoryManager';
import { PRType } from '../../github/interface';
import { PullRequestModel } from '../../github/pullRequestModel';
import { extractRepoFromQuery } from '../../github/utils';
import { PrsTreeModel } from '../prsTreeModel';
import { PRNode } from './pullRequestNode';
import { TreeNode, TreeNodeParent } from './treeNode';
import { IQueryInfo } from './workspaceFolderNode';
import { NotificationsManager } from '../../notifications/notificationsManager';

export enum PRCategoryActionType {
	Empty,
	More,
	TryOtherRemotes,
	Login,
	LoginEnterprise,
	NoRemotes,
	NoMatchingRemotes,
	ConfigureRemotes,
}

export class PRCategoryActionNode extends TreeNode implements vscode.TreeItem {
	public collapsibleState: vscode.TreeItemCollapsibleState;
	public iconPath?: { light: vscode.Uri; dark: vscode.Uri };
	public type: PRCategoryActionType;
	public command?: vscode.Command;

	constructor(parent: TreeNodeParent, type: PRCategoryActionType, node?: CategoryTreeNode) {
		super(parent);
		this.type = type;
		this.collapsibleState = vscode.TreeItemCollapsibleState.None;
		switch (type) {
			case PRCategoryActionType.Empty:
				this.label = vscode.l10n.t('0 pull requests in this category');
				break;
			case PRCategoryActionType.More:
				this.label = vscode.l10n.t('Load more');
				this.command = {
					title: vscode.l10n.t('Load more'),
					command: 'pr.loadMore',
					arguments: [node],
				};
				break;
			case PRCategoryActionType.TryOtherRemotes:
				this.label = vscode.l10n.t('Continue fetching from other remotes');
				this.command = {
					title: vscode.l10n.t('Load more'),
					command: 'pr.loadMore',
					arguments: [node],
				};
				break;
			case PRCategoryActionType.Login:
				this.label = vscode.l10n.t('Sign in');
				this.command = {
					title: vscode.l10n.t('Sign in'),
					command: 'pr.signinAndRefreshList',
					arguments: [],
				};
				break;
			case PRCategoryActionType.LoginEnterprise:
				this.label = vscode.l10n.t('Sign in with GitHub Enterprise...');
				this.command = {
					title: 'Sign in',
					command: 'pr.signinAndRefreshList',
					arguments: [],
				};
				break;
			case PRCategoryActionType.NoRemotes:
				this.label = vscode.l10n.t('No GitHub repositories found.');
				break;
			case PRCategoryActionType.NoMatchingRemotes:
				this.label = vscode.l10n.t('No remotes match the current setting.');
				break;
			case PRCategoryActionType.ConfigureRemotes:
				this.label = vscode.l10n.t('Configure remotes...');
				this.command = {
					title: vscode.l10n.t('Configure remotes'),
					command: 'pr.configureRemotes',
					arguments: [],
				};
				break;
			default:
				break;
		}
	}

	getTreeItem(): vscode.TreeItem {
		return this;
	}
}

interface PageInformation {
	pullRequestPage: number;
	hasMorePages: boolean;
}

export namespace DefaultQueries {
	export namespace Queries {
		export const LOCAL = 'Local Pull Request Branches';
		export const ALL = 'All Open';
	}
	export namespace Values {
		export const DEFAULT = 'default';
	}
}

export function isLocalQuery(queryInfo: IQueryInfo): boolean {
	return queryInfo.label === DefaultQueries.Queries.LOCAL && queryInfo.query === DefaultQueries.Values.DEFAULT;
}

export function isAllQuery(queryInfo: IQueryInfo): boolean {
	return queryInfo.label === DefaultQueries.Queries.ALL && queryInfo.query === DefaultQueries.Values.DEFAULT;
}

export class CategoryTreeNode extends TreeNode implements vscode.TreeItem {
	public collapsibleState: vscode.TreeItemCollapsibleState;
	public prs: Map<number, PullRequestModel>;
	public fetchNextPage: boolean = false;
	public repositoryPageInformation: Map<string, PageInformation> = new Map<string, PageInformation>();
	public contextValue: string;
	public resourceUri: vscode.Uri;
	public tooltip?: string | vscode.MarkdownString | undefined;
	readonly isCopilot: boolean;
	private _repo: RemoteInfo | undefined;

	constructor(
		parent: TreeNodeParent,
		readonly folderRepoManager: FolderRepositoryManager,
		private _telemetry: ITelemetry,
		public readonly type: PRType,
		private _notificationProvider: NotificationsManager,
		private _prsTreeModel: PrsTreeModel,
		_categoryLabel?: string,
		private _categoryQuery?: string,
	) {
		super(parent);

		this.prs = new Map();
		this.isCopilot = !!_categoryQuery && isCopilotQuery(_categoryQuery);

		switch (this.type) {
			case PRType.All:
				this.label = vscode.l10n.t('All Open');
				break;
			case PRType.Query:
				this.label = _categoryLabel!;
				break;
			case PRType.LocalPullRequest:
				this.label = vscode.l10n.t('Local Pull Request Branches');
				break;
			default:
				this.label = '';
				break;
		}

		this.id = parent instanceof TreeNode ? `${parent.id ?? parent.label}/${this.label}` : this.label;

		// Check if dev mode is enabled to collapse all queries
		const devMode = vscode.workspace.getConfiguration(PR_SETTINGS_NAMESPACE).get<boolean>(DEV_MODE, false);

		if (devMode) {
			this.collapsibleState = vscode.TreeItemCollapsibleState.Collapsed;
		} else {
			if ((this._prsTreeModel.expandedQueries === undefined) && (this.type === PRType.All)) {
				this.collapsibleState = vscode.TreeItemCollapsibleState.Expanded;
			} else {
				this.collapsibleState =
					this._prsTreeModel.expandedQueries?.has(this.id)
						? vscode.TreeItemCollapsibleState.Expanded
						: vscode.TreeItemCollapsibleState.Collapsed;
			}
		}

		if (this._categoryQuery) {
			this.contextValue = this.isCopilot ? 'copilot-query' : 'query';
		}

		if (this.isCopilot) {
			this.tooltip = vscode.l10n.t('Pull requests you asked the coding agent to create');
		} else if (this.type === PRType.LocalPullRequest) {
			this.tooltip = vscode.l10n.t('Pull requests for branches you have locally');
		} else if (this.type === PRType.All) {
			this.tooltip = vscode.l10n.t('All open pull requests in the current repository');
		} else if (_categoryQuery) {
			this.tooltip = new vscode.MarkdownString(vscode.l10n.t('Pull requests for query: `{0}`', _categoryQuery));
		} else {
			this.tooltip = this.label;
		}
	}

	get repo(): RemoteInfo | undefined {
		return this._repo;
	}

	private _getDescription(): string | undefined {
		if (!this.isCopilot || !this._repo) {
			return undefined;
		}
		const counts = this._prsTreeModel.getCopilotCounts(this._repo.owner, this._repo.repositoryName);
		if (counts.total === 0) {
			return undefined;
		} else if (counts.error > 0) {
			if (counts.inProgress > 0) {
				return vscode.l10n.t('{0} in progress, {1} with errors', counts.inProgress, counts.error);
			} else {
				return vscode.l10n.t('{0} with errors', counts.error);
			}
		} else if (counts.inProgress > 0) {
			return vscode.l10n.t('{0} in progress', counts.inProgress);
		} else {
			return vscode.l10n.t('done working on {0}', counts.total);
		}
	}

	public async expandPullRequest(pullRequest: PullRequestModel, retry: boolean = true): Promise<boolean> {
		if (!this._children && retry) {
			await this.getChildren();
			retry = false;
		}
		if (this._children) {
			for (const child of this._children) {
				if (child instanceof PRNode) {
					if (child.pullRequestModel.equals(pullRequest)) {
						this.reveal(child, { expand: true, select: true });
						return true;
					}
				}
			}
			// If we didn't find the PR, we might need to re-run the query
			if (retry) {
				await this.getChildren();
				return await this.expandPullRequest(pullRequest, false);
			}
		}
		return false;
	}

	override async getChildren(shouldDispose: boolean = true): Promise<TreeNode[]> {
		await super.getChildren(shouldDispose);
		if (!shouldDispose && this._children) {
			return this._children;
		}
		const isFirstLoad = !this._firstLoad;
		if (isFirstLoad) {
			this._firstLoad = this.doGetChildren();
			if (!this._prsTreeModel.hasLoaded) {
				this._firstLoad.then(() => this.refresh(this));
				return [];
			}
		}
		return isFirstLoad ? this._firstLoad! : this.doGetChildren();
	}

	private _firstLoad: Promise<TreeNode[]> | undefined;
	private async doGetChildren(): Promise<TreeNode[]> {
		let hasMorePages = false;
		let hasUnsearchedRepositories = false;
		let needLogin = false;
		const fetchNextPage = this.fetchNextPage;
		this.fetchNextPage = false;
		if (this.type === PRType.LocalPullRequest) {
			try {
				this.prs.clear();
				(await this._prsTreeModel.getLocalPullRequests(this.folderRepoManager)).items.forEach(item => this.prs.set(item.id, item));
			} catch (e) {
				vscode.window.showErrorMessage(vscode.l10n.t('Fetching local pull requests failed: {0}', formatError(e)));
				needLogin = e instanceof AuthenticationError;
			}
		} else {
			try {
				let response: ItemsResponseResult<PullRequestModel>;
				switch (this.type) {
					case PRType.All:
						response = await this._prsTreeModel.getAllPullRequests(this.folderRepoManager, fetchNextPage);
						break;
					case PRType.Query:
						response = await this._prsTreeModel.getPullRequestsForQuery(this.folderRepoManager, fetchNextPage, this._categoryQuery!);
						break;
				}
				if (!fetchNextPage) {
					this.prs.clear();
				}
				response.items.forEach(item => this.prs.set(item.id, item));
				hasMorePages = response.hasMorePages;
				hasUnsearchedRepositories = response.hasUnsearchedRepositories;
			} catch (e) {
				if (this.isCopilot && (e.response.status === 422) && e.message.includes('the users do not exist')) {
					// Skip it, it's copilot and the repo doesn't have copilot
				} else if (e.status === 404 || e.response?.status === 404) {
					// 404 errors might indicate wrong account - this is handled in folderRepositoryManager
					// but we catch it here to prevent duplicate error messages
					needLogin = e instanceof AuthenticationError;
				} else {
					const error = formatError(e);
					const actions: string[] = [];
					if (error.includes('Bad credentials')) {
						actions.push(vscode.l10n.t('Login again'));
					} else if (e.status === 403 || e.response?.status === 403) {
						// 403 forbidden - user might not have access with current account
						// Check both GitHub.com and Enterprise providers since we might have repos from either
						const isAuthenticatedGitHub = await this.folderRepoManager.credentialStore.isAuthenticatedForAccountPreferences(AuthProvider.github);
						const isAuthenticatedEnterprise = await this.folderRepoManager.credentialStore.isAuthenticatedForAccountPreferences(AuthProvider.githubEnterprise);
						if (isAuthenticatedGitHub || isAuthenticatedEnterprise) {
							actions.push(vscode.l10n.t('Check Account Preferences'));
						}
					}
					vscode.window.showErrorMessage(vscode.l10n.t('Fetching pull requests failed: {0}', formatError(e)), ...actions).then(action => {
						if (action === vscode.l10n.t('Login again')) {
							this.folderRepoManager.credentialStore.recreate(vscode.l10n.t('Your login session is no longer valid.'));
						} else if (action === vscode.l10n.t('Check Account Preferences')) {
							vscode.commands.executeCommand('_account.manageAccountPreferences', 'GitHub.vscode-pull-request-github');
						}
					});
					needLogin = e instanceof AuthenticationError;
				}
			}
		}

		if (this.prs.size > 0) {
			const nodes: (PRNode | PRCategoryActionNode)[] = Array.from(this.prs.values()).map(
				prItem => new PRNode(this, this.folderRepoManager, prItem, this.type === PRType.LocalPullRequest, this._notificationProvider, this._prsTreeModel),
			);
			if (hasMorePages) {
				nodes.push(new PRCategoryActionNode(this, PRCategoryActionType.More, this));
			} else if (hasUnsearchedRepositories) {
				nodes.push(new PRCategoryActionNode(this, PRCategoryActionType.TryOtherRemotes, this));
			}

			this._children = nodes;
			return nodes;
		} else {
			const category = needLogin ? PRCategoryActionType.Login : PRCategoryActionType.Empty;
			const result = [new PRCategoryActionNode(this, category)];

			this._children = result;
			return result;
		}
	}

	async getTreeItem(): Promise<vscode.TreeItem> {
		this.description = this._getDescription();
		if (!this._repo) {
			this._repo = await extractRepoFromQuery(this.folderRepoManager, this._categoryQuery);
		}
		this.resourceUri = toQueryUri({ remote: this._repo, isCopilot: this.isCopilot });

		// Update contextValue based on current notification state
		if (this._categoryQuery) {
			const hasNotifications = this.isCopilot && this._repo && this._prsTreeModel.getCopilotNotificationsCount(this._repo.owner, this._repo.repositoryName) > 0;
			this.contextValue = this.isCopilot ?
				(hasNotifications ? 'copilot-query-with-notifications' : 'copilot-query') :
				'query';
		}

		return this;
	}
}