"""
工作手机SDK v3.0 - AI 智能命令引擎
使用本地 Ollama 或 Gemini 免费 API 解析自然语言命令

支持：
- Ollama 本地模型（免费，默认）
- Google Gemini Flash（免费 API）
- DeepSeek（付费备选）
"""

import json
import logging
import httpx
import os
from typing import List, Dict, Optional, Any

logger = logging.getLogger(__name__)

# ========== AI 后端配置 ==========

# Ollama（本地免费，首选）
OLLAMA_URL = os.getenv("OLLAMA_URL", "http://localhost:11434")
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "qwen2.5:1.5b")

# Gemini（免费 API，备选）
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "")
GEMINI_MODEL = "gemini-2.0-flash"

# ========== 系统提示词 ==========

SYSTEM_PROMPT = """你是安卓手机操作助手。用户说指令，你只输出JSON数组。

【重要】你只能用以下10种action，禁止发明新action：
open_app, click, input, swipe, back, home, wait, screenshot, key_event, shell

APP表：微信、QQ、抖音、小红书、快手、B站、淘宝、京东、拼多多、闲鱼、支付宝、美团、饿了么、高德地图、滴滴、设置、Chrome、Termux、飞书、钉钉、企业微信、相机、电话、短信、相册、文件管理、日历、地图

操作格式（只用这些）：
{"action":"open_app","app":"APP名"}
{"action":"click","text":"按钮文字"}
{"action":"input","text":"内容"}
{"action":"swipe","direction":"up/down/left/right"}
{"action":"back"}
{"action":"home"}
{"action":"wait","seconds":2}
{"action":"screenshot"}
{"action":"key_event","key":"enter/back/home/search/del/tab"}
{"action":"shell","command":"adb shell命令"}

示例：
Q: 打开微信
A: [{"action":"open_app","app":"微信"}]

Q: 打开Chrome搜索今天天气
A: [{"action":"open_app","app":"Chrome"},{"action":"wait","seconds":2},{"action":"click","text":"搜索或输入网址"},{"action":"input","text":"今天天气"},{"action":"key_event","key":"enter"}]

Q: 打开微信发消息给文件传输助手说你好
A: [{"action":"open_app","app":"微信"},{"action":"wait","seconds":2},{"action":"click","text":"搜索"},{"action":"input","text":"文件传输助手"},{"action":"click","text":"文件传输助手"},{"action":"wait","seconds":1},{"action":"input","text":"你好"},{"action":"click","text":"发送"}]

Q: 打开抖音刷视频
A: [{"action":"open_app","app":"抖音"},{"action":"wait","seconds":3},{"action":"swipe","direction":"up"}]

Q: 打开设置看WIFI
A: [{"action":"open_app","app":"设置"},{"action":"wait","seconds":1},{"action":"click","text":"WLAN"}]

Q: 打开支付宝扫一扫
A: [{"action":"open_app","app":"支付宝"},{"action":"wait","seconds":2},{"action":"click","text":"扫一扫"}]

Q: 将系统语言设置成中文
A: [{"action":"open_app","app":"设置"},{"action":"wait","seconds":1},{"action":"click","text":"System"},{"action":"click","text":"Languages"},{"action":"click","text":"中文"}]

Q: 修改系统语言
A: [{"action":"open_app","app":"设置"},{"action":"wait","seconds":1},{"action":"shell","command":"settings put system system_locales zh-CN"},{"action":"shell","command":"setprop persist.sys.locale zh-CN"}]

Q: 点击设置里的某个选项
A: [{"action":"open_app","app":"设置"},{"action":"wait","seconds":1},{"action":"click","text":"选项文字"}]

Q: 注册微信账号
A: [{"action":"open_app","app":"微信"},{"action":"wait","seconds":3},{"action":"click","text":"注册"},{"action":"wait","seconds":2}]

规则：
1. 只输出JSON数组，无其他文字
2. action只能是上面10种之一，绝对不要发明新的action名
3. 打开APP后加wait等待加载
4. 搜索操作用input+key_event enter
5. 修改系统设置优先用shell命令
6. 不确定就拆成open_app+click步骤"""


