#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
AI 定位工作脚本 — 在宿主机上运行

功能：
  通过 Playwright 打开目标网页，提取 DOM 元素信息，
  结合步骤描述匹配最佳定位元素。
  支持调用 Claude Code CLI 做语义分析。

用法：
  # 配合后端（任务文件方式，SSH 自动调用）
  python scripts/ai_locate_worker.py --task-id TASK_ID

  # 独立使用（直接测试）
  python scripts/ai_locate_worker.py --step-desc "输入用户名" --page-url "https://192.168.5.44/login"

工作流程：
  1. 读取任务文件（或命令行参数）
  2. Playwright 打开目标网页，提取所有可交互元素
  3. 关键词算法初步筛选 + Claude Code 精确匹配
  4. 写入结果文件

作者：czj
创建日期：2026-07-29
"""

import argparse
import json
import os
import subprocess
import sys
from pathlib import Path


def extract_interactive_elements(page) -> list:
    """提取页面中所有可交互元素的信息（含 micro-app 微前端容器）"""
    elements = page.evaluate("""
        () => {
            const interactive = [];
            const tags = ['input', 'button', 'a', 'select', 'textarea', 'label'];
            const seen = new Set();

            // 在给定根节点下收集可交互元素
            function collect(root, prefix) {
                if (!root) return;
                root.querySelectorAll(tags.join(',')).forEach(el => {
                    const rect = el.getBoundingClientRect();
                    if (rect.width === 0 || rect.height === 0) return;

                    // 去重：用位置+标签作为 key
                    const key = `${el.tagName}_${Math.round(rect.x)}_${Math.round(rect.y)}_${el.placeholder||''}_${(el.textContent||'').trim().slice(0,20)}`;
                    if (seen.has(key)) return;
                    seen.add(key);

                    const info = {
                        tag: el.tagName.toLowerCase(),
                        type: el.type || '',
                        id: el.id || '',
                        name: el.name || '',
                        class: (typeof el.className === 'string' ? el.className : '') || '',
                        text: (el.textContent || '').trim().slice(0, 100),
                        placeholder: el.placeholder || '',
                        'aria-label': el.getAttribute('aria-label') || '',
                        'data-testid': el.getAttribute('data-testid') || '',
                        href: el.href || '',
                        value: el.value || '',
                        visible: el.offsetParent !== null,
                        scope: prefix || 'main',
                    };
                    interactive.push(info);
                });
            }

            // 1. 主文档
            collect(document, 'main');

            // 2. micro-app 微前端容器（jd-micro / micro-app 框架）
            document.querySelectorAll('micro-app').forEach(ma => {
                // micro-app-body 是实际内容容器
                const body = ma.querySelector('micro-app-body') || ma;
                collect(body, 'micro-app:' + (ma.getAttribute('name') || ''));
                // 也检查 shadowRoot
                if (ma.shadowRoot) {
                    collect(ma.shadowRoot, 'shadow:' + (ma.getAttribute('name') || ''));
                }
            });

            // 3. 普通 iframe
            document.querySelectorAll('iframe').forEach((iframe, idx) => {
                try {
                    const doc = iframe.contentDocument || iframe.contentWindow?.document;
                    if (doc) collect(doc, 'iframe:' + idx);
                } catch (e) { /* cross-origin */ }
            });

            return interactive;
        }
    """)
    return elements


def call_claude_code(step_description: str, elements: list, page_title: str) -> dict:
    """
    调用 Claude Code CLI 分析 DOM 元素，匹配步骤描述

    将元素信息格式化为 prompt，调用 claude -p 非交互模式
    获取 Claude 的结构化分析结果
    """
    element_summary = json.dumps(elements[:50], ensure_ascii=False, indent=2)

    prompt = f"""你是一个UI自动化测试的定位助手。请根据步骤描述，从网页元素列表中找出最匹配的元素。

步骤描述：{step_description}

页面标题：{page_title}

可交互元素列表（JSON格式）：
{element_summary}

请分析以上元素，找出与步骤描述最匹配的元素，并输出以下JSON格式（不要输出其他内容）：
{{
  "locator_type": "css 或 xpath 或 text 或 id",
  "locator_value": "定位值",
  "confidence": 0.0-1.0的小数表示置信度,
  "reason": "简要说明为什么匹配这个元素"
}}

