src/vulnerabilityTree.ts

import * as vscode from 'vscode';

export interface Vulnerability {
	file: string;
	status: string;
	error?: string;
	result?: string;
	line?: number;
	endLine?: number;
	type?: string;
	severity?: string;
	description?: string;
	originalCode?: string;
	fixedCode?: string;
}

export class VulnerabilityItem extends vscode.TreeItem {
	constructor(
		public readonly label: string,
		public readonly collapsibleState: vscode.TreeItemCollapsibleState,
		public readonly vulnerability?: Vulnerability,
		public readonly command?: vscode.Command,
		contextValue?: string,
		public readonly index?: number
	) {
		super(label, collapsibleState);

		if (contextValue) {
			this.contextValue = contextValue;
		}

		if (vulnerability?.status === 'analyzed' && vulnerability?.type) {
			// Show vulnerability type and severity with fix indicator
			const fixAvailable = vulnerability.fixedCode ? ' ― Apply Fix →' : '';
			this.description = this.getSeverityLabel(vulnerability.severity) + fixAvailable;
			this.iconPath = this.getSeverityIcon(vulnerability.severity);
			this.tooltip = this.buildTooltip(vulnerability);
		} else if (vulnerability?.status === 'analyzed' && vulnerability?.result) {
			this.description = 'Click to apply fix';
			this.iconPath = new vscode.ThemeIcon('error', new vscode.ThemeColor('errorForeground'));
			this.tooltip = vulnerability.result.substring(0, 500);
		} else if (vulnerability?.status === 'error') {
			this.description = 'Analysis error';
			this.iconPath = new vscode.ThemeIcon('warning', new vscode.ThemeColor('warningForeground'));
			this.tooltip = vulnerability.error || 'Unknown error';
		} else {
			this.iconPath = new vscode.ThemeIcon('file-code');
			this.tooltip = 'Click to open file';
		}
	}

	private getSeverityLabel(severity?: string): string {
		switch (severity) {
			case 'critical': return '🔴 Critical';
			case 'high': return '🟠 High';
			case 'medium': return '🟡 Medium';
			case 'low': return '🟢 Low';
			default: return 'Issue';
		}
	}

	private getSeverityIcon(severity?: string): vscode.ThemeIcon {
		switch (severity) {
			case 'critical':
				return new vscode.ThemeIcon('error', new vscode.ThemeColor('errorForeground'));
			case 'high':
				return new vscode.ThemeIcon('warning', new vscode.ThemeColor('problemsWarningIcon.foreground'));
			case 'medium':
				return new vscode.ThemeIcon('info', new vscode.ThemeColor('notificationsInfoIcon.foreground'));
			case 'low':
				return new vscode.ThemeIcon('pass', new vscode.ThemeColor('testing.iconPassed'));
			default:
				return new vscode.ThemeIcon('circle-outline');
		}
	}

	private buildTooltip(vuln: Vulnerability): string {
		const lines = [
			`Type: ${vuln.type || 'Unknown'}`,
			`Severity: ${vuln.severity || 'Unknown'}`,
			`Location: Line ${vuln.line}${vuln.endLine && vuln.endLine !== vuln.line ? `-${vuln.endLine}` : ''}`,
			'',
			'Original Code:',
			vuln.originalCode || '(not available)',
			'',
			'Fixed Code:',
			vuln.fixedCode || '(not available)'
		];
		return lines.join('\n');
	}
}

export class VulnerabilityTreeDataProvider implements vscode.TreeDataProvider<VulnerabilityItem> {
	private _onDidChangeTreeData: vscode.EventEmitter<VulnerabilityItem | undefined | null | void> =
		new vscode.EventEmitter<VulnerabilityItem | undefined | null | void>();
	readonly onDidChangeTreeData: vscode.Event<VulnerabilityItem | undefined | null | void> =
		this._onDidChangeTreeData.event;

	private vulnerabilities: Vulnerability[] = [];

	setVulnerabilities(vulnerabilities: Vulnerability[]): void {
		this.vulnerabilities = vulnerabilities;
		this._onDidChangeTreeData.fire(null);
	}

	getVulnerabilitiesForFile(fileName: string): Vulnerability[] {
		return this.vulnerabilities.filter((v: Vulnerability) => {
			const vulnFileName = v.file.split('\\').pop()?.split('/').pop() || v.file;
			return vulnFileName === fileName;
		});
	}

	getVulnerabilityByIndex(fileName: string, index: number): Vulnerability | undefined {
		const fileVulns = this.getVulnerabilitiesForFile(fileName);
		return fileVulns[index];
	}

	getTreeItem(element: VulnerabilityItem): vscode.TreeItem {
		return element;
	}

