DeepCodeGenius-vscode/src/extension/ChatViewProvider.ts

import * as vscode from 'vscode';
import * as path from 'path';
import * as fs from 'fs';
import * as WS from 'ws';

type ListenerParam = {
  type: string,
  data: any,
  filePath?: string
  functionName?: string
  mdPath?: string
}
type WebviewViewIns = vscode.WebviewView['webview']

export class ChatViewProvider implements vscode.WebviewViewProvider {
  public static readonly viewType = 'aiChatView';

  private _view?: vscode.WebviewView;
  private _logger: (message: string) => void;
  private _wsInstance?: WS | null
  constructor(
    private readonly _extensionUri: vscode.Uri,
    logger?: (message: string) => void
  ) {
    this._logger = logger || ((message: string) => console.log(message));
  }

  public resolveWebviewView(
    webviewView: vscode.WebviewView,
    _context: vscode.WebviewViewResolveContext,
    _token: vscode.CancellationToken,
  ) {
    this._view = webviewView;
    this._logger(`WebView视图已创建: ${ChatViewProvider.viewType}`);

    webviewView.webview.options = {
      // 启用JavaScript
      enableScripts: true,
      localResourceRoots: [
        this._extensionUri
      ]
    };
    webviewView.webview.html = this._getHtmlForWebview(webviewView.webview);
    const webview = webviewView.webview
    // 在生产环境中，使用打包后的Webview文件
    this._logger('WebView HTML内容已设置');
    let wsInstance = null
    // 处理来自Webview的消息
    webviewView.webview.onDidReceiveMessage(message => {
      this._logger(`收到WebView消息: ${message.type}`);
      if (message.type?.startsWith('ws:')) {
        this.wsConnect(message, webview)
        return
      }
      switch (message.type) {
        case 'loadMd':
          this.fetchMdFile(message, webview)
          break
        case 'openFileAndGoToFunction':
          this.openFileAndGoToFunction(message)
          break
        case 'getSettings':
          this._logger(`处理getSettings请求`);
          this._sendSettings();
          break;
        case 'saveSettings':
          this._logger(`保存设置: ${JSON.stringify(message.settings, (key, value) =>
            key === 'apiKey' ? '***' : value)}`);
          this._saveSettings(message.settings);
          break;
        case 'log':
          // 根据消息来源分类日志
          if (message.message) {
            // ChatPanel组件日志处理
            if (message.message.includes('[ChatPanel]')) {
              this._logger(`📱 ChatPanel日志: ${message.message}`);
            }
            // 为SettingsModal相关日志添加特殊标记
            else if (message.message.includes('SettingsModal')) {
              this._logger(`⚙️ SettingsModal日志: ${message.message}`);
            }
            // 为WebSocket连接相关日志添加更详细的输出
            else if (message.message.includes('WebSocket') ||
              message.message.includes('ChatStore') ||
              message.message.includes('initWebSocketService')
            ) {
              this._logger(`🔌 WebSocket日志: ${message.message}`);
            }
            // 为重要日志添加特殊标记
            else if (message.message.includes('【重要】')) {
              this._logger(`⚠️ ${message.message}`);
            }
            else {
              this._logger(`WebView日志: ${message.message}`);
            }
          } else {
            this._logger(`WebView日志: ${JSON.stringify(message)}`);
          }
          break;
        case 'error':
          this._logger(`WebView错误: ${message.message}`);
          if (message.stack) {
            this._logger(`错误堆栈: ${message.stack}`);
          }
          break;
      }
    });

    // 初始化时发送主题颜色
    this._sendThemeColors();

    // 监听主题变化
    vscode.window.onDidChangeActiveColorTheme(() => {
      this._sendThemeColors();
    });
  }