class AICommandEngine:
    """AI 命令解析引擎"""

    def __init__(self):
        self._backend = None  # 缓存已检测到的后端
        self._available = None

    async def detect_backend(self, force=False) -> str:
        """检测可用的 AI 后端"""
        if self._backend and self._backend != "none" and not force:
            return self._backend

        # 1. Ollama（本地免费）
        try:
            async with httpx.AsyncClient(timeout=5.0) as c:
                r = await c.get(f"{OLLAMA_URL}/api/tags")
                if r.status_code == 200:
                    models = [m["name"] for m in r.json().get("models", [])]
                    if models:
                        self._backend = "ollama"
                        self._available = models
                        logger.info(f"AI 后端: Ollama ({', '.join(models)})")
                        return "ollama"
                    else:
                        logger.info("Ollama 在线但无模型（可能在下载中）")
        except Exception:
            pass

        # 2. Gemini 免费
        if GEMINI_API_KEY:
            self._backend = "gemini"
            logger.info("AI 后端: Gemini Flash (免费)")
            return "gemini"

        self._backend = "none"
        logger.warning("无可用 AI 后端（Ollama 未运行/无模型，Gemini 未配置）")
        return "none"

    async def parse(self, user_input: str, context: Optional[Dict] = None) -> List[Dict[str, Any]]:
        """
        解析自然语言命令，返回操作列表。

        Args:
            user_input: 用户输入的自然语言
            context: 可选，当前屏幕上下文（当前APP等）

        Returns:
            操作列表 [{"action": "...", ...}, ...]
        """
        backend = await self.detect_backend()

        # 构建 prompt
        prompt = user_input
        if context:
            prompt = f"当前APP: {context.get('current_app', '未知')}\n用户指令: {user_input}"

        if backend == "ollama":
            return await self._call_ollama(prompt)
        elif backend == "gemini":
            return await self._call_gemini(prompt)
        else:
            return []

    async def _call_ollama(self, prompt: str) -> List[Dict]:
        """调用本地 Ollama"""
        try:
            async with httpx.AsyncClient(timeout=60.0) as c:
                r = await c.post(f"{OLLAMA_URL}/api/chat", json={
                    "model": OLLAMA_MODEL,
                    "messages": [
                        {"role": "system", "content": SYSTEM_PROMPT},
                        {"role": "user", "content": prompt},
                    ],
                    "stream": False,
                    "options": {"temperature": 0.3, "num_predict": 500},
                })
                if r.status_code != 200:
                    logger.error(f"Ollama 错误: {r.status_code} {r.text[:200]}")
                    return []
                text = r.json().get("message", {}).get("content", "").strip()
                logger.info(f"Ollama 原始回复: {text[:200]}")
                return self._extract_actions(text)
        except Exception as e:
            logger.error(f"Ollama 调用失败: {type(e).__name__}: {e}")
            return []

    async def _call_gemini(self, prompt: str) -> List[Dict]:
        """调用 Google Gemini 免费 API"""
        try:
            url = f"https://generativelanguage.googleapis.com/v1beta/models/{GEMINI_MODEL}:generateContent?key={GEMINI_API_KEY}"
            async with httpx.AsyncClient(timeout=30.0) as c:
                r = await c.post(url, json={
                    "contents": [{"parts": [{"text": f"{SYSTEM_PROMPT}\n\n用户: {prompt}\n返回:"}]}],
                    "generationConfig": {"temperature": 0.3, "maxOutputTokens": 500},
                })
                if r.status_code != 200:
                    logger.error(f"Gemini 错误: {r.status_code} {r.text[:200]}")
                    return []
                text = r.json()["candidates"][0]["content"]["parts"][0]["text"].strip()
                return self._extract_actions(text)
        except Exception as e:
            logger.error(f"Gemini 调用失败: {e}")
            return []

    # 有效的 action 类型白名单
    VALID_ACTIONS = {
        "open_app", "click", "input", "swipe", "back", "home",
        "wait", "screenshot", "key_event", "shell", "close_app",
    }

    # 无效 action 的智能修复映射
    ACTION_FIX_MAP = {
        # 系统设置相关
        "settings": "open_app",       # settings -> 打开设置
        "setlocale": "shell",         # setlocale -> shell 命令
        "set_language": "shell",      # set_language -> shell 命令
        "change_language": "shell",
        "system_setting": "open_app",
        "setting": "open_app",
        # 操作类别
        "tap": "click",
        "press": "key_event",
        "type": "input",
        "write": "input",
        "scroll": "swipe",
        "scroll_up": "swipe",
        "scroll_down": "swipe",
        "launch": "open_app",
        "start": "open_app",
        "open": "open_app",
        "close": "close_app",
        "stop": "close_app",
        "kill": "close_app",
        "go_back": "back",
        "go_home": "home",
        "return": "back",
        "capture": "screenshot",
        "delay": "wait",
        "sleep": "wait",
        "pause": "wait",
        "send": "click",             # send -> 点击发送按钮
        "search": "click",           # search -> 点击搜索
        "register": "click",         # register -> 点击注册
        "login": "click",            # login -> 点击登录
    }

    def _fix_action(self, action: Dict) -> Dict:
        """修复无效的 action 类型，尽量保留用户意图"""
        act = action.get("action", "")

        if act in self.VALID_ACTIONS:
            return action  # 已经有效

        # 尝试映射修复
        fixed_act = self.ACTION_FIX_MAP.get(act.lower())

        if fixed_act:
            logger.info(f"AI action 修复: {act} -> {fixed_act}")
            new_action = dict(action)
            new_action["action"] = fixed_act

            # 根据修复后的类型补充参数
            if fixed_act == "open_app" and "app" not in new_action:
                # settings/setting -> 打开设置
                if "setting" in act.lower():
                    new_action["app"] = "设置"
            elif fixed_act == "shell":
                if "command" not in new_action:
                    # setlocale/set_language -> 设置中文
                    if "locale" in act.lower() or "language" in act.lower():
                        new_action["command"] = "settings put system system_locales zh-CN"
            elif fixed_act == "click":
                if "text" not in new_action:
                    # 用 action 名称本身作为按钮文字
                    text_map = {
                        "send": "发送", "search": "搜索",
                        "register": "注册", "login": "登录",
                    }
                    new_action["text"] = text_map.get(act.lower(), act)
            elif fixed_act == "swipe":
                if "direction" not in new_action:
                    if "up" in act.lower():
                        new_action["direction"] = "up"
                    elif "down" in act.lower():
                        new_action["direction"] = "down"
                    else:
                        new_action["direction"] = "up"

            return new_action

        # 最后尝试：含有关键词的智能推断
        act_lower = act.lower()
        if any(kw in act_lower for kw in ["app", "open", "launch", "start"]):
            action["action"] = "open_app"
            if "app" not in action:
                action["app"] = "设置"
            return action
        if any(kw in act_lower for kw in ["click", "tap", "press", "button"]):
            action["action"] = "click"
            return action
        if any(kw in act_lower for kw in ["type", "input", "text", "write"]):
            action["action"] = "input"
            return action

        logger.warning(f"AI 生成了无法修复的操作: {act}，跳过")
        return None  # 无法修复

    def _extract_actions(self, text: str) -> List[Dict]:
        """从 AI 返回的文本中提取 JSON 操作列表"""
        text = text.strip()
        # 去除 markdown 代码块
        if "```json" in text:
            text = text.split("```json")[-1].split("```")[0].strip()
        elif "```" in text:
            text = text.split("```")[1].split("```")[0].strip() if text.count("```") >= 2 else text
        # 找到 JSON 数组
        start = text.find("[")
        end = text.rfind("]")
        if start >= 0 and end > start:
            text = text[start:end + 1]
        try:
            data = json.loads(text)
            if isinstance(data, dict):
                data = [data]
            if isinstance(data, list):
                raw_actions = [d for d in data if isinstance(d, dict) and "action" in d]
                # 校验并修复每个 action
                fixed_actions = []
                for a in raw_actions:
                    fixed = self._fix_action(a)
                    if fixed is not None:
                        fixed_actions.append(fixed)
                return fixed_actions
        except json.JSONDecodeError:
            logger.warning(f"AI 返回无法解析: {text[:200]}")
        return []

    async def get_status(self) -> Dict:
        """获取 AI 引擎状态（每次重新检测）"""
        backend = await self.detect_backend(force=True)
        return {
            "backend": backend,
            "model": OLLAMA_MODEL if backend == "ollama" else GEMINI_MODEL if backend == "gemini" else "无",
            "available": backend != "none",
            "models": self._available or [],
        }


# 全局实例
ai_engine = AICommandEngine()