#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
模块名称：ai_locator.py
模块描述：AI 定位路由 — 通过宿主机 Claude Code 获取元素定位信息

功能：
1. POST /api/ai/locate — 创建定位任务，SSH 调用宿主机执行定位脚本
2. GET /api/ai/locate/{task_id}/result — 查询定位结果

工作流程：
1. 前端请求定位 → 后端创建任务文件 (data/ai_locate_tasks/{task_id}.json)
2. 后端通过 SSH 调用宿主机脚本：python scripts/ai_locate_worker.py --task-id TASK_ID
3. 宿主机脚本：Playwright 开网页 → 提取 DOM → Claude Code 分析匹配 → 写结果文件
4. 前端轮询获取结果

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

import json
import logging
import os
import subprocess
import uuid
from pathlib import Path
from typing import Optional, List, Dict, Any

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)

router = APIRouter()

# ==================== 目录配置 ====================

# 任务文件目录：容器内 /app/data 映射到宿主机 /data/third_party/plat-auto-test/data
# 优先使用共享 data 目录，保证容器与宿主机脚本读写同一位置
_SHARED_DATA = Path("/app/data/ai_locate_tasks")
_LOCAL_DATA = Path(__file__).parent.parent.parent / "data" / "ai_locate_tasks"
TASK_DIR = _SHARED_DATA if _SHARED_DATA.parent.exists() else _LOCAL_DATA

# 部署目录（服务器上）
DEPLOY_DIR = "/data/third_party/plat-auto-test"

# 宿主机 SSH 配置
HOST_USER = "ubains"
HOST = "192.168.5.60"

# ==================== Schema 定义 ====================


class LocateRequest(BaseModel):
    """定位请求"""

    step_description: str = Field(..., description="步骤描述，如「输入用户名」")
    page_url: str = Field(..., description="目标网页 URL")
    expected: str = Field(default="", description="预期结果（可选）")


class LocateCandidate(BaseModel):
    """定位候选"""

    locator_type: str = Field(..., description="定位类型: css/xpath/text/id/name")
    locator_value: str = Field(..., description="定位值")
    confidence: float = Field(..., ge=0, le=1, description="置信度 0-1")
    element_info: Dict[str, Any] = Field(default_factory=dict, description="元素信息")


class LocateTaskResponse(BaseModel):
    """创建任务响应"""

    task_id: str = Field(..., description="任务ID")
    status: str = Field(..., description="任务状态: pending/completed/failed")
    message: str = Field(default="", description="提示信息")


class LocateResultResponse(BaseModel):
    """查询结果响应"""

    task_id: str = Field(..., description="任务ID")
    status: str = Field(..., description="任务状态: pending/completed/failed")
    success: bool = Field(False, description="是否成功")
    candidates: List[LocateCandidate] = Field(default_factory=list, description="定位候选列表")
    message: str = Field(default="", description="提示信息")


# ==================== 路由 ====================


@router.post("/locate", response_model=LocateTaskResponse)
async def create_locate_task(data: LocateRequest):
    """
    创建定位任务

    1. 创建任务文件
    2. 通过 SSH 调用宿主机 Claude Code 执行定位脚本
    3. 返回 task_id，前端轮询结果
    """
    try:
        # 确保任务目录存在
        TASK_DIR.mkdir(parents=True, exist_ok=True)

        # 生成唯一任务 ID
        task_id = f"locate_{uuid.uuid4().hex[:12]}"

        # 创建任务文件
        task_file = TASK_DIR / f"{task_id}.json"
        task_data = {
            "step_description": data.step_description,
            "page_url": data.page_url,
            "expected": data.expected,
        }
        with open(task_file, "w", encoding="utf-8") as f:
            json.dump(task_data, f, ensure_ascii=False, indent=2)

        logger.info(f"AI定位任务已创建: {task_id}, 步骤: {data.step_description}")

        # SSH 调用宿主机执行定位脚本（异步，不阻塞）
        ssh_cmd = (
            f"ssh {HOST_USER}@{HOST} "
            f"\"cd {DEPLOY_DIR}/backend && python3 scripts/ai_locate_worker.py --task-id {task_id}\""
        )
        try:
            subprocess.Popen(
                ssh_cmd,
                shell=True,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
            logger.info(f"已通过 SSH 调用宿主机定位脚本: {task_id}")
        except Exception as e:
            logger.warning(f"SSH 调用宿主机失败: {e}，任务文件已创建，可手动执行")
            # SSH 失败不影响任务创建，用户可以手动执行脚本

        return LocateTaskResponse(
            task_id=task_id,
            status="pending",
            message="定位任务已创建",
        )

    except Exception as e:
        logger.error(f"创建定位任务失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"创建定位任务失败: {str(e)}")


@router.get("/locate/{task_id}/result", response_model=LocateResultResponse)
async def get_locate_result(task_id: str):
    """
    查询定位结果

    根据 task_id 读取结果文件，返回定位候选列表。
    """
    try:
        result_file = TASK_DIR / f"{task_id}_result.json"

        if not result_file.exists():
            # 结果文件不存在，检查任务文件是否存在
            task_file = TASK_DIR / f"{task_id}.json"
            if not task_file.exists():
                return LocateResultResponse(
                    task_id=task_id,
                    status="failed",
                    success=False,
                    candidates=[],
                    message="任务不存在",
                )

            return LocateResultResponse(
                task_id=task_id,
                status="pending",
                success=False,
                candidates=[],
                message="任务正在执行中",
            )

        # 读取结果文件
        with open(result_file, "r", encoding="utf-8") as f:
            result = json.load(f)

        # 构建候选列表
        candidates = []
        for c in result.get("candidates", []):
            candidates.append(LocateCandidate(
                locator_type=c.get("locator_type", ""),
                locator_value=c.get("locator_value", ""),
                confidence=c.get("confidence", 0),
                element_info=c.get("element_info", {}),
            ))

        return LocateResultResponse(
            task_id=task_id,
            status=result.get("status", "completed"),
            success=result.get("success", False),
            candidates=candidates,
            message=result.get("message", ""),
        )

    except Exception as e:
        logger.error(f"查询定位结果失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"查询定位结果失败: {str(e)}")