  // 发送主题颜色
  private _sendThemeColors() {
    if (!this._view) return;

    // 获取当前主题颜色
    const theme = vscode.window.activeColorTheme;
    const isDark = theme.kind === vscode.ColorThemeKind.Dark;
    const isLight = theme.kind === vscode.ColorThemeKind.Light;

    // 获取常用颜色
    const colors = {
      // 基础颜色
      foreground: this._getColor('foreground'),
      background: this._getColor('editor.background'),

      // 文本颜色
      descriptionForeground: this._getColor('descriptionForeground'),
      errorForeground: this._getColor('errorForeground'),

      // 按钮颜色
      buttonBackground: this._getColor('button.background'),
      buttonForeground: this._getColor('button.foreground'),
      buttonHoverBackground: this._getColor('button.hoverBackground'),

      // 输入框颜色
      inputBackground: this._getColor('input.background'),
      inputForeground: this._getColor('input.foreground'),
      inputPlaceholderForeground: this._getColor('input.placeholderForeground'),
      inputBorder: this._getColor('input.border'),

      // 链接颜色
      linkForeground: this._getColor('textLink.foreground'),
      linkActiveForeground: this._getColor('textLink.activeForeground'),

      // 其他界面元素
      panelBorder: this._getColor('panel.border'),

      // 主题类型
      isDark,
      isLight
    };

    this._view.webview.postMessage({
      type: 'themeColors',
      colors
    });
  }

  // 获取主题颜色
  private _getColor(colorId: string): string | undefined {
    // 提供一组默认颜色值
    const defaultColors: Record<string, string> = {
      'foreground': '#cccccc',
      'editor.background': '#1e1e1e',
      'descriptionForeground': '#999999',
      'errorForeground': '#f48771',
      'button.background': '#0e639c',
      'button.foreground': '#ffffff',
      'button.hoverBackground': '#1177bb',
      'input.background': '#3c3c3c',
      'input.foreground': '#cccccc',
      'input.placeholderForeground': '#a6a6a6',
      'input.border': '#3c3c3c',
      'textLink.foreground': '#3794ff',
      'textLink.activeForeground': '#3794ff',
      'panel.border': '#80808059'
    };

    try {
      // 尝试从配置中获取颜色
      const themeOverride = vscode.workspace.getConfiguration('workbench')
        .get<Record<string, string>>('colorCustomizations', {});

      // 首先尝试从用户配置的覆盖颜色中获取
      if (themeOverride[colorId]) {
        return themeOverride[colorId];
      }

      // 然后从默认颜色中获取
      return defaultColors[colorId] || '#cccccc';
    } catch (error) {
      this._logger(`获取颜色失败 ${colorId}: ${error}`);
      // 回退到默认颜色
      return defaultColors[colorId] || '#cccccc';
    }
  }

  // 发送设置
  private _sendSettings() {
    if (!this._view) return;

    const config = vscode.workspace.getConfiguration('aiChat');
    const settings = {
      apiHost: config.get('apiHost', ''),
      apiKey: config.get('apiKey', '')
    };

    this._view.webview.postMessage({
      type: 'settings',
      settings
    });
  }

  // 保存设置
  private _saveSettings(settings: { apiHost: string; apiKey: string }) {
    if (!settings) return;

    this._logger(`准备保存设置: apiHost=${settings.apiHost}, apiKey=***`);

    const config = vscode.workspace.getConfiguration('aiChat');
    let hasChanges = false;

    // 保存成功标志
    let success = true;
    let errorMessage = '';

    // 检查并更新API主机
    const updateHost = async () => {
      if (settings.apiHost === config.get('apiHost')) return;

      hasChanges = true;
      this._logger('更新API主机设置');

      try {
        await config.update('apiHost', settings.apiHost, vscode.ConfigurationTarget.Global);
        this._logger('已更新API主机设置');
      } catch (error) {
        success = false;
        errorMessage = error instanceof Error ? error.message : '更新API主机失败';
        this._logger(`更新API主机设置失败: ${errorMessage}`);
      }
    };

    // 检查并更新API密钥
    const updateKey = async () => {
      if (settings.apiKey === config.get('apiKey')) return;

      hasChanges = true;
      this._logger('更新API密钥设置');

      try {
        await config.update('apiKey', settings.apiKey, vscode.ConfigurationTarget.Global);
        this._logger('已更新API密钥设置');
      } catch (error) {
        success = false;
        errorMessage = error instanceof Error ? error.message : '更新API密钥失败';
        this._logger(`更新API密钥设置失败: ${errorMessage}`);
      }
    };

    // 执行所有更新操作
    const updateAll = async () => {
      await updateHost();
      await updateKey();

      // 通知前端设置已保存
      if (this._view) {
        if (success) {
          this._logger('所有设置保存成功');
          this._view.webview.postMessage({
            type: 'settingsSaved',
            success: true
          });
        } else {
          this._logger(`保存设置失败: ${errorMessage}`);
          this._view.webview.postMessage({
            type: 'settingsSaved',
            success: false,
            error: `保存设置失败: ${errorMessage}`
          });
        }
      }
    };

    // 启动更新过程
    updateAll().catch(error => {
      this._logger(`保存设置过程中发生未捕获异常: ${error}`);

      // 确保前端收到失败消息
      if (this._view) {
        this._view.webview.postMessage({
          type: 'settingsSaved',
          success: false,
          error: '保存设置过程中发生错误'
        });
      }
    });
  }