如果找不到匹配的元素，输出：
{{"locator_type": "", "locator_value": "", "confidence": 0, "reason": "未找到匹配元素"}}
"""

    try:
        result = subprocess.run(
            ['claude', '-p', prompt, '--print'],
            capture_output=True,
            text=True,
            timeout=30,
            env={**os.environ, 'CLAUDE_CODE_HEADLESS': '1'}
        )

        if result.returncode == 0:
            output = result.stdout.strip()
            # 提取 JSON 部分
            json_start = output.find('{')
            json_end = output.rfind('}') + 1
            if json_start >= 0 and json_end > json_start:
                json_str = output[json_start:json_end]
                return json.loads(json_str)

        return {"locator_type": "", "locator_value": "", "confidence": 0, "reason": f"Claude调用失败: {result.stderr[:200]}"}

    except subprocess.TimeoutExpired:
        return {"locator_type": "", "locator_value": "", "confidence": 0, "reason": "Claude调用超时"}
    except FileNotFoundError:
        # Claude Code CLI 未安装，降级为纯算法匹配
        return {"locator_type": "", "locator_value": "", "confidence": 0, "reason": "Claude CLI 未安装，使用算法匹配"}
    except Exception as e:
        return {"locator_type": "", "locator_value": "", "confidence": 0, "reason": f"异常: {str(e)}"}


def build_locator_candidates(step_description: str, elements: list, page_title: str) -> list:
    """
    通过算法 + Claude Code 分析，生成定位候选列表

    先用关键词算法做初步筛选，置信度不足时调用 Claude Code 做精确匹配
    """
    desc_lower = step_description.lower()
    candidates = []

    # 1. 关键词算法初步筛选
    for el in elements:
        score = 0.0
        reasons = []

        # id 精确匹配
        if el['id'] and (el['id'].lower() in desc_lower or desc_lower in el['id'].lower()):
            score += 0.4
            reasons.append(f"id匹配: {el['id']}")

        # placeholder 匹配
        if el['placeholder']:
            for kw in desc_lower.split():
                if len(kw) > 1 and kw in el['placeholder'].lower():
                    score += 0.3
                    reasons.append(f"placeholder匹配: {el['placeholder']}")
                    break

        # aria-label 匹配
        if el['aria-label']:
            for kw in desc_lower.split():
                if len(kw) > 1 and kw in el['aria-label'].lower():
                    score += 0.3
                    break

        # 文本内容匹配
        if el['text']:
            for kw in desc_lower.split():
                if len(kw) > 1 and kw in el['text'].lower():
                    score += 0.25
                    reasons.append(f"文本匹配: {el['text']}")
                    break

        # name 属性匹配
        if el['name'] and (el['name'].lower() in desc_lower or desc_lower in el['name'].lower()):
            score += 0.2

        # data-testid 匹配
        if el['data-testid'] and el['data-testid'].lower() in desc_lower:
            score += 0.3

        # 语义推断：输入类动作 → input/textarea 优先
        if any(kw in desc_lower for kw in ['输入', '填写', '键入', '填入']):
            if el['tag'] in ['input', 'textarea']:
                score += 0.15

        # 语义推断：点击类动作 → button/a 优先
        if any(kw in desc_lower for kw in ['点击', '按下', '选择', '单击']):
            if el['tag'] in ['button', 'a']:
                score += 0.15

        if score > 0.15:
            locator = _build_locator(step_description, el)
            candidates.append({
                **locator,
                'confidence': round(score, 2),
                'element_info': el,
                'reasons': reasons,
            })

    # 2. 如果最高置信度不足 0.6，调用 Claude Code 做精确分析
    needs_claude = not candidates or candidates[0]['confidence'] < 0.6
    if needs_claude:
        claude_result = call_claude_code(step_description, elements, page_title)
        if claude_result.get('locator_type') and claude_result.get('confidence', 0) > 0.5:
            claude_match = {
                'locator_type': claude_result['locator_type'],
                'locator_value': claude_result['locator_value'],
                'confidence': claude_result.get('confidence', 0.8),
                'element_info': {},
                'reasons': [claude_result.get('reason', 'Claude Code 分析匹配')],
            }
            # 检查是否与已有候选重复
            is_duplicate = any(
                c['locator_value'] == claude_match['locator_value']
                for c in candidates
            )
            if not is_duplicate:
                candidates.insert(0, claude_match)

    # 按置信度降序排列
    candidates.sort(key=lambda x: x['confidence'], reverse=True)
    return candidates[:5]


def _build_locator(description: str, el: dict) -> dict:
    """根据元素信息生成最佳定位方式"""
    if el['id']:
        return {'locator_type': 'css', 'locator_value': f"#{el['id']}"}
    if el['data-testid']:
        return {'locator_type': 'css', 'locator_value': f"[data-testid='{el['data-testid']}']"}
    if el['aria-label']:
        return {'locator_type': 'css', 'locator_value': f"[aria-label='{el['aria-label']}']"}
    if el['name']:
        return {'locator_type': 'css', 'locator_value': f"{el['tag']}[name='{el['name']}']"}
    if el['placeholder']:
        return {'locator_type': 'css', 'locator_value': f"[placeholder='{el['placeholder']}']"}
    if el['text'] and el['tag'] in ['button', 'a']:
        return {'locator_type': 'text', 'locator_value': el['text']}
    if el['id']:
        return {'locator_type': 'xpath', 'locator_value': f"//{el['tag']}[@id='{el['id']}']"}
    return {'locator_type': 'css', 'locator_value': ''}


def run_locate(step_description: str, page_url: str, expected: str = "") -> dict:
    """执行定位——主流程"""
    print(f"[AI定位] 步骤描述: {step_description}")
    print(f"[AI定位] 目标页面: {page_url}")

    try:
        from playwright.sync_api import sync_playwright
    except ImportError:
        print("[AI定位] ❌ 未安装 Playwright，请执行: pip install playwright && playwright install chromium")
        return {"success": False, "candidates": [], "message": "Playwright 未安装"}

    with sync_playwright() as p:
        browser = p.chromium.launch(
            headless=True,
            args=[
                '--disable-blink-features=AutomationControlled',
                '--disable-automation',
                '--no-sandbox',
                '--disable-dev-shm-usage',
                '--disable-gpu',
            ]
        )
        # 内网自签证书需忽略 HTTPS 错误
        context = browser.new_context(
            user_agent='Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 Chrome/126.0.0.0',
            locale='zh-CN',
            timezone_id='Asia/Shanghai',
            ignore_https_errors=True,
        )
        page = context.new_page()
        page.add_init_script("""
            Object.defineProperty(navigator, 'webdriver', { get: () => undefined });
        """)

        # 打开目标页面
        # SPA / 微前端可能需要额外等待渲染
        page.goto(page_url, wait_until='networkidle', timeout=30000)
        page.wait_for_timeout(2000)  # 等待 micro-app 等微前端框架渲染
        page_title = page.title()
        print(f"[AI定位] 页面加载完成: {page_title} (url={page.url})")

        # 提取可交互元素
        elements = extract_interactive_elements(page)
        print(f"[AI定位] 提取到 {len(elements)} 个可交互元素")

        browser.close()

    # 匹配元素
    candidates = build_locator_candidates(step_description, elements, page_title)
    print(f"[AI定位] 匹配到 {len(candidates)} 个候选元素")

    result = {
        'success': len(candidates) > 0,
        'candidates': candidates,
        'message': f"定位成功，找到 {len(candidates)} 个候选" if candidates else "未找到匹配元素",
    }

    if candidates:
        best = candidates[0]
        print(f"\n[AI定位] 最佳匹配:")
        print(f"  定位类型: {best['locator_type']}")
        print(f"  定位值: {best['locator_value']}")
        print(f"  置信度: {best['confidence']}")
        if best['element_info']:
            info = best['element_info']
            print(f"  元素标签: {info.get('tag', '')}")
            print(f"  元素属性: id={info.get('id','')}, name={info.get('name','')}, placeholder={info.get('placeholder','')}")
    else:
        print(f"\n[AI定位] 未找到匹配元素")

    return result


def main():
    parser = argparse.ArgumentParser(description='AI 定位工作脚本（服务器端）')
    parser.add_argument('--task-id', help='任务ID（从后端获取）')
    parser.add_argument('--step-desc', help='步骤描述（直接传入）')
    parser.add_argument('--page-url', help='目标网页 URL')
    parser.add_argument('--expected', default='', help='预期结果')
    args = parser.parse_args()

    if args.task_id:
        # 任务文件方式：配合后端使用
        # 优先使用与容器共享的 data 目录（/data/third_party/plat-auto-test/data）
        # 容器内 TASK_DIR 映射到宿主机 DEPLOY_DIR/data/ai_locate_tasks
        deploy_data_dir = Path('/data/third_party/plat-auto-test/data/ai_locate_tasks')
        script_data_dir = Path(__file__).parent.parent / 'data' / 'ai_locate_tasks'
        task_dir = deploy_data_dir if deploy_data_dir.exists() else script_data_dir
        task_file = task_dir / f"{args.task_id}.json"
        result_file = task_dir / f"{args.task_id}_result.json"

        if not task_file.exists():
            print(f"❌ 任务文件不存在: {task_file}")
            return

        with open(task_file, 'r', encoding='utf-8') as f:
            task = json.load(f)

        result = run_locate(
            step_description=task['step_description'],
            page_url=task['page_url'],
            expected=task.get('expected', ''),
        )

        result['task_id'] = args.task_id
        result['status'] = 'completed'
        task_dir.mkdir(parents=True, exist_ok=True)
        with open(result_file, 'w', encoding='utf-8') as f:
            json.dump(result, f, ensure_ascii=False, indent=2)

        print(f"\n✅ 定位结果已写入: {result_file}")

    else:
        # 直接运行
        result = run_locate(
            step_description=args.step_desc,
            page_url=args.page_url,
            expected=args.expected,
        )
        print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == '__main__':
    main()