param(
    [ValidateSet("install", "check", "remove")]
    [string] $Mode = "install"
)

$ErrorActionPreference = "Stop"
$root = (Resolve-Path (Join-Path $PSScriptRoot "..")).Path
$client = Join-Path $root "repokey-client.cmd"
$manifestPath = Join-Path $root "register-stack-mcp-config.json"

$servers = @(
    @{ Name = "vitis";          TomlName = "vitis";          Args = @("vitis", "serve") },
    @{ Name = "pl-qor";         TomlName = "pl_qor";         Args = @("pl-qor", "serve") },
    @{ Name = "vivado-context"; TomlName = "vivado_context"; Args = @("context", "serve") }
)

function Get-ConfigCandidates {
    $candidates = @()
    if ($env:APPDATA) {
        $claudePath = Join-Path $env:APPDATA "Claude\claude_desktop_config.json"
        $candidates += [pscustomobject]@{ Name = "Claude Desktop"; Kind = "json"; Path = $claudePath; DetectPath = (Split-Path -Parent $claudePath) }
        $clinePath = Join-Path $env:APPDATA "Code\User\globalStorage\saoudrizwan.claude-dev\settings\cline_mcp_settings.json"
        $candidates += [pscustomobject]@{ Name = "Cline"; Kind = "json"; Path = $clinePath; DetectPath = (Split-Path -Parent (Split-Path -Parent $clinePath)) }
    }
    if ($env:USERPROFILE) {
        $cursorPath = Join-Path $env:USERPROFILE ".cursor\mcp.json"
        $candidates += [pscustomobject]@{ Name = "Cursor"; Kind = "json"; Path = $cursorPath; DetectPath = (Split-Path -Parent $cursorPath) }
        $windsurfPath = Join-Path $env:USERPROFILE ".codeium\windsurf\mcp_config.json"
        $candidates += [pscustomobject]@{ Name = "Windsurf"; Kind = "json"; Path = $windsurfPath; DetectPath = (Split-Path -Parent $windsurfPath) }
        $claudeCodePath = Join-Path $env:USERPROFILE ".claude.json"
        $candidates += [pscustomobject]@{ Name = "Claude Code"; Kind = "json"; Path = $claudeCodePath; DetectPath = (Join-Path $env:USERPROFILE ".claude") }
        $codexPath = Join-Path $env:USERPROFILE ".codex\config.toml"
        $candidates += [pscustomobject]@{ Name = "Codex"; Kind = "toml"; Path = $codexPath; DetectPath = (Join-Path $env:USERPROFILE ".codex") }
    }
    return $candidates
}

function Test-CandidateInstalled {
    param([pscustomobject] $Candidate)
    return (Test-Path -LiteralPath $Candidate.Path) -or (Test-Path -LiteralPath $Candidate.DetectPath)
}

function Test-ServersAvailable {
    return (Test-Path -LiteralPath $client)
}

function Get-ServerContainer {
    param([psobject] $Config)
    foreach ($propertyName in @("mcpServers", "servers", "mcp_servers")) {
        $property = $Config.PSObject.Properties[$propertyName]
        if ($null -ne $property -and $null -ne $property.Value) {
            return [pscustomobject]@{ Name = $propertyName; Value = $property.Value }
        }
    }
    $container = [pscustomobject]@{}
    $Config | Add-Member -NotePropertyName "mcpServers" -NotePropertyValue $container
    return [pscustomobject]@{ Name = "mcpServers"; Value = $container }
}

function Test-SameJsonEntry {
    param([psobject] $Entry, [string[]] $Args)
    if ($null -eq $Entry) { return $false }
    $entryArgs = @($Entry.args | ForEach-Object { [string]$_ })
    return ($entryArgs -join "`n" -eq ($Args -join "`n"))
}

function New-JsonEntry {
    param([string[]] $Args)
    return [pscustomobject]@{
        command = $client
        args = $Args
        optional = $true
    }
}

function Write-Manifest {
    $manifest = [ordered]@{
        name = "repoKeyStackMcps"
        servers = @($servers | ForEach-Object {
            [ordered]@{
                name = $_.Name
                command = $client
                args = $_.Args
                optional = $true
            }
        })
        note = "AI-stack MCP servers, served by the RepoKey client (no Python needed)."
    }
    $manifest | ConvertTo-Json -Depth 8 | Set-Content -Encoding UTF8 -LiteralPath $manifestPath
}