	getChildren(element?: VulnerabilityItem): Thenable<VulnerabilityItem[]> {
		if (!element) {
			// Root level - show list of files with vulnerabilities
			const fileMap = new Map<string, Vulnerability[]>();

			this.vulnerabilities.forEach((vuln: Vulnerability) => {
				if (!fileMap.has(vuln.file)) {
					fileMap.set(vuln.file, []);
				}
				fileMap.get(vuln.file)!.push(vuln);
			});

			const items = Array.from(fileMap.entries()).map(
				([file, vulns]) => {
					const errorCount = vulns.filter(v => v.status === 'error').length;
					const issueCount = vulns.filter(v => v.status === 'analyzed').length;
					const description = errorCount > 0
						? `${issueCount} issues, ${errorCount} errors`
						: `${issueCount} issue${issueCount !== 1 ? 's' : ''}`;

					const item = new VulnerabilityItem(
						this.getFileName(file),
						vscode.TreeItemCollapsibleState.Expanded,
						undefined,
						{
							command: 'vscode.open',
							title: 'Open file',
							arguments: [vscode.Uri.file(file)],
						},
						'file'
					);
					item.description = description;
					return item;
				}
			);

			return Promise.resolve(items);
		} else if (element.vulnerability === undefined) {
			// File level - show vulnerabilities in this file
			const fileName = element.label;
			const fileVulns = this.vulnerabilities.filter((v: Vulnerability) =>
				this.getFileName(v.file) === fileName
			);

			const items = fileVulns.map(
				(vuln, index) => {
					// Build a descriptive label
					let label: string;
					if (vuln.type) {
						label = vuln.type;
						if (vuln.line) {
							label += ` (Line ${vuln.line})`;
						}
					} else {
						label = `Issue ${index + 1}${vuln.status === 'error' ? ' (Error)' : ''}`;
					}

					return new VulnerabilityItem(
						label,
						vscode.TreeItemCollapsibleState.Collapsed,
						vuln,
						undefined,
						vuln.status === 'analyzed' ? 'vulnerability' : undefined,
						index
					);
				}
			);

			return Promise.resolve(items);
		} else {
			// Vulnerability level - show details
			const vuln = element.vulnerability;
			const items: VulnerabilityItem[] = [];

			if (vuln.status === 'analyzed') {
				// Show vulnerability details
				if (vuln.description) {
					const descItem = new VulnerabilityItem(
						`📋 ${vuln.description}`,
						vscode.TreeItemCollapsibleState.None,
						undefined
					);
					items.push(descItem);
				}

				// Show original code
				if (vuln.originalCode) {
					const origHeader = new VulnerabilityItem(
						'❌ Vulnerable Code:',
						vscode.TreeItemCollapsibleState.None,
						undefined
					);
					items.push(origHeader);

					const origLines = vuln.originalCode.split('\n').slice(0, 5);
					const origPreview = origLines.map(l => `  ${l}`).join('\n');
					const origItem = new VulnerabilityItem(
						origPreview + (vuln.originalCode.split('\n').length > 5 ? '\n  ...' : ''),
						vscode.TreeItemCollapsibleState.None,
						undefined
					);
					items.push(origItem);
				}

				// Show fixed code
				if (vuln.fixedCode) {
					const fixHeader = new VulnerabilityItem(
						'✅ Fixed Code:',
						vscode.TreeItemCollapsibleState.None,
						undefined
					);
					items.push(fixHeader);

					const fixLines = vuln.fixedCode.split('\n').slice(0, 5);
					const fixPreview = fixLines.map(l => `  ${l}`).join('\n');
					const fixItem = new VulnerabilityItem(
						fixPreview + (vuln.fixedCode.split('\n').length > 5 ? '\n  ...' : ''),
						vscode.TreeItemCollapsibleState.None,
						undefined
					);
					items.push(fixItem);
				} else if (vuln.result) {
					// Fallback to old format
					const lines = vuln.result.split('\n');
					const previewLines = lines.slice(0, 10);
					const preview = previewLines.join('\n') + (lines.length > 10 ? '\n...' : '');

					items.push(
						new VulnerabilityItem(
							'Fixed Code Preview:',
							vscode.TreeItemCollapsibleState.None,
							undefined
						)
					);

					items.push(
						new VulnerabilityItem(
							preview,
							vscode.TreeItemCollapsibleState.None,
							undefined
						)
					);
				}
			} else if (vuln.error) {
				items.push(
					new VulnerabilityItem(
						`Error: ${vuln.error}`,
						vscode.TreeItemCollapsibleState.None,
						undefined
					)
				);
			}

			return Promise.resolve(items);
		}
	}

	private getFileName(filePath: string): string {
		return filePath.split('\\').pop()?.split('/').pop() || filePath;
	}
}