#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""修复quick_login - 使用池中浏览器时直接登录，不尝试加载cookies"""

import re

with open('/www/wwwroot/zsglpt/playwright_automation.py', 'r', encoding='utf-8') as f:
    content = f.read()

# 修复 quick_login 方法 - 当已有浏览器时直接登录
old_quick_login = '''    def quick_login(self, username: str, password: str, remember: bool = True):
        """快速登录 - 优先使用cookies，失败则正常登录"""
        # 尝试使用cookies
        if self.load_cookies(username):
            self.log(f"尝试使用已保存的登录态...")
            if self.check_login_state():
                self.log(f"✓ 登录态有效，跳过登录")
                return {"success": True, "message": "使用已保存的登录态", "used_cookies": True}
            else:
                self.log(f"登录态已失效，重新登录")
                # 关闭当前context，重新登录
                try:
                    if self.context:
                        self.context.close()
                    if self.browser:
                        self.browser.close()
                    if self.playwright:
                        self.playwright.stop()
                except:
                    pass

        # 正常登录
        result = self.login(username, password, remember)

        # 登录成功后保存cookies
        if result.get('success'):
            self.save_cookies(username)
            result['used_cookies'] = False

        return result'''

new_quick_login = '''    def quick_login(self, username: str, password: str, remember: bool = True):
        """快速登录 - 使用池中浏览器时直接登录，否则尝试cookies"""
        # 如果已有浏览器实例（从池中获取），直接使用该浏览器登录
        # 不尝试加载cookies，因为load_cookies会创建新浏览器覆盖池中的
        if self.browser and self.browser.is_connected():
            self.log("使用池中浏览器，直接登录")
            result = self.login(username, password, remember)
            if result.get('success'):
                self.save_cookies(username)
                result['used_cookies'] = False
            return result

        # 无现有浏览器时，尝试使用cookies
        if self.load_cookies(username):
            self.log(f"尝试使用已保存的登录态...")
            if self.check_login_state():
                self.log(f"✓ 登录态有效，跳过登录")
                return {"success": True, "message": "使用已保存的登录态", "used_cookies": True}
            else:
                self.log(f"登录态已失效，重新登录")
                # ��闭当前context，重新登录
                try:
                    if self.context:
                        self.context.close()
                    if self.browser:
                        self.browser.close()
                    if self.playwright:
                        self.playwright.stop()
                except:
                    pass

        # 正常登录
        result = self.login(username, password, remember)

        # 登录成功后保存cookies
        if result.get('success'):
            self.save_cookies(username)
            result['used_cookies'] = False

        return result'''

if old_quick_login in content:
    content = content.replace(old_quick_login, new_quick_login)
    print("OK - quick_login已修复")
else:
    print("WARNING - 未找到old_quick_login")

with open('/www/wwwroot/zsglpt/playwright_automation.py', 'w', encoding='utf-8') as f:
    f.write(content)

print("完成")