Prompter/src/providers/prompterTreeProvider.ts

import * as vscode from 'vscode';
import * as path from 'path';
import { FileTreeItem, SettingTreeItem } from '../models/treeItems';
import { PrompterSettings } from '../models/settings';
import { FileSelectionManager } from './fileSelectionManager';
import { SettingsManager } from './settingsManager';

/**
 * Tree provider for the Prompter extension
 * Handles both file browsing and settings views
 */
export class PrompterTreeProvider implements vscode.TreeDataProvider<FileTreeItem | SettingTreeItem> {
    private _onDidChangeTreeData: vscode.EventEmitter<FileTreeItem | SettingTreeItem | undefined | null | void> = 
        new vscode.EventEmitter<FileTreeItem | SettingTreeItem | undefined | null | void>();

    readonly onDidChangeTreeData: vscode.Event<FileTreeItem | SettingTreeItem | undefined | null | void> = 
        this._onDidChangeTreeData.event;

    private fileSelectionManager: FileSelectionManager;
    private settingsManager: SettingsManager;
    private showingSettings: boolean = false;

    constructor(private workspaceRoot: string | undefined) {
        this.fileSelectionManager = new FileSelectionManager();
        this.settingsManager = new SettingsManager();
    }

    // Toggle between file view and settings view
    toggleSettingsView(): void {
        this.showingSettings = !this.showingSettings;
        this.refresh();
    }

    showSettingsView(): void {
        this.showingSettings = true;
        this.refresh();
    }

    showFilesView(): void {
        this.showingSettings = false;
        this.refresh();
    }

    getTreeItem(element: FileTreeItem | SettingTreeItem): vscode.TreeItem {
        // Return the element as is if it's a SettingTreeItem
        if (element instanceof SettingTreeItem) {
            return element;
        }

        // Handle FileTreeItem
        const fileItem = element as FileTreeItem;
        if (fileItem.resourceUri && this.fileSelectionManager.isSelected(fileItem.resourceUri.fsPath)) {
            fileItem.checkboxState = vscode.TreeItemCheckboxState.Checked;
        } else {
            fileItem.checkboxState = vscode.TreeItemCheckboxState.Unchecked;
        }
        return fileItem;
    }

    async getChildren(element?: FileTreeItem | SettingTreeItem): Promise<(FileTreeItem | SettingTreeItem)[]> {
        // If we're showing settings, return settings items
        if (this.showingSettings && !element) {
            return this.getSettingsItems();
        }

        // Otherwise show files
        if (!this.workspaceRoot) {
            vscode.window.showInformationMessage('No workspace folder is opened');
            return Promise.resolve([]);
        }

        if (element) {
            // Make sure we're dealing with a FileTreeItem that has a resourceUri
            if (!(element instanceof SettingTreeItem) && element.resourceUri) {
                return this.getFilesInDirectory(element.resourceUri.fsPath);
            }
            return Promise.resolve([]);
        } else {
            return this.getFilesInDirectory(this.workspaceRoot);
        }
    }

    private async getFilesInDirectory(dirPath: string): Promise<FileTreeItem[]> {
        try {
            const files = await vscode.workspace.fs.readDirectory(vscode.Uri.file(dirPath));

            return files.map(([name, type]) => {
                const filePath = path.join(dirPath, name);
                const uri = vscode.Uri.file(filePath);

                return new FileTreeItem(
                    uri,
                    type === vscode.FileType.Directory
                        ? vscode.TreeItemCollapsibleState.Collapsed
                        : vscode.TreeItemCollapsibleState.None
                );
            });
        } catch (error) {
            console.error(`Error reading directory: ${dirPath}`, error);
            return [];
        }
    }

    // Refresh the tree view
    refresh(): void {
        this._onDidChangeTreeData.fire();
    }

    // Add a file to the selection
    addToSelection(item: FileTreeItem): void {
        if (item.resourceUri) {
            if (this.isDirectory(item)) {
                this.fileSelectionManager.addDirectory(item.resourceUri.fsPath);
            } else {
                this.fileSelectionManager.addFile(item.resourceUri.fsPath);
            }
        }
    }

    // Remove a file from the selection
    removeFromSelection(item: FileTreeItem): void {
        if (item.resourceUri) {
            if (this.isDirectory(item)) {
                this.fileSelectionManager.removeDirectory(item.resourceUri.fsPath);
            } else {
                this.fileSelectionManager.removeFile(item.resourceUri.fsPath);
            }
        }
    }

    private isDirectory(item: FileTreeItem): boolean {
        return item.collapsibleState === vscode.TreeItemCollapsibleState.Collapsed || 
               item.collapsibleState === vscode.TreeItemCollapsibleState.Expanded;
    }

    // Toggle a file's selection status
    toggleSelection(item: FileTreeItem): void {
        if (item.resourceUri) {
            const filePath = item.resourceUri.fsPath;
            if (this.fileSelectionManager.isSelected(filePath)) {
                this.removeFromSelection(item);
            } else {
                this.addToSelection(item);
            }
            this.refresh();
        }
    }

    // Get the selected files
    getSelectedFiles(): Set<string> {
        return this.fileSelectionManager.getSelectedFiles();
    }

    // Get settings items for the tree view
    private getSettingsItems(): SettingTreeItem[] {
        const settings = this.settingsManager.getSettings();
        return [
            new SettingTreeItem('Include Formatting Instructions', 'includeFormattingInstructions', settings.includeFormattingInstructions),
            new SettingTreeItem('Token Calculation', 'tokenCalculationEnabled', settings.tokenCalculationEnabled),
            new SettingTreeItem('Include File Map', 'includeFileMap', settings.includeFileMap)
        ];
    }

    // Update a setting value
    updateSetting(key: keyof PrompterSettings, value: boolean): void {
        this.settingsManager.updateSetting(key, value);
        this.refresh();
    }

    // Check if showing settings view
    isShowingSettings(): boolean {
        return this.showingSettings;
    }

    // Toggle formatting instructions setting
    toggleXmlEdits(): void {
        this.settingsManager.toggleFormattingInstructions();
        this.refresh();
    }

    // Check if formatting instructions are enabled
    isXmlEditsEnabled(): boolean {
        return this.settingsManager.isFormattingInstructionsEnabled();
    }

    // Get all settings
    getSettings(): PrompterSettings {
        return this.settingsManager.getSettings();
    }

    // Update all settings at once
    updateSettings(newSettings: PrompterSettings): void {
        this.settingsManager.updateAllSettings(newSettings);
        this.refresh();
    }
}