# -*- coding: utf-8 -*-
"""两件事（用户要求）：

1. 跑全量时开的窗口，跑完要**自己关上**
   测试会弹真实的 Tk 窗口（首次设置窗口、投图窗口、屏幕共享提示等）。
   正常路径下测试自己会关，但一旦某个用例失败/超时，
   进程挂着，窗口就留在屏幕上不走了。
   做法：在 verify_all 结尾加一个收尾步骤，按进程把窗口关掉
   （先 CloseMainWindow 好声好气关，再强杀兜底）。

2. 把"锁屏检查"从全量里删掉
   之前我在 test_agent_protocol.py 里加了 desktop_locked()，
   锁屏时打印一大段提示。用户要求以后都别做这个检查。
"""
from __future__ import annotations

import ast
import re
import sys
from pathlib import Path

ROOT = Path(r"C:\Users\shu\Desktop\deepseekagent\LanRemote")
VERIFY = ROOT / "_tests" / "verify_all.ps1"
PROTOCOL = ROOT / "_tests" / "test_agent_protocol.py"

# ===================== 1. 删掉锁屏检查 =====================

LOCK_BLOCK = '''def desktop_locked() -> bool:
    """屏幕锁了没。锁屏时 Windows 不给抓屏，相关断言必然失败 ——
    那是环境，不是产品问题，要单独说清楚。"""
    try:
        from PIL import ImageGrab
        ImageGrab.grab()
        return False
    except Exception:
        return True


def main() -> int:
    locked = desktop_locked()
    if locked:
        print()
        print("  ！检测到桌面锁了（锁屏/无桌面会话）。")
        print("    截屏和镜像相关的断言会失败，这是环境限制，不是产品问题。")
        print("    解锁屏幕后重跑即可。")
        print()'''

NEW_MAIN = '''def main() -> int:'''

# ===================== 2. 收尾关窗口 =====================

CLEANUP_FUNC = '''function Close-TestWindows {
    # 测试会弹真实的 Tk 窗口（首次设置、投图、屏幕共享提示…）。
    # 正常路径下测试自己会关；但用例失败/超时时进程还挂着，
    # 窗口就留在屏幕上不走了。这里统一收尾。
    $closed = 0
    Get-Process python, pythonw -ErrorAction SilentlyContinue |
        Where-Object { $_.Id -ne $PID } |
        ForEach-Object {
            $line = (Get-CimInstance Win32_Process -Filter "ProcessId=$($_.Id)" -ErrorAction SilentlyContinue).CommandLine
            if ($line -and ($line -match '_tests' -or $line -match 'run_agent' `
                    -or $line -match 'relay\\.py')) {
                # 先好好关主窗口
                try { $_.CloseMainWindow() | Out-Null } catch { }
                Start-Sleep -Milliseconds 250
                if (-not $_.HasExited) {
                    Stop-Process -Id $_.Id -Force -ErrorAction SilentlyContinue
                }
                $closed++
            }
        }

    # 被控端自己（测试起过 ykxmz.exe）
    Get-Process ykxmz -ErrorAction SilentlyContinue | ForEach-Object {
        try { $_.CloseMainWindow() | Out-Null } catch { }
        Start-Sleep -Milliseconds 250
        if (-not $_.HasExited) {
            Stop-Process -Id $_.Id -Force -ErrorAction SilentlyContinue
        }
        $closed++
    }

    # Java 测试进程
    Get-Process java -ErrorAction SilentlyContinue | ForEach-Object {
        $line = (Get-CimInstance Win32_Process -Filter "ProcessId=$($_.Id)" -ErrorAction SilentlyContinue).CommandLine
        if ($line -and ($line -match 'JavaClientTest' -or $line -match 'JavaRelayTest')) {
            Stop-Process -Id $_.Id -Force -ErrorAction SilentlyContinue
            $closed++
        }
    }

    if ($closed -gt 0) {
        Write-Host "  关了 $closed 个测试进程/窗口" -ForegroundColor DarkGray
    }
    Start-Sleep -Milliseconds 400
}

'''

SUMMARY_ANCHOR = """Banner "验证汇总\""""


def main() -> int:
    ok = True

    # ---- 1. 删锁屏检查 ----
    text = PROTOCOL.read_text(encoding="utf-8")
    if "desktop_locked" in text:
        if LOCK_BLOCK in text:
            text = text.replace(LOCK_BLOCK, NEW_MAIN, 1)
            PROTOCOL.write_text(text, encoding="utf-8", newline="\n")
            print("  已删掉 test_agent_protocol.py 里的锁屏检查")
        else:
            # 锚点不完全一致时，用正则兜底
            text = re.sub(
                r"def desktop_locked\(\).*?\n\n\n(?=def main)", "", text,
                count=1, flags=re.S)
            text = re.sub(
                r"def main\(\) -> int:\n    locked = desktop_locked\(\)"
                r".*?\n\n", "def main() -> int:\n", text, count=1, flags=re.S)
            if "desktop_locked" in text:
                print("  [警告] 锁屏检查没删干净，还有残留")
                ok = False
            else:
                PROTOCOL.write_text(text, encoding="utf-8", newline="\n")
                print("  已删掉锁屏检查（正则兜底）")
    else:
        print("  锁屏检查已经没有了")

    # ---- 2. verify_all 加收尾 ----
    raw = VERIFY.read_bytes()
    had_bom = raw.startswith(b"\xef\xbb\xbf")
    vtext = raw.decode("utf-8-sig")

    if "Close-TestWindows" in vtext:
        print("  verify_all 已经有 Close-TestWindows")
    else:
        # 函数插在 Stop-Strays 前面
        anchor = "function Stop-Strays {"
        if anchor in vtext:
            vtext = vtext.replace(anchor, CLEANUP_FUNC + anchor, 1)
            print("  加了 Close-TestWindows 函数")
        else:
            print("  [失败] 找不到插入位置")
            ok = False

        # 在汇总之前调用
        if SUMMARY_ANCHOR in vtext:
            vtext = vtext.replace(
                SUMMARY_ANCHOR,
                '# 收尾：把测试过程中弹出来的窗口关掉（用户要求）\n'
                'Close-TestWindows\n\n' + SUMMARY_ANCHOR, 1)
            print("  在汇总前调用了 Close-TestWindows")
        else:
            print("  [失败] 找不到汇总锚点")
            ok = False

        VERIFY.write_bytes((b"\xef\xbb\xbf" if had_bom else b"")
                           + vtext.replace("\r\n", "\n").encode("utf-8"))

    ast.parse(PROTOCOL.read_text(encoding="utf-8"))
    print("  test_agent_protocol.py 语法 OK")
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())