function Backup-Config {
    param([string] $Path)
    $stamp = Get-Date -Format "yyyyMMdd-HHmmss"
    $backup = "$Path.repokey-stack.$stamp.bak"
    Copy-Item -LiteralPath $Path -Destination $backup -Force
    return $backup
}

function Update-JsonConfig {
    param([pscustomobject] $Candidate)
    $path = $Candidate.Path
    $config = $null
    if (Test-Path -LiteralPath $path) {
        try {
            $config = Get-Content -Raw -LiteralPath $path | ConvertFrom-Json
        } catch {
            Write-Host "[WARN] $($Candidate.Name) config is not valid JSON; skipped: $path"
            return "invalid"
        }
    } else {
        $parent = Split-Path -Parent $path
        New-Item -ItemType Directory -Force -Path $parent | Out-Null
        $config = [pscustomobject]@{}
    }

    $containerInfo = Get-ServerContainer $config
    $changed = $false
    foreach ($server in $servers) {
        $existing = $containerInfo.Value.PSObject.Properties[$server.Name]
        if (Test-SameJsonEntry $existing.Value $server.Args) { continue }
        if (-not $changed -and (Test-Path -LiteralPath $path)) {
            $backup = Backup-Config $path
            Write-Host "[OK] Backed up $($Candidate.Name) config: $backup"
        }
        $containerInfo.Value | Add-Member -Force -NotePropertyName $server.Name -NotePropertyValue (New-JsonEntry $server.Args)
        $changed = $true
    }
    if ($changed) {
        $config | ConvertTo-Json -Depth 12 | Set-Content -Encoding UTF8 -LiteralPath $path
        Write-Host "[OK] Registered AI-stack MCPs for $($Candidate.Name)."
    } else {
        Write-Host "[OK] $($Candidate.Name) already has AI-stack MCPs configured."
    }
}

function Get-CodexBlock {
    param([string] $Text, [string] $TomlName)
    $pattern = "(?ms)^\[mcp_servers\.(?:$TomlName|`"$TomlName`")\]\s*$.*?(?=^\[|\z)"
    return [regex]::Match($Text, $pattern).Value
}

function Test-SameCodexBlock {
    param([string] $Block, [string[]] $Args)
    if ([string]::IsNullOrWhiteSpace($Block)) { return $false }
    $needle = "'" + ($Args -join "', '") + "'"
    return ($Block -match [regex]::Escape($needle))
}

function New-CodexBlock {
    param([string] $TomlName, [string[]] $Args)
    $clientPosix = $client.Replace("\\", "/")
    $argList = ($Args | ForEach-Object { "'$_'" }) -join ", "
    $block = "`n`n[mcp_servers.$TomlName]`n"
    $block += "command = 'cmd'`n"
    $block += "args = ['/c', '$clientPosix', $argList]`n"
    $block += "startup_timeout_sec = 120`n"
    return $block
}

function Update-CodexConfig {
    param([pscustomobject] $Candidate)
    $path = $Candidate.Path
    $text = ""
    if (Test-Path -LiteralPath $path) {
        $text = [System.IO.File]::ReadAllText($path)
    } else {
        New-Item -ItemType Directory -Force -Path (Split-Path -Parent $path) | Out-Null
    }

    $changed = $false
    foreach ($server in $servers) {
        $existingBlock = Get-CodexBlock $text $server.TomlName
        if (Test-SameCodexBlock $existingBlock $server.Args) { continue }
        if (-not $changed -and (Test-Path -LiteralPath $path)) {
            $backup = Backup-Config $path
            Write-Host "[OK] Backed up Codex config: $backup"
        }
        if (-not [string]::IsNullOrWhiteSpace($existingBlock)) {
            $text = $text.Replace($existingBlock, "").TrimEnd()
        }
        $text = $text.TrimEnd() + (New-CodexBlock $server.TomlName $server.Args) + [Environment]::NewLine
        $changed = $true
    }
    if ($changed) {
        $utf8 = New-Object System.Text.UTF8Encoding($false)
        [System.IO.File]::WriteAllText($path, $text, $utf8)
        Write-Host "[OK] Registered AI-stack MCPs for Codex."
    } else {
        Write-Host "[OK] Codex already has AI-stack MCPs configured."
    }
}

