class PromptTemplate {
    [string]$Name
    [string]$Template
    [hashtable]$Variables = @{}
    [string]$SystemContext

    PromptTemplate([string]$Name, [string]$Template) {
        $this.Name = $Name
        $this.Template = $Template
    }

    [void] SetVariable([string]$Key, [object]$Value) {
        $this.Variables[$Key] = $Value
    }

    [void] SetSystemContext([string]$Context) {
        $this.SystemContext = $Context
    }

    [string] Render() {
        $rendered = $this.Template
        foreach ($key in $this.Variables.Keys) {
            $placeholder = "{{$key}}"
            $rendered = $rendered -replace [regex]::Escape($placeholder), $this.Variables[$key]
        }
        return $rendered
    }

    [hashtable] ToChatMessages() {
        $messages = @()
        if ($this.SystemContext) {
            $messages += @{
                role    = 'system'
                content = $this.SystemContext
            }
        }
        $messages += @{
            role    = 'user'
            content = $this.Render()
        }
        return @{ messages = $messages }
    }
}

# 创建运维场景的 Prompt 模板库
$templates = @{}

$templates['error-diagnosis'] = [PromptTemplate]::new(
    'error-diagnosis',
    '请分析以下 PowerShell 错误信息，给出可能的原因和修复建议。' +
    "`n`n错误类型：{ErrorType}" +
    "`n错误消息：{ErrorMessage}" +
    '`n发生环境：{Environment}'
)

$templates['script-generation'] = [PromptTemplate]::new(
    'script-generation',
    '请生成一个 PowerShell 脚本，要求如下：' +
    "`n目标：{Goal}" +
    "`n约束条件：{Constraints}" +
    "`n目标平台：{Platform}" +
    '`n请使用 PowerShell 7 兼容语法，包含错误处理和注释。'
)

# 设置系统上下文：定义 LLM 的角色和行为规范
$systemContext = @"
你是一位资深的 PowerShell 和 Azure 运维专家。
你的回答需要：
1. 使用 PowerShell 7 兼容语法
2. 包含完整的错误处理（try/catch/finally）
3. 遵循最佳实践（Approved Verbs、强类型、注释）
4. 输出可直接执行的脚本代码
"@

$templates['error-diagnosis'].SetSystemContext($systemContext)
$templates['script-generation'].SetSystemContext($systemContext)

# 使用模板：填充变量并渲染
$tpl = $templates['error-diagnosis']
$tpl.SetVariable('ErrorType', 'InvalidOperationException')
$tpl.SetVariable('ErrorMessage', 'Collection was modified; enumeration operation may not execute.')
$tpl.SetVariable('Environment', 'Windows Server 2022, PowerShell 7.4')

$renderedPrompt = $tpl.Render()
$chatPayload = $tpl.ToChatMessages()

Write-Host "=== 渲染后的 Prompt ===" -ForegroundColor Cyan
Write-Host $renderedPrompt
Write-Host "`n=== Chat API 载荷 ===" -ForegroundColor Cyan
$chatPayload | ConvertTo-Json -Depth 3

function Invoke-StructuredCompletion {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [string]$UserPrompt,

        [Parameter(Mandatory)]
        [string]$JsonSchema,

        [string]$Model = 'gpt-4o-mini',

        [int]$MaxRetries = 3,

        [string]$Endpoint = 'http://localhost:11434/v1/chat/completions'
    )

    # 构建系统提示词，强制 JSON 输出并给出 Schema
    $systemPrompt = @"
你是一个 PowerShell 运维助手。用户会提出运维相关的问题，你需要给出结构化的回答。

