"""
Hook 模块管理服务

实现能力：
1. 模块管理（增删改查/启停/scope）
2. 脚本管理（保存/下载/部署）
3. 设备模块状态（加载/重载/日志）
4. Hook 事件总线（历史 + 实时订阅）
"""

from __future__ import annotations

import asyncio
import base64
import hashlib
import json
import os
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Set


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


class HookModuleService:
    def __init__(self) -> None:
        app_dir = Path(__file__).resolve().parent.parent
        self.data_dir = app_dir / "data" / "hook"
        self.scripts_dir = self.data_dir / "scripts"
        self.modules_file = self.data_dir / "modules.json"
        self.events_file = self.data_dir / "events.jsonl"
        self.device_state_file = self.data_dir / "device_modules.json"
        self._lock = asyncio.Lock()
        self._event_clients: Set[Any] = set()
        self._ensure_store()

    def _ensure_store(self) -> None:
        self.scripts_dir.mkdir(parents=True, exist_ok=True)
        if not self.modules_file.exists():
            self.modules_file.write_text("{}", encoding="utf-8")
        if not self.device_state_file.exists():
            self.device_state_file.write_text("{}", encoding="utf-8")
        if not self.events_file.exists():
            self.events_file.write_text("", encoding="utf-8")

    def _read_json(self, path: Path, default: Any) -> Any:
        try:
            return json.loads(path.read_text(encoding="utf-8"))
        except Exception:
            return default

    def _write_json(self, path: Path, data: Any) -> None:
        path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")

    async def list_modules(self, enabled: Optional[bool] = None, scope: Optional[str] = None) -> List[dict]:
        async with self._lock:
            modules = self._read_json(self.modules_file, {})
        result = list(modules.values())
        if enabled is not None:
            result = [m for m in result if bool(m.get("enabled")) == enabled]
        if scope:
            result = [m for m in result if scope in (m.get("scopes") or [])]
        return sorted(result, key=lambda x: x.get("updated_at", ""), reverse=True)

    async def get_module(self, module_id: str) -> Optional[dict]:
        async with self._lock:
            modules = self._read_json(self.modules_file, {})
            return modules.get(module_id)

    async def upsert_module(self, payload: dict) -> dict:
        module_id = payload["module_id"]
        now = _now_iso()
        script_content = payload.pop("script_content", None)
        script_url = payload.get("script_url")
        script_hash = payload.get("script_hash")
        if script_content and not script_url:
            decoded = base64.b64decode(script_content)
            script_name = f"{module_id}.js"
            target = self.scripts_dir / script_name
            target.write_bytes(decoded)
            digest = hashlib.sha256(decoded).hexdigest()
            script_url = f"/api/v3/scripts/{module_id}"
            script_hash = f"sha256:{digest}"

        async with self._lock:
            modules = self._read_json(self.modules_file, {})
            old = modules.get(module_id, {})
            data = {
                "module_id": module_id,
                "name": payload.get("name", old.get("name", module_id)),
                "version": payload.get("version", old.get("version", "1.0.0")),
                "description": payload.get("description", old.get("description", "")),
                "enabled": bool(payload.get("enabled", old.get("enabled", True))),
                "scopes": payload.get("scopes", old.get("scopes", ["com.tencent.mm"])),
                "capabilities": payload.get("capabilities", old.get("capabilities", [])),
                "min_frida_version": payload.get("min_frida_version", old.get("min_frida_version", "16.0.0")),
                "script_url": script_url or old.get("script_url"),
                "script_hash": script_hash or old.get("script_hash"),
                "script_id": payload.get("script_id", old.get("script_id", module_id)),
                "device_count": old.get("device_count", 0),
                "error_count": old.get("error_count", 0),
                "created_at": old.get("created_at", now),
                "updated_at": now,
            }
            modules[module_id] = data
            self._write_json(self.modules_file, modules)
            return data

    async def delete_module(self, module_id: str) -> bool:
        async with self._lock:
            modules = self._read_json(self.modules_file, {})
            if module_id not in modules:
                return False
            del modules[module_id]
            self._write_json(self.modules_file, modules)
        return True

    async def set_scope(self, module_id: str, scopes: List[str]) -> Optional[dict]:
        async with self._lock:
            modules = self._read_json(self.modules_file, {})
            item = modules.get(module_id)
            if not item:
                return None
            item["scopes"] = scopes
            item["updated_at"] = _now_iso()
            modules[module_id] = item
            self._write_json(self.modules_file, modules)
            return item

    async def set_enabled(self, module_id: str, enabled: bool) -> Optional[dict]:
        async with self._lock:
            modules = self._read_json(self.modules_file, {})
            item = modules.get(module_id)
            if not item:
                return None
            item["enabled"] = enabled
            item["updated_at"] = _now_iso()
            modules[module_id] = item
            self._write_json(self.modules_file, modules)
            return item

    async def list_scripts(self) -> List[dict]:
        out: List[dict] = []
        for p in sorted(self.scripts_dir.glob("*.js"), key=lambda x: x.stat().st_mtime, reverse=True):
            raw = p.read_bytes()
            out.append(
                {
                    "script_id": p.stem,
                    "filename": p.name,
                    "size": len(raw),
                    "hash": f"sha256:{hashlib.sha256(raw).hexdigest()}",
                    "url": f"/api/v3/scripts/{p.stem}",
                    "updated_at": datetime.fromtimestamp(p.stat().st_mtime, tz=timezone.utc).isoformat(),
                }
            )
        return out

    async def save_script(self, script_id: str, content: bytes) -> dict:
        target = self.scripts_dir / f"{script_id}.js"
        target.write_bytes(content)
        digest = hashlib.sha256(content).hexdigest()
        return {
            "script_id": script_id,
            "url": f"/api/v3/scripts/{script_id}",
            "hash": f"sha256:{digest}",
            "size": len(content),
        }

    async def get_script_path(self, script_id: str) -> Optional[Path]:
        target = self.scripts_dir / f"{script_id}.js"
        return target if target.exists() else None

    async def get_device_modules(self, device_id: str) -> dict:
        async with self._lock:
            state = self._read_json(self.device_state_file, {})
            item = state.get(device_id, {})
            modules = item.get("modules", [])
            return {
                "device_id": device_id,
                "supports_hook": item.get("supports_hook", False),
                "frida_version": item.get("frida_version", ""),
                "hook_framework": item.get("hook_framework", "frida-server"),
                "root_status": item.get("root_status", False),
                "modules": modules,
            }

    async def update_device_probe(self, device_id: str, probe: dict) -> None:
        async with self._lock:
            state = self._read_json(self.device_state_file, {})
            item = state.get(device_id, {})
            item.update(
                {
                    "supports_hook": probe.get("supports_hook", False),
                    "frida_version": probe.get("frida_version", ""),
                    "root_status": probe.get("root_status", False),
                    "hook_framework": "frida-server",
                    "updated_at": _now_iso(),
                }
            )
            item.setdefault("modules", [])
            state[device_id] = item
            self._write_json(self.device_state_file, state)

    async def deploy_script(self, script_id: str, device_ids: List[str], auto_reload: bool = True) -> dict:
        async with self._lock:
            state = self._read_json(self.device_state_file, {})
            deployed: List[str] = []
            failed: List[dict] = []
            reloaded: List[str] = []
            now = _now_iso()
            for did in device_ids:
                item = state.setdefault(
                    did,
                    {"supports_hook": False, "frida_version": "", "root_status": False, "modules": []},
                )
                modules = item.setdefault("modules", [])
                exists = None
                for m in modules:
                    if m.get("module_id") == script_id:
                        exists = m
                        break
                if exists:
                    exists["status"] = "loaded"
                    exists["loaded_at"] = now
                else:
                    modules.append(
                        {
                            "module_id": script_id,
                            "version": "1.0.0",
                            "status": "loaded",
                            "loaded_at": now,
                            "target_process": "com.tencent.mm",
                            "target_pid": 0,
                            "rpc_methods": [],
                            "last_error": None,
                            "events_today": 0,
                        }
                    )
                deployed.append(did)
                if auto_reload:
                    reloaded.append(did)
            self._write_json(self.device_state_file, state)
            return {"deployed": deployed, "failed": failed, "reloaded": reloaded}

    async def reload_device_modules(self, device_id: str, module_ids: List[str], force: bool = False) -> dict:
        async with self._lock:
            state = self._read_json(self.device_state_file, {})
            item = state.get(device_id)
            if not item:
                return {"reloaded": [], "failed": module_ids}
            mods = item.get("modules", [])
            found, failed = [], []
            now = _now_iso()
            for mid in module_ids:
                hit = None
                for m in mods:
                    if m.get("module_id") == mid:
                        hit = m
                        break
                if hit:
                    hit["status"] = "loaded"
                    hit["loaded_at"] = now
                    if force:
                        hit["last_error"] = None
                    found.append(mid)
                else:
                    failed.append(mid)
            item["modules"] = mods
            state[device_id] = item
            self._write_json(self.device_state_file, state)
            return {"reloaded": found, "failed": failed}

    async def add_device_log(self, device_id: str, module_id: str, line: str) -> None:
        async with self._lock:
            state = self._read_json(self.device_state_file, {})
            item = state.setdefault(device_id, {"modules": []})
            logs = item.setdefault("logs", {})
            bucket = logs.setdefault(module_id, [])
            bucket.append({"ts": _now_iso(), "line": line})
            if len(bucket) > 300:
                logs[module_id] = bucket[-300:]
            state[device_id] = item
            self._write_json(self.device_state_file, state)

    async def get_device_logs(self, device_id: str, module_id: str) -> List[dict]:
        async with self._lock:
            state = self._read_json(self.device_state_file, {})
            item = state.get(device_id, {})
            logs = item.get("logs", {}).get(module_id, [])
            return logs[-200:]

    async def add_event(self, event: dict) -> dict:
        data = {"event_id": f"evt_{int(datetime.now().timestamp()*1000)}", **event}
        if "timestamp" not in data:
            data["timestamp"] = _now_iso()
        async with self._lock:
            with self.events_file.open("a", encoding="utf-8") as f:
                f.write(json.dumps(data, ensure_ascii=False) + "\n")
        await self._broadcast_event({"type": "hook_event", "data": data})
        return data

    async def list_events(
        self,
        event_type: Optional[str] = None,
        device_id: Optional[str] = None,
        platform: Optional[str] = None,
        limit: int = 100,
    ) -> List[dict]:
        events: List[dict] = []
        if not self.events_file.exists():
            return events
        lines = self.events_file.read_text(encoding="utf-8").splitlines()[-2000:]
        for line in reversed(lines):
            if len(events) >= limit:
                break
            try:
                item = json.loads(line)
            except Exception:
                continue
            if event_type and item.get("event_type") != event_type:
                continue
            if device_id and item.get("device_id") != device_id:
                continue
            if platform and item.get("platform") != platform:
                continue
            events.append(item)
        return events

    async def attach_ws_client(self, ws: Any) -> None:
        self._event_clients.add(ws)

    async def detach_ws_client(self, ws: Any) -> None:
        self._event_clients.discard(ws)

    async def _broadcast_event(self, payload: dict) -> None:
        if not self._event_clients:
            return
        dead = []
        for ws in list(self._event_clients):
            try:
                await ws.send_json(payload)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self._event_clients.discard(ws)


hook_module_service = HookModuleService()