function Check-CodexConfig {
    param([pscustomobject] $Candidate)
    if (-not (Test-Path -LiteralPath $Candidate.Path)) {
        Write-Host "[INFO] Codex: configuration not created yet."
        return
    }
    $text = [System.IO.File]::ReadAllText($Candidate.Path)
    foreach ($server in $servers) {
        $state = if (Test-SameCodexBlock (Get-CodexBlock $text $server.TomlName) $server.Args) { "registered" } else { "not registered" }
        Write-Host "[INFO] Codex: $($server.Name) MCP $state."
    }
}

function Remove-CodexConfig {
    param([pscustomobject] $Candidate)
    if (-not (Test-Path -LiteralPath $Candidate.Path)) { return }
    $text = [System.IO.File]::ReadAllText($Candidate.Path)
    $backup = $null
    foreach ($server in $servers) {
        $block = Get-CodexBlock $text $server.TomlName
        if ([string]::IsNullOrWhiteSpace($block)) { continue }
        if ($null -eq $backup) { $backup = Backup-Config $Candidate.Path }
        $text = $text.Replace($block, "").TrimEnd()
    }
    if ($null -ne $backup) {
        $utf8 = New-Object System.Text.UTF8Encoding($false)
        [System.IO.File]::WriteAllText($Candidate.Path, $text + [Environment]::NewLine, $utf8)
        Write-Host "[OK] Removed AI-stack MCPs from Codex; backup: $backup"
    }
}

Write-Manifest
Write-Host "RepoKey AI-stack MCP setup"
Write-Host "[OK] Registration manifest: $manifestPath"

if (-not (Test-ServersAvailable)) {
    Write-Host "[INFO] repokey-client.cmd not found; skipping AI-stack MCP registration."
    exit 0
}

$candidates = @(Get-ConfigCandidates | Where-Object { Test-CandidateInstalled $_ })
if ($candidates.Count -eq 0) {
    Write-Host "[INFO] No supported AI client configuration was detected."
    exit 0
}

if ($Mode -eq "check") {
    foreach ($candidate in $candidates) {
        if ($candidate.Kind -eq "toml") { Check-CodexConfig $candidate; continue }
        if (-not (Test-Path -LiteralPath $candidate.Path)) {
            Write-Host "[INFO] $($candidate.Name): configuration not created yet."
            continue
        }
        try {
            $config = Get-Content -Raw -LiteralPath $candidate.Path | ConvertFrom-Json
            $containerInfo = Get-ServerContainer $config
            foreach ($server in $servers) {
                $state = if (Test-SameJsonEntry $containerInfo.Value.PSObject.Properties[$server.Name].Value $server.Args) { "registered" } else { "not registered" }
                Write-Host "[INFO] $($candidate.Name): $($server.Name) MCP $state."
            }
        } catch {
            Write-Host "[WARN] $($candidate.Name): invalid or unsupported JSON config."
        }
    }
    exit 0
}

if ($Mode -eq "remove") {
    foreach ($candidate in $candidates) {
        if ($candidate.Kind -eq "toml") { Remove-CodexConfig $candidate; continue }
        if (-not (Test-Path -LiteralPath $candidate.Path)) { continue }
        try {
            $config = Get-Content -Raw -LiteralPath $candidate.Path | ConvertFrom-Json
            $containerInfo = Get-ServerContainer $config
            $changed = $false
            foreach ($server in $servers) {
                $property = $containerInfo.Value.PSObject.Properties[$server.Name]
                if ($null -ne $property) {
                    if (-not $changed) { $backup = Backup-Config $candidate.Path }
                    $containerInfo.Value.PSObject.Properties.Remove($server.Name)
                    $changed = $true
                }
            }
            if ($changed) {
                $config | ConvertTo-Json -Depth 12 | Set-Content -Encoding UTF8 -LiteralPath $candidate.Path
                Write-Host "[OK] Removed AI-stack MCPs from $($candidate.Name); backup: $backup"
            }
        } catch {
            Write-Host "[WARN] Could not update $($candidate.Name) config."
        }
    }
    exit 0
}

foreach ($candidate in $candidates) {
    if ($candidate.Kind -eq "toml") { [void](Update-CodexConfig $candidate) }
    else { [void](Update-JsonConfig $candidate) }
}
Write-Host "[OK] AI-stack MCP registration finished."
Write-Host "[INFO] Restart the AI client once so it reloads MCP configuration."