严格要求：
1. 只输出合法的 JSON，不要包含任何其他文字说明
2. 严格遵循以下 JSON Schema：
$JsonSchema
3. 不要输出 Markdown 代码块标记，直接输出 JSON
4. 字符串值使用 UTF-8 编码
"@

    $attempt = 0
    while ($attempt -lt $MaxRetries) {
        $attempt++
        Write-Verbose "尝试第 $attempt 次..."

        $body = @{
            model       = $Model
            messages    = @(
                @{ role = 'system'; content = $systemPrompt },
                @{ role = 'user';   content = $UserPrompt }
            )
            temperature = 0.1
        } | ConvertTo-Json -Depth 5

        try {
            $response = Invoke-RestMethod -Uri $Endpoint -Method Post `
                -ContentType 'application/json' `
                -Body ([System.Text.Encoding]::UTF8.GetBytes($body)) `
                -ErrorAction Stop

            $content = $response.choices[0].message.content.Trim()

            # 去除可能的 Markdown 代码块标记
            if ($content -match '^\x60{3}json?\s*([\s\S]*?)\x60{3}$') {
                $content = $Matches[1].Trim()
            }

            # 尝试解析 JSON
            $parsed = $content | ConvertFrom-Json -ErrorAction Stop

            Write-Verbose "JSON 解析成功"
            return [PSCustomObject]@{
                Success   = $true
                Data      = $parsed
                Raw       = $content
                Attempts  = $attempt
            }
        }
        catch {
            Write-Warning "第 $attempt 次尝试失败：$($_.Exception.Message)"
            if ($attempt -ge $MaxRetries) {
                return [PSCustomObject]@{
                    Success  = $false
                    Error    = $_.Exception.Message
                    Attempts = $attempt
                }
            }
            Start-Sleep -Milliseconds ($attempt * 500)
        }
    }
}

# 定义输出 Schema：服务器健康检查结果
$schema = @'
{
  "type": "object",
  "properties": {
    "server_name": { "type": "string" },
    "overall_status": { "type": "string", "enum": ["healthy", "warning", "critical"] },
    "checks": {
      "type": "array",
      "items": {
        "type": "object",
        "properties": {
          "item": { "type": "string" },
          "status": { "type": "string" },
          "detail": { "type": "string" },
          "recommendation": { "type": "string" }
        },
        "required": ["item", "status", "detail", "recommendation"]
      }
    },
    "summary": { "type": "string" }
  },
  "required": ["server_name", "overall_status", "checks", "summary"]
}
'@

$result = Invoke-StructuredCompletion -UserPrompt @'
WebSrv01 服务器状态：CPU 使用率 92%，内存剩余 2GB/32GB，
磁盘 C: 剩余 15%，事件日志中发现 37 个 Error 级别事件（过去 1 小时），
IIS 应用池 MyAppPool 已崩溃 3 次。
'@ -JsonSchema $schema -Verbose

if ($result.Success) {
    Write-Host "服务器：$($result.Data.server_name)" -ForegroundColor Cyan
    Write-Host "状态：$($result.Data.overall_status)" -ForegroundColor Yellow
    Write-Host "`n检查项："
    $result.Data.checks | Format-Table -Property item, status, detail -Wrap
    Write-Host "摘要：$($result.Data.summary)"
}

class OpsCopilot {
    [string]$Model
    [string]$Endpoint
    [hashtable]$Templates
    [string]$SystemInfo

    OpsCopilot([string]$Endpoint, [string]$Model) {
        $this.Endpoint = $Endpoint
        $this.Model = $Model
        $this.Templates = @{}
        $this.SystemInfo = $this.CollectSystemInfo()
        $this.RegisterTemplates()
    }

    [string] CollectSystemInfo() {
        $info = @{
            hostname     = $env:COMPUTERNAME
            os           = (Get-CimInstance Win32_OperatingSystem).Caption
            psVersion    = $PSVersionTable.PSVersion.ToString()
            dotnet       = [System.Environment]::Version.ToString()
            cpu          = (Get-CimInstance Win32_Processor).Name
            totalMemGB   = [math]::Round(
                (Get-CimInstance Win32_ComputerSystem).TotalPhysicalMemory / 1GB, 1
            )
            drives       = Get-CimInstance Win32_LogicalDisk -Filter 'DriveType=3' |
                ForEach-Object {
                    @{
                        drive     = $_.DeviceID
                        freeGB    = [math]::Round($_.FreeSpace / 1GB, 1)
                        totalGB   = [math]::Round($_.Size / 1GB, 1)
                    }
                }
        }
        return ($info | ConvertTo-Json -Depth 3)
    }

    [void] RegisterTemplates() {
        $this.Templates['ask'] = @'
基于以下系统信息回答运维问题。

系统信息：
{SystemInfo}

问题：{Question}

请给出：
1. 问题分析
2. 具体的 PowerShell 命令或脚本
3. 注意事项
'@

        $this.Templates['diagnose'] = @'
基于以下系统信息诊断错误。

系统信息：
{SystemInfo}

错误详情：
{ErrorDetail}

请分析：
1. 错误的根本原因
2. 修复步骤（含 PowerShell 命令）
3. 预防措施
'@
    }

    [string] RenderTemplate([string]$TemplateName, [hashtable]$Vars) {
        $rendered = $this.Templates[$TemplateName]
        $rendered = $rendered -replace '\{SystemInfo\}', $this.SystemInfo
        foreach ($key in $Vars.Keys) {
            $rendered = $rendered -replace "\{$key\}", $Vars[$key]
        }
        return $rendered
    }

    [string] Chat([string]$UserMessage) {
        $systemPrompt = @"
你是一位精通 PowerShell 和 Windows Server 的运维专家。
当前系统环境：
$($this.SystemInfo)

回答要求：
- 优先使用 PowerShell 7 原生命令
- 脚本包含 try/catch 错误处理
- 给出可立即执行的命令
- 用中文回答
"@

        $body = @{
            model       = $this.Model
            messages    = @(
                @{ role = 'system'; content = $systemPrompt },
                @{ role = 'user';   content = $UserMessage }
            )
            temperature = 0.2
        } | ConvertTo-Json -Depth 5

        $response = Invoke-RestMethod -Uri $this.Endpoint -Method Post `
            -ContentType 'application/json' `
            -Body ([System.Text.Encoding]::UTF8.GetBytes($body))

        return $response.choices[0].message.content
    }

    [string] Ask([string]$Question) {
        $prompt = $this.RenderTemplate('ask', @{
            Question = $Question
        })
        return $this.Chat($prompt)
    }

    [string] Diagnose([string]$ErrorDetail) {
        $prompt = $this.RenderTemplate('diagnose', @{
            ErrorDetail = $ErrorDetail
        })
        return $this.Chat($prompt)
    }
}

# 初始化 Copilot（使用本地 Ollama 或远程 API）
$copilot = [OpsCopilot]::new('http://localhost:11434/v1/chat/completions', 'qwen3:8b')

# 示例 1：提出运维问题
Write-Host "=== 运维问答 ===" -ForegroundColor Cyan
$answer = $copilot.Ask('如何查找占用磁盘空间最大的前 10 个文件夹？')
Write-Host $answer

# 示例 2：诊断错误
Write-Host "`n=== 错误诊断 ===" -ForegroundColor Cyan
$diagnosis = $copilot.Diagnose(@"
执行 Get-EventLog -LogName Application -Newest 100 时报错：
"Requested registry access is not allowed"
当前以普通用户身份运行 PowerShell。
"@)
Write-Host $diagnosis

function Get-OllamaModel {
    <# 获取本地已安装的模型列表 #>
    $response = Invoke-RestMethod -Uri 'http://localhost:11434/api/tags' -Method Get
    $response.models | ForEach-Object {
        [PSCustomObject]@{
            Name       = $_.name
            SizeMB     = [math]::Round($_.size / 1MB, 1)
            ModifiedAt = $_.modified_at
            Family     = $_.details.family
            Parameter  = $_.details.parameter_size
        }
    } | Sort-Object Name | Format-Table -AutoSize
}

function Invoke-OllamaChat {
    <#
    .SYNOPSIS
        向本地 Ollama 模型发送对话请求
    .PARAMETER Model
        模型名称，如 qwen2.5:7b、llama3.1:8b
    .PARAMETER Messages
        消息数组，每条包含 role 和 content
    .PARAMETER Stream
        是否启用流式输出（默认 $false）
    .PARAMETER Temperature
        生成温度，0.0 到 1.0 之间
    #>
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [string]$Model,

        [Parameter(Mandatory)]
        [hashtable[]]$Messages,

        [switch]$Stream,

        [ValidateRange(0.0, 1.0)]
        [double]$Temperature = 0.7
    )

    $body = @{
        model      = $Model
        messages   = $Messages
        stream     = $Stream.IsPresent
        options    = @{
            temperature = $Temperature
        }
    } | ConvertTo-Json -Depth 5

    if ($Stream) {
        # 流式输出：逐行读取 SSE 响应
        $request = [System.Net.Http.HttpRequestMessage]::new(
            [System.Net.Http.HttpMethod]::Post,
            'http://localhost:11434/api/chat'
        )
        $request.Content = [System.Net.Http.StringContent]::new(
            $body, [System.Text.Encoding]::UTF8, 'application/json'
        )

        $client = [System.Net.Http.HttpClient]::new()
        $response = $client.Send($request, [System.Net.Http.HttpCompletionOption]::ResponseHeadersRead)
        $reader = [System.IO.StreamReader]::new($response.Content.ReadAsStream())

        $fullResponse = [System.Text.StringBuilder]::new()
        while (-not $reader.EndOfStream) {
            $line = $reader.ReadLine()
            if ($line) {
                $chunk = $line | ConvertFrom-Json
                if ($chunk.message.content) {
                    Write-Host $chunk.message.content -NoNewline
                    [void]$fullResponse.Append($chunk.message.content)
                }
                if ($chunk.done) { break }
            }
        }
        Write-Host ''
        $reader.Close()
        $client.Dispose()
        return $fullResponse.ToString()
    }
    else {
        $result = Invoke-RestMethod -Uri 'http://localhost:11434/api/chat' `
            -Method Post -Body $body -ContentType 'application/json'
        return $result.message.content
    }
}

# 用法示例：查看已安装模型
Get-OllamaModel

# 用法示例：单轮对话
$messages = @(
    @{ role = 'system'; content = '你是一个 PowerShell 专家，回答简洁精准。' }
    @{ role = 'user';   content = '如何获取系统中占用磁盘最大的 10 个目录？' }
)

$answer = Invoke-OllamaChat -Model 'qwen2.5:7b' -Messages $messages -Temperature 0.3
Write-Host "`n--- 模型回复 ---`n$answer"

class LocalVectorStore {
    <# 本地向量存储，用于 RAG 检索 #>
    [System.Collections.Generic.List[hashtable]]$Documents

    LocalVectorStore() {
        $this.Documents = [System.Collections.Generic.List[hashtable]]::new()
    }

    [void] AddChunk([string]$Content, [string]$Source, [double[]]$Embedding) {
        $this.Documents.Add(@{
            content  = $Content
            source   = $Source
            embedding = $Embedding
        })
    }

    [double] CosineSimilarity([double[]]$A, [double[]]$B) {
        $dot = 0.0; $normA = 0.0; $normB = 0.0
        for ($i = 0; $i -lt $A.Count; $i++) {
            $dot += $A[$i] * $B[$i]
            $normA += $A[$i] * $A[$i]
            $normB += $B[$i] * $B[$i]
        }
        if ($normA -eq 0 -or $normB -eq 0) { return 0 }
        return $dot / ([math]::Sqrt($normA) * [math]::Sqrt($normB))
    }

    [hashtable[]] Search([double[]]$QueryEmbedding, [int]$TopK = 3) {
        return $this.Documents | ForEach-Object {
            [PSCustomObject]@{
                Content    = $_.content
                Source     = $_.source
                Similarity = $this.CosineSimilarity($QueryEmbedding, $_.embedding)
            }
        } | Sort-Object Similarity -Descending | Select-Object -First $TopK |
            ForEach-Object { @{ content = $_.Content; source = $_.Source; score = $_.Similarity } }
    }
}

function Get-LocalEmbedding {
    param([string]$Text)
    $body = @{ model = 'nomic-embed-text'; prompt = $Text } | ConvertTo-Json
    $result = Invoke-RestMethod -Uri 'http://localhost:11434/api/embeddings' `
        -Method Post -Body $body -ContentType 'application/json'
    return $result.embedding
}

function Split-Document {
    param(
        [string]$Content,
        [int]$ChunkSize = 500,
        [int]$Overlap = 50
    )
    $chunks = [System.Collections.Generic.List[string]]::new()
    $words = $Content -split '\s+'
    $i = 0
    while ($i -lt $words.Count) {
        $end = [math]::Min($i + $ChunkSize, $words.Count)
        $chunkText = ($words[$i..($end - 1)] -join ' ')
        [void]$chunks.Add($chunkText)
        $i += $ChunkSize - $Overlap
    }
    return $chunks
}

function New-RAGIndex {
    <# 从指定目录的文档构建 RAG 索引 #>
    param(
        [string]$Path = '.\docs',
        [int]$ChunkSize = 500
    )

    $store = [LocalVectorStore]::new()
    $files = Get-ChildItem -Path $Path -Include '*.md','*.txt' -Recurse

    foreach ($file in $files) {
        $content = Get-Content $file.FullName -Raw -Encoding UTF8
        $chunks = Split-Document -Content $content -ChunkSize $ChunkSize
        $chunkIndex = 0
        foreach ($chunk in $chunks) {
            $embedding = Get-LocalEmbedding -Text $chunk
            $store.AddChunk($chunk, "$($file.Name)#chunk-$chunkIndex", $embedding)
            $chunkIndex++
        }
        Write-Host "已索引: $($file.Name) ($($chunks.Count) 个切片)" -ForegroundColor Cyan
    }
    Write-Host "索引完成，共 $($store.Documents.Count) 个文档切片" -ForegroundColor Green
    return $store
}

function Invoke-RAGQuery {
    <# 基于 RAG 索引进行检索增强查询 #>
    param(
        [LocalVectorStore]$Store,
        [string]$Question,
        [string]$Model = 'qwen2.5:7b',
        [int]$TopK = 3
    )

    $queryEmbedding = Get-LocalEmbedding -Text $Question
    $results = $Store.Search($queryEmbedding, $TopK)

    $context = ($results | ForEach-Object {
        "【来源: $($_.source) | 相关度: $($_.score):N3】`n$($_.content)"
    }) -join "`n`n---`n`n"

    $messages = @(
        @{
            role    = 'system'
            content = '根据以下参考资料回答用户问题。如果资料中没有相关信息，请明确说明。'
        }
        @{
            role    = 'user'
            content = "参考资料：`n$context`n`n问题：$Question"
        }
    )

    Write-Host "`n检索到 $($results.Count) 条相关文档:" -ForegroundColor Yellow
    $results | ForEach-Object {
        Write-Host "  - $($_.source) (相关度: $($_.score):N3)" -ForegroundColor DarkGray
    }

    return Invoke-OllamaChat -Model $Model -Messages $messages -Temperature 0.2
}

# 用法示例：构建索引并查询
$store = New-RAGIndex -Path '.\knowledge-base'
$answer = Invoke-RAGQuery -Store $store -Question '公司的密码策略要求是什么？'
Write-Host "`n$answer"

function Invoke-AILogAnalysis {
    <# 使用本地 LLM 分析系统日志 #>
    param(
        [string]$LogPath = 'C:\Windows\System32\winevt\Logs\System.evtx',
        [int]$MaxEvents = 50,
        [string]$Model = 'qwen2.5:7b'
    )

    Write-Host "正在读取最近的 $MaxEvents 条事件日志..." -ForegroundColor Cyan
    $events = Get-WinEvent -Path $LogPath -MaxEvents $MaxEvents -ErrorAction SilentlyContinue

    if (-not $events) {
        $events = Get-WinEvent -LogName System -MaxEvents $MaxEvents
    }

    $errorEvents = $events | Where-Object { $_.Level -le 3 } |
        Select-Object TimeCreated, Id, LevelDisplayName, Message |
        Format-Table -AutoSize | Out-String

    $prompt = @"
以下是最近的事件日志（仅含错误和警告）：

$errorEvents

请分析以下内容：
1. 是否存在异常模式或重复性错误
2. 可能的根因分析
3. 推荐的处理步骤
"@

    $messages = @(
        @{ role = 'system'; content = '你是一位资深的 Windows 系统管理员，擅长日志分析和故障排查。' }
        @{ role = 'user';   content = $prompt }
    )

    return Invoke-OllamaChat -Model $Model -Messages $messages -Temperature 0.3
}

function Get-AICommandSuggestion {
    <# 根据自然语言描述推荐 PowerShell 命令 #>
    param(
        [Parameter(Mandatory)]
        [string]$Description,

        [string]$Model = 'qwen2.5:7b'
    )

    $messages = @(
        @{
            role    = 'system'
            content = '你是一个 PowerShell 命令助手。根据用户的自然语言描述，给出 1-3 个 PowerShell 命令方案。每个方案包含命令、参数说明和使用场景。只返回 PowerShell 代码和简短说明。'
        }
        @{
            role    = 'user'
            content = $Description
        }
    )

    return Invoke-OllamaChat -Model $Model -Messages $messages -Temperature 0.4
}

function Repair-AIScriptError {
    <# 使用本地 LLM 诊断和修复脚本错误 #>
    param(
        [Parameter(Mandatory)]
        [string]$ScriptContent,

        [Parameter(Mandatory)]
        [string]$ErrorMessage,

        [string]$Model = 'qwen2.5:7b'
    )

    $messages = @(
        @{
            role    = 'system'
            content = '你是一个 PowerShell 调试专家。分析脚本代码和错误信息，给出问题原因和修复后的完整脚本。'
        }
        @{
            role    = 'user'
            content = "脚本代码：`n$ScriptContent`n`n错误信息：`n$ErrorMessage"
        }
    )

    return Invoke-OllamaChat -Model $Model -Messages $messages -Temperature 0.2
}

# 用法示例 1：日志分析
$analysis = Invoke-AILogAnalysis -MaxEvents 100
Write-Host $analysis

# 用法示例 2：命令推荐
$suggestion = Get-AICommandSuggestion -Description '查找过去 24 小时内被修改的大于 100MB 的文件'
Write-Host $suggestion

# 用法示例 3：脚本错误修复
$badScript = @'
Get-Process | Where-Object { $_.WorkingSet -gt 100MB } | Select Name, CPU
'@

$errorMsg = 'Cannot compare "System.Byte[]" to "System.Int32"'

$fix = Repair-AIScriptError -ScriptContent $badScript -ErrorMessage $errorMsg
Write-Host $fix

## 推荐方案

### 方案一：使用 Get-ChildItem
Get-ChildItem -Path C:\ -Recurse -File -ErrorAction SilentlyContinue |
    Where-Object { $_.Length -gt 100MB -and $_.LastWriteTime -gt (Get-Date).AddHours(-24) } |
    Select-Object FullName, @{N='SizeMB';E={[math]::Round($_.Length/1MB,2)}}, LastWriteTime |
    Sort-Object SizeMB -Descending

适用于全盘扫描，结果详细但速度较慢。

using namespace System.Collections.Generic
using namespace System.Text.Json

# MCP 服务端核心类
class McpServer {
    [string] $ServerName
    [string] $Version
    [List[hashtable]] $Tools
    [Dictionary[string, scriptblock]] $Handlers

    McpServer([string]$Name, [string]$Ver) {
        $this.ServerName = $Name
        $this.Version = $Ver
        $this.Tools = [List[hashtable]]::new()
        $this.Handlers = [Dictionary[string, scriptblock]]::new()
    }

    # 注册工具
    [void] RegisterTool(
        [string]$ToolName,
        [string]$Description,
        [hashtable]$Parameters,
        [scriptblock]$Handler
    ) {
        $toolDef = @{
            name        = $ToolName
            description = $Description
            inputSchema = @{
                type       = 'object'
                properties = $Parameters
                required   = @($Parameters.Keys)
            }
        }
        $this.Tools.Add($toolDef)
        $this.Handlers[$ToolName] = $Handler
    }

    # 处理 JSON-RPC 请求
    [string] HandleRequest([string]$JsonLine) {
        $request = $JsonLine | ConvertFrom-Json -AsHashtable
        $response = @{ jsonrpc = '2.0'; id = $request.id }

        switch ($request.method) {
            'initialize' {
                $response.result = @{
                    protocolVersion = '2025-03-26'
                    capabilities    = @{ tools = @{} }
                    serverInfo      = @{
                        name    = $this.ServerName
                        version = $this.Version
                    }
                }
            }
            'tools/list' {
                $response.result = @{ tools = $this.Tools }
            }
            'tools/call' {
                $toolName = $request.params.name
                $args     = $request.params.arguments
                if ($this.Handlers.ContainsKey($toolName)) {
                    $result = & $this.Handlers[$toolName] $args
                    $response.result = @{
                        content = @(
                            @{ type = 'text'; text = ($result | ConvertTo-Json -Depth 5) }
                        )
                    }
                } else {
                    $response.error = @{
                        code    = -32601
                        message = "Tool not found: $toolName"
                    }
                }
            }
            default {
                $response.error = @{
                    code    = -32601
                    message = "Method not found: $($request.method)"
                }
            }
        }
        return $response | ConvertTo-Json -Depth 10 -Compress
    }

    # 启动 stdio 监听循环
    [void] Start() {
        Write-Host "[$($this.ServerName)] MCP Server started on stdio" `
            -ForegroundColor Green
        while ($null -ne ($line = [Console]::In.ReadLine())) {
            $reply = $this.HandleRequest($line)
            [Console]::Out.WriteLine($reply)
            [Console]::Out.Flush()
        }
    }
}

# 创建服务端实例并注册示例工具
$server = [McpServer]::new('ps-mcp-server', '1.0.0')

$server.RegisterTool(
    'get-process-info',
    '获取指定进程的详细信息',
    @{ name = @{ type = 'string'; description = '进程名称' } },
    {
        param($args)
        $procs = Get-Process -Name $args.name -ErrorAction SilentlyContinue
        if ($procs) {
            $procs | Select-Object Name, Id, CPU, WorkingSet64, Path |
                Format-Table -AutoSize | Out-String
        } else {
            "未找到进程: $($args.name)"
        }
    }
)

$server.RegisterTool(
    'calculate',
    '执行数学表达式计算',
    @{
        expression = @{ type = 'string'; description = '数学表达式' }
    },
    {
        param($args)
        $result = Invoke-Expression $args.expression
        "计算结果: $($args.expression) = $result"
    }
)

# 如需 stdio 模式，取消下行注释
# $server.Start()

class McpClient {
    [System.Diagnostics.Process] $Process
    [int] $RequestId = 0
    [hashtable] $ServerCapabilities
    [array] $AvailableTools

    # 启动服务端进程
    [void] Connect([string]$Command, [string[]]$Arguments) {
        $psi = [System.Diagnostics.ProcessStartInfo]::new()
        $psi.FileName            = $Command
        $psi.Arguments           = $Arguments -join ' '
        $psi.UseShellExecute     = $false
        $psi.RedirectStandardInput  = $true
        $psi.RedirectStandardOutput = $true
        $psi.CreateNoWindow      = $true

        $this.Process = [System.Diagnostics.Process]::Start($psi)
        Write-Host "已连接到 MCP 服务端 (PID: $($this.Process.Id))" `
            -ForegroundColor Cyan

        # 发送 initialize 请求
        $initResult = $this.SendRequest('initialize', @{
            protocolVersion = '2025-03-26'
            capabilities    = @{}
            clientInfo      = @{
                name    = 'ps-mcp-client'
                version = '1.0.0'
            }
        })
        $this.ServerCapabilities = $initResult.capabilities
        Write-Host "服务端: $($initResult.serverInfo.name) v$($initResult.serverInfo.version)"

        # 发送 initialized 通知
        $this.SendNotification('notifications/initialized', @{})

        # 缓存可用工具列表
        $toolsResult = $this.SendRequest('tools/list', @{})
        $this.AvailableTools = $toolsResult.tools
        Write-Host "已发现 $($this.AvailableTools.Count) 个工具:"
        foreach ($tool in $this.AvailableTools) {
            Write-Host ("  - {0}: {1}" -f $tool.name, $tool.description) `
                -ForegroundColor Yellow
        }
    }

    # 发送 JSON-RPC 请求并等待响应
    [hashtable] SendRequest([string]$Method, [hashtable]$Params) {
        $this.RequestId++
        $request = @{
            jsonrpc = '2.0'
            id      = $this.RequestId
            method  = $Method
            params  = $Params
        }
        $jsonLine = $request | ConvertTo-Json -Depth 10 -Compress
        $this.Process.StandardInput.WriteLine($jsonLine)
        $this.Process.StandardInput.Flush()

        $responseLine = $this.Process.StandardOutput.ReadLine()
        $response = $responseLine | ConvertFrom-Json -AsHashtable
        if ($response.error) {
            throw "MCP Error [$($response.error.code)]: $($response.error.message)"
        }
        return $response.result
    }

    # 发送 JSON-RPC 通知（无 id，无响应）
    [void] SendNotification([string]$Method, [hashtable]$Params) {
        $notification = @{
            jsonrpc = '2.0'
            method  = $Method
            params  = $Params
        }
        $jsonLine = $notification | ConvertTo-Json -Depth 10 -Compress
        $this.Process.StandardInput.WriteLine($jsonLine)
        $this.Process.StandardInput.Flush()
    }

    # 调用指定工具
    [string] InvokeTool([string]$ToolName, [hashtable]$Arguments) {
        $result = $this.SendRequest('tools/call', @{
            name      = $ToolName
            arguments = $Arguments
        })
        $textParts = $result.content | Where-Object { $_.type -eq 'text' }
        return ($textParts | ForEach-Object { $_.text }) -join "`n"
    }

    # 断开连接
    [void] Disconnect() {
        if ($this.Process -and -not $this.Process.HasExited) {
            $this.Process.StandardInput.Close()
            $this.Process.WaitForExit(5000)
            if (-not $this.Process.HasExited) {
                $this.Process.Kill()
            }
        }
        Write-Host "已断开 MCP 连接" -ForegroundColor DarkGray
    }
}

$sysServer = [McpServer]::new('ps-sys-tools', '1.0.0')

# 工具 1：系统概览
$sysServer.RegisterTool(
    'system-overview',
    '获取操作系统、CPU、内存等系统概览信息',
    @{},
    {
        param($args)
        $os     = [System.Environment]::OSVersion
        $cpu    = Get-CimInstance Win32_Processor | Select-Object -First 1
        $mem    = Get-CimInstance Win32_OperatingSystem
        $totalMem  = [math]::Round($mem.TotalVisibleMemorySize / 1MB, 2)
        $freeMem   = [math]::Round($mem.FreePhysicalMemory / 1MB, 2)
        $usedPct   = [math]::Round(($totalMem - $freeMem) / $totalMem * 100, 1)
        @{
            OS               = "$($os.Platform) $($os.Version)"
            MachineName      = [System.Environment]::MachineName
            Processor        = $cpu.Name
            CPUUtilization   = "$([math]::Round($cpu.LoadPercentage, 1))%"
            TotalMemory_GB   = $totalMem
            FreeMemory_GB    = $freeMem
            MemoryUsage      = "$usedPct%"
            PowerShellVersion= $PSVersionTable.PSVersion.ToString()
            Uptime           = (Get-Date) - $mem.LastBootUpTime | Out-String
        }
    }
)

# 工具 2：搜索文件内容
$sysServer.RegisterTool(
    'search-files',
    '在指定目录中搜索包含指定文本的文件',
    @{
        directory = @{ type = 'string'; description = '搜索目录路径' }
        pattern   = @{ type = 'string'; description = '搜索的文本模式' }
        extension = @{
            type        = 'string'
            description = '文件扩展名过滤（可选）'
        }
    },
    {
        param($args)
        $params = @{
            Path    = $args.directory
            Pattern = $args.pattern
            Recurse = $true
        }
        if ($args.extension) {
            $params.Include = "*.$($args.extension)"
        }
        $results = Select-String @params |
            Select-Object -First 20 |
            ForEach-Object {
                @{
                    File   = $_.Path
                    Line   = $_.LineNumber
                    Match  = $_.Line.Trim()
                }
            }
        if (-not $results) {
            return @{ message = '未找到匹配结果'; matches = @() }
        }
        @{ matches = $results; totalFound = $results.Count }
    }
)

# 工具 3：网络连通性诊断
$sysServer.RegisterTool(
    'network-diag',
    '对指定主机执行网络连通性诊断（Ping + 端口检测）',
    @{
        hostname = @{ type = 'string'; description = '目标主机名或 IP' }
        ports    = @{
            type        = 'string'
            description = '要检测的端口，逗号分隔（如 80,443,22）'
        }
    },
    {
        param($args)
        $results = @{}

        # Ping 测试
        $ping = Test-Connection -TargetName $args.hostname `
            -Count 4 -ErrorAction SilentlyContinue
        if ($ping) {
            $latencies = $ping | ForEach-Object { $_.Latency }
            $results.ping = @{
                success   = $true
                avgMs     = [math]::Round(($latencies | Measure-Object -Average).Average, 2)
                minMs     = ($latencies | Measure-Object -Minimum).Minimum
                maxMs     = ($latencies | Measure-Object -Maximum).Maximum
            }
        } else {
            $results.ping = @{ success = $false; message = '主机不可达' }
        }

        # 端口检测
        $portResults = @{}
        foreach ($port in $args.port -split ',') {
            $port = $port.Trim()
            $tcp = [System.Net.Sockets.TcpClient]::new()
            try {
                $connect = $tcp.ConnectAsync($args.hostname, [int]$port)
                $wait = $connect.Wait(3000)
                $portResults[$port] = if ($wait -and $connect.IsCompletedSuccessfully) {
                    'OPEN'
                } else {
                    'CLOSED/TIMEOUT'
                }
            } finally {
                $tcp.Close()
            }
        }
        $results.ports = $portResults
        $results
    }
)

# 工具 4：服务状态查询
$sysServer.RegisterTool(
    'get-service-status',
    '查询指定服务的运行状态',
    @{
        name = @{ type = 'string'; description = '服务名称（支持通配符）' }
    },
    {
        param($args)
        $services = Get-Service -Name $args.name -ErrorAction SilentlyContinue
        if (-not $services) {
            return @{ message = "未找到服务: $($args.name)" }
        }
        $services | ForEach-Object {
            @{
                Name       = $_.Name
                DisplayName= $_.DisplayName
                Status     = $_.Status.ToString()
                StartType  = $_.StartType.ToString()
            }
        }
    }
)

Write-Host "已注册 $($sysServer.Tools.Count) 个系统管理工具" `
    -ForegroundColor Green
Write-Host "工具列表:"
$sysServer.Tools | ForEach-Object {
    Write-Host ("  {0,-20} {1}" -f $_.name, $_.description)
}

function Invoke-LLMChat {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [array]$Messages,

        [Parameter()]
        [string]$Model = 'gpt-4o-mini',

        [Parameter()]
        [string]$BaseUrl = 'https://api.openai.com/v1',

        [Parameter()]
        [string]$ApiKey = $env:OPENAI_API_KEY,

        [Parameter()]
        [double]$Temperature = 0.3,

        [Parameter()]
        [array]$Tools
    )

    $headers = @{
        'Content-Type'  = 'application/json'
        'Authorization' = "Bearer $ApiKey"
    }

    $body = @{
        model       = $Model
        messages    = $Messages
        temperature = $Temperature
    }

    if ($Tools) {
        $body['tools'] = $Tools
        $body['tool_choice'] = 'auto'
    }

    $uri = "$BaseUrl/chat/completions"
    $response = Invoke-RestMethod -Uri $uri -Method Post -Headers $headers -Body ($body | ConvertTo-Json -Depth 10)

    return $response.choices[0].message
}

# 构建系统提示词，定义 Agent 的角色和行为规范
$systemPrompt = @"
你是一个 PowerShell 运维 Agent。你可以使用提供的工具来执行系统管理任务。
请根据用户的请求，选择合适的工具进行操作。每次只调用一个工具。
如果任务需要多个步骤，请逐步完成。操作完成后请给出简洁的总结。
"@

# 初始化对话历史
$script:conversationHistory = @(
    @{ role = 'system'; content = $systemPrompt }
)