  // 生成Webview HTML内容
  private _getHtmlForWebview(webview: vscode.Webview) {
    // 在生产环境中，使用打包后的Webview文件
    const distPath = path.join(this._extensionUri.fsPath, 'webview', 'dist');

    // 生成随机nonce以增强安全性
    const nonce = this._getNonce();

    // 检查dist目录是否存在
    if (fs.existsSync(distPath)) {
      try {
        // 读取生产环境的index.html
        let indexHtml = fs.readFileSync(
          path.join(distPath, 'index.html'),
          'utf-8'
        );

        // 添加CSP策略和nonce
        if (!indexHtml.includes('content-security-policy')) {
          indexHtml = indexHtml.replace(
            '<head>',
            `<head>
            <meta http-equiv="Content-Security-Policy" content="default-src 'none'; connect-src ws://localhost:8080; style-src ${webview.cspSource} 'unsafe-inline'; img-src ${webview.cspSource} https: data:; script-src 'nonce-${nonce}';">`
          );
        }

        // 确保有VSCode API脚本
        if (!indexHtml.includes('acquireVsCodeApi')) {
          indexHtml = indexHtml.replace(
            '<head>',
            `<head>
            <script nonce="${nonce}">
              const vscode = acquireVsCodeApi();
              window.vscodeApi = vscode;
            </script>`
          );
        }

        // 将资源路径替换为Webview可访问的URI
        indexHtml = indexHtml.replace(
          /(href|src)="([^"]*)"/g,
          (_, attribute, url) => {
            // 跳过外部URL或数据URI
            if (url.startsWith('http') || url.startsWith('data:')) {
              return `${attribute}="${url}"`;
            }

            // 转换为Webview URI
            const resourceUri = webview.asWebviewUri(
              vscode.Uri.joinPath(this._extensionUri, 'webview', 'dist', url)
            );

            // 为脚本添加nonce
            if (attribute === 'src' && url.endsWith('.js')) {
              return `${attribute}="${resourceUri}" nonce="${nonce}"`;
            }

            return `${attribute}="${resourceUri}"`;
          }
        );

        this._logger('WebView HTML内容已生成，包含VSCode API初始化脚本');
        return indexHtml;
      } catch (error) {
        this._logger(`生成WebView HTML时出错: ${error}`);
        return this._getFallbackHtml(webview, nonce);
      }
    } else {
      return this._getFallbackHtml(webview, nonce);
    }
  }

  // 获取回退HTML内容
  private _getFallbackHtml(webview: vscode.Webview, nonce: string) {
    // 开发环境下使用占位符HTML
    return `
      <!DOCTYPE html>
      <html lang="zh-CN">
        <head>
          <meta charset="UTF-8">
          <meta name="viewport" content="width=device-width, initial-scale=1.0">
          <meta http-equiv="Content-Security-Policy" content="default-src 'none'; connect-src ws://localhost:8080; style-src ${webview.cspSource} 'unsafe-inline'; script-src 'nonce-${nonce}';">
          <title>AI Chat</title>
          <script nonce="${nonce}">
            const vscode = acquireVsCodeApi();
            window.vscodeApi = vscode;
          </script>
          <style>
            body {
              background-color: var(--vscode-editor-background);
              color: var(--vscode-editor-foreground);
              font-family: var(--vscode-font-family);
              padding: 20px;
              display: flex;
              flex-direction: column;
              align-items: center;
              justify-content: center;
              height: 100vh;
              margin: 0;
            }
            h1 {
              margin-bottom: 20px;
            }
            p {
              max-width: 600px;
              text-align: center;
            }
          </style>
        </head>
        <body>
          <h1>AI Chat</h1>
          <p>
            开发模式：请运行 "npm run webview:build" 来构建Webview前端，或在开发时使用 "npm run dev" 进行热重载开发。
          </p>
        </body>
      </html>
    `;
  }

  // 生成随机nonce
  private _getNonce() {
    let text = '';
    const possible = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789';
    for (let i = 0; i < 32; i++) {
      text += possible.charAt(Math.floor(Math.random() * possible.length));
    }
    return text;
  }

  wsConnect(message: ListenerParam, webview: WebviewViewIns): void | false {
    switch (message.type) {
      case 'ws:connect':
        // 建立 WebSocket 连接
        this._wsInstance = new WS('ws://47.117.75.243:8080/ws', {
          headers: { 'X-Api-Key': 'simpletest2025_094' }
        })
        // 连接打开
        this._wsInstance.on('open', () => {
          webview.postMessage({ type: 'ws:open', data: true })
        })
        // 转发服务器消息到 Webview
        this._wsInstance.on('message', (data) => {
          this._logger('ws:message ===== ' + data.toString())
          webview.postMessage({
            type: 'ws:message',
            data: data
          })
        })

        // 错误处理
        this._wsInstance.on('error', (err) => {
          webview.postMessage({
            type: 'ws:error',
            data: '连接失败: ' + err.message
          })
        })

        // 连接关闭处理
        this._wsInstance.on('close', () => {
          webview.postMessage({ type: 'ws:close' })
        })
        break

      case 'ws:send':
        // 转发 Webview 的消息到服务器
        if (this._wsInstance) {
          this._logger('ws:send' + message.data)
          this._wsInstance.send(message.data)
        }
        break
      case 'ws:disconnect':
        if (this._wsInstance) {
          this._wsInstance.close()
          this._wsInstance = null
        }
        break
      default:
        return false
    }
  }
  findFunction(symbols: vscode.DocumentSymbol[], name: string): vscode.DocumentSymbol | false {
    for (const symbol of symbols) {
        if (symbol.kind === vscode.SymbolKind.Function && symbol.name === name) {
            return symbol;
        }
        if (symbol.children) {
            const found = this.findFunction(symbol.children, name);
            if (found) return found;
        }
    }
    return false;
  }
  async openFileAndGoToFunction(message: ListenerParam): Promise<void | false> {
    if (message.type !== 'openFileAndGoToFunction') { return false };
    const filePath = message.filePath; // 确保路径正确，可以是绝对路径或相对于工作区
    const targetFunc = message.functionName
    if (!filePath || !targetFunc){
      vscode.window.showWarningMessage(`未提供文件: ${filePath} 或函数: ${targetFunc}`)
      return
    }
    try {
      // 打开文件
      const uri = vscode.Uri.file(filePath)
      const doc = await vscode.workspace.openTextDocument(uri)
      const editor = await vscode.window.showTextDocument(doc)

      // 获取文档符号
      const symbols = await vscode.commands.executeCommand<vscode.DocumentSymbol[]>(
        'vscode.executeDocumentSymbolProvider',
        doc.uri
      )

      if (symbols) {
        const funcSymbol = this.findFunction(symbols, targetFunc)
        if (funcSymbol) {
          // 跳转到目标位置
          editor.revealRange(funcSymbol.range, vscode.TextEditorRevealType.InCenter)
          editor.selection = new vscode.Selection(funcSymbol.range.start, funcSymbol.range.end)
        } else {
          vscode.window.showWarningMessage(`未找到函数 ${targetFunc}`)
        }
      }
    } catch (error) {
      vscode.window.showErrorMessage(`无法打开文件: ${error}`)
    }
  }
  // 获取本地markdown文件
  fetchMdFile(message: ListenerParam, webview: WebviewViewIns): false | void {
    if (message.type !== 'loadMd') { return false }
    if (message.mdPath) {
      const filePath = path.join(this._extensionUri.fsPath, 'webview', 'dist', message.mdPath!.replace('__localWebViewPath/', ''))
      try {
        const content = fs.readFileSync(filePath, 'utf-8')
        webview.postMessage({ type: 'mdContent', data: content })
      } catch (error) {
        this._logger('获取 md 内容失败' + error)
        webview.postMessage({ type: 'mdContent', data: '获取 md 内容失败' })
      }
      return
    }
    webview.postMessage({ type: 'mdContent', data: '未正确传入md文件路径' })
  }
}