#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
金山文档上传测试 - 完整自动登录版本
自动处理：登录并加入编译 → 扫码 → 确认登录
"""

import os
import sys
import time
import base64
from datetime import datetime
from io import BytesIO
from PIL import Image

# 添加项目路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

try:
    from playwright.sync_api import sync_playwright
except ImportError:
    print("错误: 需要安装 playwright")
    print("请运行: pip install playwright")
    sys.exit(1)


def log(message, level='INFO'):
    """日志输出"""
    timestamp = datetime.now().strftime("%H:%M:%S")
    print(f"[{timestamp}] {level}: {message}")


def pause(msg="按Enter键继续..."):
    """等待用户按键"""
    input(f"\n{msg}")


def ask_yes_no(question, default='n'):
    """询问用户是/否问题"""
    if default == 'y':
        prompt = f"{question} (Y/n): "
    else:
        prompt = f"{question} (y/N): "

    answer = input(prompt).strip().lower()
    if not answer:
        answer = default
    return answer == 'y'


def save_qr_code(qr_image_bytes, filename="qr_code.png"):
    """保存二维码图片"""
    try:
        with open(filename, 'wb') as f:
            f.write(qr_image_bytes)
        log(f"[OK] 二维码已保存到: {filename}", 'SUCCESS')
        return filename
    except Exception as e:
        log(f"✗ 保存二维码失败: {str(e)}", 'ERROR')
        return None


def click_login_join_button(page):
    """点击'登录并加入编辑'按钮"""
    log("查找'登录并加入编辑'按钮...", 'INFO')

    # 多种可能的按钮选择器
    login_selectors = [
        "text=登录并加入编辑",
        "text=登录并加入编译",
        "button:has-text('登录')",
        "text=立即登录",
        "[class*='login']",
        "[id*='login']"
    ]

    for selector in login_selectors:
        try:
            button = page.locator(selector).first
            if button.is_visible(timeout=3000):
                log(f"[OK] 找到登录按钮: {selector}", 'SUCCESS')
                button.click()
                log("[OK] 已点击登录按钮", 'SUCCESS')
                return True
        except Exception:
            continue

    log("✗ 未找到登录按钮", 'ERROR')
    return False


def wait_for_qr_code(page, timeout=30):
    """等待二维码出现"""
    log("等待二维码加载...", 'INFO')

    start_time = time.time()
    while time.time() - start_time < timeout:
        try:
            # 查找二维码元素
            qr_selectors = [
                "canvas",
                "img[src*='qr']",
                "img[alt*='二维码']",
                "[class*='qr']",
                "[id*='qr']",
                "div[class*='qrcode']",
                "img[src*='wechat']"
            ]

            for selector in qr_selectors:
                try:
                    elements = page.query_selector_all(selector)
                    for i, element in enumerate(elements):
                        try:
                            # 尝试截图
                            screenshot = element.screenshot()
                            if len(screenshot) > 500:  # 足够大的图片
                                filename = f"qr_code_{i}.png"
                                save_qr_code(screenshot, filename)
                                log(f"[OK] 找到二维码元素: {selector}[{i}]", 'SUCCESS')
                                return True
                        except Exception:
                            continue
                except Exception:
                    continue

            time.sleep(1)

        except Exception as e:
            log(f"检查二维码时出错: {str(e)}", 'WARNING')
            time.sleep(1)

    return False


def wait_for_confirm_login(page, timeout=120):
    """等待'确认登录'按钮出现并点击"""
    log("等待用户扫码...", 'INFO')
    log("请使用手机微信扫描二维码", 'INFO')
    log("扫码完成后，程序会自动检测并点击'确认登录'", 'INFO')

    start_time = time.time()
    check_interval = 2  # 每2秒检查一次

    while time.time() - start_time < timeout:
        try:
            # 查找确认登录按钮
            confirm_selectors = [
                "text=确认登录",
                "text=确认登陆",
                "button:has-text('确认')",
                "text=登录",
                "[class*='confirm']",
                "[id*='confirm']"
            ]

            for selector in confirm_selectors:
                try:
                    button = page.locator(selector).first
                    if button.is_visible(timeout=1000):
                        log(f"[OK] 找到确认按钮: {selector}", 'SUCCESS')
                        button.click()
                        log("[OK] 已点击确认登录按钮", 'SUCCESS')
                        return True
                except Exception:
                    continue

            # 如果没找到按钮，显示等待信息
            elapsed = int(time.time() - start_time)
            if elapsed % 10 == 0:  # 每10秒显示一次
                log(f"等待中... ({elapsed}秒)", 'INFO')

            time.sleep(check_interval)

        except Exception as e:
            log(f"检查确认按钮时出错: {str(e)}", 'WARNING')
            time.sleep(check_interval)

    return False


def wait_for_document_loaded(page, timeout=30):
    """等待文档页面加载完成"""
    log("等待文档页面加载...", 'INFO')

    start_time = time.time()
    while time.time() - start_time < timeout:
        try:
            current_url = page.url
            log(f"当前URL: {current_url}", 'INFO')

            # 检查是否进入文档页面
            if "kdocs.cn" in current_url and "/spreadsheet/" in current_url:
                log("[OK] 已进入文档页面", 'SUCCESS')
                return True

            # 检查表格元素
            try:
                canvas_count = page.locator("canvas").count()
                if canvas_count > 0:
                    log(f"[OK] 检测到 {canvas_count} 个表格元素", 'SUCCESS')
                    return True
            except:
                pass

            time.sleep(2)

        except Exception as e:
            log(f"检查页面状态时出错: {str(e)}", 'WARNING')
            time.sleep(2)

    return False


def main():
    """主函数"""
    print("=" * 70)
    print("[LOCK] 金山文档上传测试 - 完整自动登录版本")
    print("=" * 70)
    print()
    print("特点:")
    print("  [OK] 自动点击'登录并加入编译'")
    print("  [OK] 自动捕获二维码")
    print("  [OK] 自动等待并点击'确认登录'")
    print("  [OK] 自动检测文档加载")
    print()

    # 配置
    doc_url = input("请输入金山文档URL (或按Enter使用默认): ").strip()
    if not doc_url:
        doc_url = "https://kdocs.cn/l/cpwEOo5ynKX4"

    print(f"\n使用URL: {doc_url}")
    print()

    if not ask_yes_no("确认开始测试?"):
        print("测试已取消")
        return

    print("\n" + "=" * 70)
    print("开始测试流程")
    print("=" * 70)

    playwright = None
    browser = None
    context = None
    page = None

    try:
        # ===== 步骤1: 启动浏览器 =====
        print("\n" + "=" * 50)
        print("步骤1: 启动浏览器")
        print("=" * 50)

        log("正在启动Playwright...", 'INFO')
        playwright = sync_playwright().start()
        log("[OK] Playwright启动成功", 'SUCCESS')

        log("正在启动浏览器...", 'INFO')
        browser = playwright.chromium.launch(headless=False)
        log("[OK] 浏览器启动成功", 'SUCCESS')

        log("正在创建上下文...", 'INFO')
        context = browser.new_context()
        log("[OK] 上下文创建成功", 'SUCCESS')

        log("正在创建页面...", 'INFO')
        page = context.new_page()
        page.set_default_timeout(30000)
        log("[OK] 页面创建成功", 'SUCCESS')

        pause("浏览器已启动，请观察浏览器窗口")

        log("额外等待5秒确保浏览器完全就绪...", 'INFO')
        time.sleep(5)

        # ===== 步骤2: 打开文档页面 =====
        print("\n" + "=" * 50)
        print("步骤2: 打开文档页面")
        print("=" * 50)

        log(f"正在导航到: {doc_url}", 'INFO')
        page.goto(doc_url, wait_until='domcontentloaded')
        log("[OK] 页面导航完成", 'SUCCESS')

        log("等待8秒让页面完全加载...", 'INFO')
        time.sleep(8)

        current_url = page.url
        log(f"当前URL: {current_url}", 'INFO')

        # ===== 步骤3: 自动点击登录按钮 =====
        print("\n" + "=" * 50)
        print("步骤3: 点击登录按钮")
        print("=" * 50)

        log("检测页面状态...", 'INFO')
        log("等待页面元素完全加载...", 'INFO')

        # 额外的等待确保页面完全加载
        log("额外等待5秒确保页面完全加载...", 'INFO')
        time.sleep(5)

        # 尝试等待特定元素出现
        try:
            page.wait_for_selector("text=登录并加入", timeout=15000)
            log("[OK] 检测到'登录并加入编辑'页面", 'SUCCESS')
            login_button_found = True
        except:
            log("⚠ 未检测到登录按钮，继续等待...", 'WARNING')
            time.sleep(5)
            login_button_found = False

        # 最终检测页面内容
        page_content = page.content()
        if "登录并加入" in page_content:
            log("[OK] 检测到'登录并加入编辑'页面", 'SUCCESS')
            login_button_found = True
        else:
            log("⚠ 未检测到'登录并加入编辑'页面", 'WARNING')
            login_button_found = False

        # 执行点击操作
        if login_button_found:
            if click_login_join_button(page):
                log("[OK] 已点击登录按钮，等待跳转到扫码页面...", 'SUCCESS')
                time.sleep(5)  # 增加等待时间
            else:
                log("✗ 点击登录按钮失败", 'ERROR')
                return
        else:
            # 检查是否已经直接进入登录页面
            if "login" in page.url.lower() or "account" in page.url.lower():
                log("[OK] 已直接进入登录页面", 'SUCCESS')
            else:
                log("⚠ 页面状态不明确，请手动检查浏览器窗口", 'WARNING')

        # ===== 步骤4: 等待二维码 =====
        print("\n" + "=" * 50)
        print("步骤4: 等待二维码")
        print("=" * 50)

        if wait_for_qr_code(page, timeout=90):
            log("[OK] 二维码加载完成", 'SUCCESS')
        else:
            log("⚠ 未检测到二维码，可能页面结构有变化", 'WARNING')

        # ===== 步骤5: 等待确认登录 =====
        print("\n" + "=" * 50)
        print("步骤5: 等待确认登录")
        print("=" * 50)

        log("扫码流程:", 'INFO')
        log("1. 请使用手机微信扫描二维码", 'INFO')
        log("2. 扫码后点击'确认登录'", 'INFO')
        log("3. 程序会自动检测并处理", 'INFO')

        if wait_for_confirm_login(page, timeout=180):
            log("[OK] 登录确认完成", 'SUCCESS')
        else:
            log("⚠ 未检测到确认登录操作", 'WARNING')

        # ===== 步骤6: 等待文档加载 =====
        print("\n" + "=" * 50)
        print("步骤6: 等待文档加载")
        print("=" * 50)

        if wait_for_document_loaded(page, timeout=60):
            log("[OK] 文档页面加载完成", 'SUCCESS')

            # 验证表格元素
            try:
                canvas_count = page.locator("canvas").count()
                log(f"[OK] 检测到 {canvas_count} 个表格元素", 'SUCCESS')

                # 尝试读取名称框
                try:
                    name_box = page.locator("input.edit-box").first
                    if name_box.is_visible():
                        value = name_box.input_value()
                        log(f"[OK] 名称框可见，当前值: '{value}'", 'SUCCESS')
                except:
                    pass

            except Exception as e:
                log(f"检查表格元素时出错: {str(e)}", 'WARNING')
        else:
            log("⚠ 文档页面加载超时", 'WARNING')

        # ===== 步骤7: 表格功能测试 =====
        print("\n" + "=" * 50)
        print("步骤7: 表格功能测试")
        print("=" * 50)

        # 测试搜索功能
        test_name = input("请输入要搜索的姓名 (默认: 张三): ").strip()
        if not test_name:
            test_name = "张三"

        log(f"搜索姓名: {test_name}", 'INFO')

        try:
            page.keyboard.press("Control+f")
            time.sleep(0.5)

            page.keyboard.type(test_name)
            time.sleep(0.3)

            page.keyboard.press("Enter")
            time.sleep(1)

            page.keyboard.press("Escape")
            time.sleep(0.3)

            log("[OK] 搜索测试完成", 'SUCCESS')
            log("请查看浏览器窗口，检查搜索结果", 'INFO')

        except Exception as e:
            log(f"✗ 搜索测试失败: {str(e)}", 'ERROR')

        pause("搜索测试完成")

        # ===== 步骤8: 图片上传测试 =====
        print("\n" + "=" * 50)
        print("步骤8: 图片上传测试")
        print("=" * 50)

        if ask_yes_no("是否进行图片上传测试?"):
            image_path = input("请输入测试图片的完整路径: ").strip()

            if not image_path or not os.path.exists(image_path):
                log("图片文件不存在，跳过上传测试", 'WARNING')
            else:
                log(f"选中的图片: {image_path}", 'INFO')

                try:
                    # 导航到D3
                    name_box = page.locator("input.edit-box").first
                    name_box.click()
                    name_box.fill("D3")
                    name_box.press("Enter")
                    time.sleep(0.5)

                    log("[OK] 已导航到D3单元格")

                    # 点击插入
                    insert_btn = page.locator("text=插入").first
                    insert_btn.click()
                    time.sleep(0.5)

                    log("[OK] 已点击插入按钮")

                    # 点击图片
                    image_btn = page.locator("text=图片").first
                    image_btn.click()
                    time.sleep(0.5)

                    log("[OK] 已点击图片按钮")

                    # 选择本地
                    local_option = page.locator("text=本地").first
                    local_option.click()

                    log("[OK] 已选择本地图片")

                    # 上传文件
                    with page.expect_file_chooser() as fc_info:
                        pass

                    file_chooser = fc_info.value
                    file_chooser.set_files(image_path)

                    log("[OK] 文件上传命令已发送")

                    time.sleep(3)

                    log("[OK] 图片上传测试完成", 'SUCCESS')

                except Exception as e:
                    log(f"✗ 图片上传测试失败: {str(e)}", 'ERROR')

        pause("所有测试完成")

        # ===== 测试完成 =====
        print("\n" + "=" * 70)
        log("🎉 所有测试完成!", 'SUCCESS')
        print("=" * 70)

    except KeyboardInterrupt:
        print("\n")
        log("测试被用户中断", 'WARNING')
    except Exception as e:
        print("\n")
        log(f"测试过程中出现错误: {str(e)}", 'ERROR')
        import traceback
        traceback.print_exc()
    finally:
        # 清理资源
        print("\n" + "=" * 70)
        print("清理资源...")
        print("=" * 70)

        try:
            if page:
                page.close()
                log("[OK] 页面已关闭", 'SUCCESS')
        except:
            pass

        try:
            if context:
                context.close()
                log("[OK] 上下文已关闭", 'SUCCESS')
        except:
            pass

        try:
            if browser:
                browser.close()
                log("[OK] 浏览器已关闭", 'SUCCESS')
        except:
            pass

        try:
            if playwright:
                playwright.stop()
                log("[OK] Playwright已停止", 'SUCCESS')
        except:
            pass

        log("测试结束", 'SUCCESS')
        print("=" * 70)


if __name__ == "__main__":
    main()