one_dragon_context.py

import inspect
import logging
import threading
from enum import Enum
from functools import cached_property
from pathlib import Path

import cv2
from pynput import keyboard

from one_dragon.base.config.basic_model_config import BasicModelConfig
from one_dragon.base.config.custom_config import UILanguageEnum
from one_dragon.base.controller.controller_base import ControllerBase
from one_dragon.base.controller.pc_button.pc_button_listener import PcButtonListener
from one_dragon.base.matcher.ocr.ocr_matcher import OcrMatcher
from one_dragon.base.matcher.ocr.ocr_service import OcrService
from one_dragon.base.matcher.ocr.onnx_ocr_matcher import (
    DEFAULT_OCR_MODEL_NAME,
    PPOCRV6_MODEL_NAME,
    OnnxOcrMatcher,
    OnnxOcrParam,
    get_final_file_list,
)
from one_dragon.base.matcher.template_matcher import TemplateMatcher
from one_dragon.base.operation.application.application_factory_manager import (
    ApplicationFactoryManager,
)
from one_dragon.base.operation.application.application_group_manager import (
    ApplicationGroupManager,
)
from one_dragon.base.operation.application.application_run_context import (
    ApplicationRunContext,
)
from one_dragon.base.operation.application.plugin_info import PluginSource
from one_dragon.base.operation.context_event_bus import ContextEventBus
from one_dragon.base.operation.context_lazy_signal import ContextLazySignal
from one_dragon.base.operation.one_dragon_env_context import (
    ONE_DRAGON_CONTEXT_EXECUTOR,
    OneDragonEnvContext,
)
from one_dragon.base.operation.overlay_debug_bus import OverlayDebugBus
from one_dragon.base.push.push_service import PushService
from one_dragon.base.screen.screen_loader import ScreenContext
from one_dragon.base.screen.template_loader import TemplateLoader
from one_dragon.utils import debug_utils, file_utils, i18_utils, log_utils, thread_utils
from one_dragon.utils.log_utils import log


class ContextKeyboardEventEnum(Enum):

    PRESS = 'context_keyboard_press'


class ContextInstanceEventEnum(Enum):

    instance_active = 'instance_active'


class OneDragonContext(ContextEventBus, OneDragonEnvContext):

    def __init__(self):
        ContextEventBus.__init__(self)
        OneDragonEnvContext.__init__(self)

        self.signal: ContextLazySignal = ContextLazySignal()

        if self.one_dragon_config.current_active_instance is None:
            self.one_dragon_config.create_new_instance(True)
        self.current_instance_idx = self.one_dragon_config.current_active_instance.idx
        self.overlay_debug_bus: OverlayDebugBus = OverlayDebugBus()

        self.screen_loader: ScreenContext = ScreenContext()
        self.template_loader: TemplateLoader = TemplateLoader()
        self.tm: TemplateMatcher = TemplateMatcher(self.template_loader)
        self.tm.overlay_debug_bus = self.overlay_debug_bus

        self.ocr: OcrMatcher = OnnxOcrMatcher(
            OnnxOcrParam(
                use_gpu=self.model_config.ocr_use_gpu,
                det_limit_side_len=max(self.project_config.screen_standard_width, self.project_config.screen_standard_height),
            )
        )
        self.ocr.overlay_debug_bus = self.overlay_debug_bus
        self.ocr_service: OcrService = OcrService(ocr_matcher=self.ocr)
        self._ocr_v6_downloading: bool = False  # 后台下载 V6 是否进行中 防止重复启动
        self.controller: ControllerBase | None = None

        self.keyboard_controller = keyboard.Controller()
        self.btn_listener = PcButtonListener(on_button_tap=self._on_key_press, listen_keyboard=True, listen_mouse=True)
        self.btn_listener.start()

        # 注册应用
        self.run_context: ApplicationRunContext = ApplicationRunContext(self)
        self.app_group_manager: ApplicationGroupManager = ApplicationGroupManager(self)

        self.push_service: PushService = PushService(self)

        # 初始化相关
        self._init_lock = threading.Lock()
        self._application_registered: bool = False  # 应用是否注册完毕
        self.ready_for_application: bool = False  # 初始化完成 可以运行应用了

    #------------------- 需要懒加载的都使用 @cached_property -------------------#

    #------------------- 以下是 应用工厂相关 -------------------#

    @cached_property
    def application_plugin_dirs(self) -> list[tuple[Path, PluginSource]]:
        """
        应用插件目录列表

        默认返回：
        1. 子类所在目录的同级 'application' 目录（内置应用）
        2. 项目根目录下的 'plugins' 目录（外部插件，支持相对导入）

        例如：如果子类在 zzz_od/context/zzz_context.py，则返回：
        - (zzz_od/application, BUILTIN)
        - ({project_root}/plugins, THIRD_PARTY)

        Returns:
            list[tuple[Path, PluginSource]]: 应用插件目录列表
        """
        dirs: list[tuple[Path, PluginSource]] = []

        # 获取实际子类的定义文件
        cls_file = inspect.getfile(self.__class__)
        parent_dir = Path(cls_file).parent.parent

        # 计算 application 目录：子类文件所在目录的上级目录下的 application 目录
        # 例如：zzz_od/context/zzz_context.py -> zzz_od/application
        application_dir = parent_dir / 'application'
        if application_dir.is_dir():
            dirs.append((application_dir, PluginSource.BUILTIN))

        # 计算项目根目录下的 plugins 目录（外部插件）
        # 从 src 目录往上一级就是项目根目录
        src_dir = file_utils.find_src_dir(cls_file)
        if src_dir is not None:
            project_root = src_dir.parent
            plugins_dir = project_root / 'plugins'
            if plugins_dir.is_dir():
                dirs.append((plugins_dir, PluginSource.THIRD_PARTY))

        return dirs

    @cached_property
    def factory_manager(self) -> ApplicationFactoryManager:
        """应用工厂管理器"""
        return ApplicationFactoryManager(self, self.application_plugin_dirs)

    #------------------- 以下是 游戏/脚本级别的 -------------------#

    @cached_property
    def one_dragon_config(self):
        from one_dragon.base.config.one_dragon_config import OneDragonConfig
        return OneDragonConfig()

    @cached_property
    def custom_config(self):
        from one_dragon.base.config.custom_config import CustomConfig
        return CustomConfig()

    @cached_property
    def pip_config(self):
        from one_dragon.base.config.pip_config import PipConfig
        return PipConfig()

    @cached_property
    def cv_service(self):
        from one_dragon.base.cv_process.cv_service import CvService
        return CvService(self)

    @cached_property
    def model_config(self) -> BasicModelConfig:
        return BasicModelConfig()

    #------------------- 以下是 账号实例级别的 需要在 reload_instance_config 中刷新 -------------------#

    @cached_property
    def game_account_config(self):
        from one_dragon.base.config.game_account_config import GameAccountConfig
        return GameAccountConfig(self.current_instance_idx)

    @cached_property
    def notify_config(self):
        from one_dragon.base.config.notify_config import NotifyConfig
        return NotifyConfig(self.current_instance_idx, self.run_context.notify_app_map)

    @cached_property
    def standalone_app_config(self):
        """应用运行界面的配置（保存用户添加的应用列表）"""
        from one_dragon.base.config.standalone_app_config import StandaloneAppConfig
        return StandaloneAppConfig(self.current_instance_idx)

    #------------------- 以下是 应用注册相关 -------------------#

    def register_application_factory(self) -> None:
        """注册应用

        使用工厂管理器自动扫描和注册应用工厂。
        """
        # 发现并注册应用
        non_default_factories, default_factories = self.factory_manager.discover_factories()

        if non_default_factories:
            self.run_context.registry_application(non_default_factories, default_group=False)

        if default_factories:
            self.run_context.registry_application(default_factories, default_group=True)

    def _load_plugin_screens(self) -> None:
        """加载插件的 screen_info

        遍历所有已注册插件，从其 screen_info/ 子目录加载 screen YAML。
        未设 app_id 的 screen 自动使用插件的 app_id。
        应在 screen_loader.reload() 之后调用。
        """
        for plugin_info in self.factory_manager.plugin_infos:
            if plugin_info.plugin_dir is None:
                continue
            screen_dir = plugin_info.plugin_dir / 'screen_info'
            self.screen_loader.load_extra_screen_dir(
                str(screen_dir), default_app_id=plugin_info.app_id
            )

    def refresh_application_registration(self) -> None:
        """刷新应用注册

        重新扫描插件目录，刷新所有应用的注册。
        可在运行时调用以加载新的应用或更新已有应用。
        """
        log.info("开始刷新应用注册...")

        # 清空现有注册
        self.run_context.clear_applications()

        # 重新发现并注册
        non_default_factories, default_factories = self.factory_manager.discover_factories(reload_modules=True)

        if non_default_factories:
            self.run_context.registry_application(non_default_factories, default_group=False)

        if default_factories:
            self.run_context.registry_application(default_factories, default_group=True)

        # 重新加载插件 screen
        self.screen_loader.reload()
        self._load_plugin_screens()

        # 更新默认应用组
        self.app_group_manager.set_default_apps(self.run_context.default_group_apps)

        # 清除应用组配置缓存，使其重新加载
        self.app_group_manager.clear_config_cache()

        # 刷新通知配置中的应用映射
        if 'notify_config' in self.__dict__:
            del self.__dict__['notify_config']

        log.info("应用注册刷新完成")

    #------------------- 以下是 初始化相关 -------------------#

    def init(self) -> None:
        if not self._init_lock.acquire(blocking=False):
            return

        try:
            self.ready_for_application = False

            if self.custom_config.ui_language == UILanguageEnum.AUTO.value.value:
                i18_utils.detect_and_set_default_language()
            else:
                i18_utils.update_default_lang(self.custom_config.ui_language)

            log_utils.set_log_level(logging.DEBUG if self.env_config.is_debug else logging.INFO)

            if not self._application_registered:  # 只需要注册一次
                self.register_application_factory()
                self.app_group_manager.set_default_apps(self.run_context.default_group_apps)
                self._application_registered = True

            self.init_ocr()

            self.screen_loader.reload()
            self._load_plugin_screens()

            # 账号实例层级的配置 不是应用特有的配置
            self.reload_instance_config()

            # 初始化控制器
            self.init_controller()

            self.init_for_application()

            self.ready_for_application = True

            self.run_context.check_and_update_all_run_record(self.current_instance_idx)

            self.push_service.init_push_channels()

            # 只有在配置了 ghproxy 代理时才更新代理地址
            if self.env_config.is_gh_proxy:
                self.gh_proxy_service.update_proxy_url()

            self.init_others()
        except Exception:
            log.error('初始化出错', exc_info=True)
        finally:
            self._init_lock.release()

    def init_controller(self) -> None:
        """
        初始化控制器
        由子类自行实现
        """
        pass

    def init_for_application(self) -> None:
        """
        执行应用前 还需要做的初始化
        应该加载游戏级别
        由子类自行实现
        """
        pass

    def init_others(self) -> None:
        """
        其他非必要的初始化
        由子类自行实现
        """
        pass

    def init_async(self) -> None:
        """
        异步初始化
        """
        f = ONE_DRAGON_CONTEXT_EXECUTOR.submit(self.init)
        f.add_done_callback(thread_utils.handle_future_result)

    def _on_key_press(self, key: str):
        """
        按键时触发 抛出事件，事件体为按键
        :param key: 按键
        :return:
        """
        # log.info('按键 %s' % key)
        if key == self.key_start_running:
            self.run_context.switch_context_pause_and_run()
        elif key == self.key_stop_running:
            self.run_context.stop_running()
        elif key == self.key_screenshot:
            self.screenshot_and_save_debug(self.env_config.copy_screenshot)

        self.dispatch_event(ContextKeyboardEventEnum.PRESS.value, key)

    @property
    def is_game_window_ready(self) -> bool:
        """
        游戏窗口是否已经出现
        :return:
        """
        return self.controller is not None and self.controller.is_game_window_ready

    @property
    def key_start_running(self) -> str:
        return self.env_config.key_start_running

    @property
    def key_stop_running(self) -> str:
        return self.env_config.key_stop_running

    @property
    def key_screenshot(self) -> str:
        return self.env_config.key_screenshot

    @property
    def key_debug(self) -> str:
        return self.env_config.key_debug

    def screenshot_and_save_debug(self, copy_screenshot: bool) -> None:
        """
        截图 保存到debug
        """
        if self.controller is None or not self.controller.is_game_window_ready:
            return
        if self.controller.game_win is not None:
            self.controller.game_win.active()
        _, img = self.controller.screenshot(independent=True)
        base_file_name = debug_utils.save_debug_image(img, copy_screenshot=copy_screenshot)
        self._save_overlay_patched_image(img, base_file_name)

    def _save_overlay_patched_image(self, base_image, base_file_name: str) -> None:
        try:
            from one_dragon_qt.overlay.overlay_config import OverlayConfig
        except Exception:
            return

        config = OverlayConfig()
        if not config.patched_capture_enabled:
            return

        overlay_rgba = self._capture_overlay_rgba()
        if overlay_rgba is None:
            return

        patched = self._compose_overlay_patched_image(base_image, overlay_rgba)
        if patched is None:
            return

        debug_utils.save_debug_image(
            patched,
            file_name=f"{base_file_name}{config.patched_capture_suffix}",
            copy_screenshot=False,
        )

    @staticmethod
    def _compose_overlay_patched_image(base_image, overlay_rgba):
        if base_image is None or overlay_rgba is None:
            return None
        if len(overlay_rgba.shape) != 3 or overlay_rgba.shape[2] < 4:
            return None

        target_h, target_w = base_image.shape[:2]
        if target_h <= 0 or target_w <= 0:
            return None

        if overlay_rgba.shape[0] != target_h or overlay_rgba.shape[1] != target_w:
            overlay_rgba = cv2.resize(
                overlay_rgba,
                (target_w, target_h),
                interpolation=cv2.INTER_LINEAR,
            )

        overlay_rgb = overlay_rgba[:, :, :3].astype("float32")
        overlay_bgr = overlay_rgb[:, :, ::-1]  # RGBA -> BGR 通道顺序
        alpha = (overlay_rgba[:, :, 3:4].astype("float32") / 255.0).clip(0.0, 1.0)
        base_rgb = base_image.astype("float32")
        patched = base_rgb * (1.0 - alpha) + overlay_bgr * alpha
        return patched.astype("uint8")

    @staticmethod
    def _capture_overlay_rgba():
        try:
            from one_dragon_qt.overlay.overlay_manager import OverlayManager
        except Exception:
            return None
        manager = OverlayManager.instance()
        if manager is None:
            return None
        return manager.capture_overlay_rgba()

    def switch_instance(self, instance_idx: int) -> None:
        """
        切换实例
        :param instance_idx:
        :return:
        """
        self.one_dragon_config.active_instance(instance_idx)
        self.current_instance_idx = self.one_dragon_config.current_active_instance.idx
        self.reload_instance_config()
        self.on_switch_instance()
        self.dispatch_event(ContextInstanceEventEnum.instance_active.value, instance_idx)

    def on_switch_instance(self) -> None:
        """
        切换实例后的回调，用于更新 controller 配置
        由子类实现具体逻辑
        """
        pass

    def reload_instance_config(self):
        """
        重新加载账号实例相关的配置，不是单个应用特有的配置
        注意如果有缓存需要清理缓存
        子类需要继承加载更多的配置
        """
        log.info('开始加载实例配置 %d', self.current_instance_idx)

        to_clear_props = [
            'game_account_config',
            'notify_config',
            'standalone_app_config',
        ]
        for prop in to_clear_props:
            if prop in self.__dict__:
                del self.__dict__[prop]

    def init_ocr(self) -> None:
        """
        初始化OCR
        :return:
        """
        ocr_model_name = self._decide_ocr_model_name()

        # 清理旧实例资源
        if hasattr(self, 'ocr') and self.ocr is not None:
            if hasattr(self.ocr, 'cleanup'):
                self.ocr.cleanup()

        self.ocr = OnnxOcrMatcher(
            OnnxOcrParam(
                ocr_model_name=ocr_model_name,
                use_gpu=self.model_config.ocr_use_gpu,
                det_limit_side_len=max(self.project_config.screen_standard_width, self.project_config.screen_standard_height),
            )
        )
        self.ocr.overlay_debug_bus = self.overlay_debug_bus
        self.ocr_service.ocr_matcher = self.ocr
        if 'cv_service' in self.__dict__:
            self.cv_service.ocr = self.ocr
        self.ocr.init_model(
            ghproxy_url=self.env_config.gh_proxy_url if self.env_config.is_gh_proxy else None,
            proxy_url=self.env_config.personal_proxy if self.env_config.is_personal_proxy else None,
        )

        # 正常模式下 向 V6 收敛
        if not self.env_config.is_debug:
            if self._is_ocr_model_ready(PPOCRV6_MODEL_NAME):
                # V6 已就绪（启动就有 或 刚同步下载完成） 落盘配置
                self.model_config.ocr = PPOCRV6_MODEL_NAME
            elif ocr_model_name == DEFAULT_OCR_MODEL_NAME:
                # 当前用 V5 顶住 后台下载 V6 成功后自动切换
                self._download_ocr_v6_in_background()

    def _decide_ocr_model_name(self) -> str:
        """
        决定本次初始化使用的 OCR 模型名

        调试模式: 直接使用配置里的模型 不自动切换
        正常模式: 向 V6 收敛
          - V6 文件齐全 -> 用 V6
          - V6 不齐 但 V5 齐全 -> 先用 V5 顶住 后台下载 V6
          - 都没有 -> 直接用 V6（会触发下载）
        """
        if self.env_config.is_debug:
            return self.model_config.ocr

        if self._is_ocr_model_ready(PPOCRV6_MODEL_NAME):
            return PPOCRV6_MODEL_NAME

        if self._is_ocr_model_ready(DEFAULT_OCR_MODEL_NAME):
            return DEFAULT_OCR_MODEL_NAME

        return PPOCRV6_MODEL_NAME

    @staticmethod
    def _is_ocr_model_ready(ocr_model_name: str) -> bool:
        """
        判断某个 OCR 模型的文件是否已经全部就绪
        """
        return all(Path(f).exists() for f in get_final_file_list(ocr_model_name))

    def _download_ocr_v6_in_background(self) -> None:
        """
        后台下载 V6 模型 下载成功后落盘配置 下次启动自动生效
        已有下载任务进行中时 不重复启动
        """
        if self._ocr_v6_downloading:
            return
        self._ocr_v6_downloading = True

        def download_task() -> None:
            try:
                v6_matcher = OnnxOcrMatcher(
                    OnnxOcrParam(
                        ocr_model_name=PPOCRV6_MODEL_NAME,
                    )
                )
                done = v6_matcher.download(
                    download_by_github=True,
                    download_by_gitee=False,
                    download_by_mirror_chan=False,
                    ghproxy_url=self.env_config.gh_proxy_url if self.env_config.is_gh_proxy else None,
                    proxy_url=self.env_config.personal_proxy if self.env_config.is_personal_proxy else None,
                )
                if done:
                    # 只落盘配置 不立刻切换 避免切换失败导致当前可用的 V5 失效
                    self.model_config.ocr = PPOCRV6_MODEL_NAME
                    log.info('OCR V6 后台下载完成 配置已更新 下次启动自动切换')
                else:
                    log.error('OCR V6 后台下载失败 保持当前 V5 可用 下次启动再试')
            except Exception:
                log.error('OCR V6 后台下载异常 保持当前 V5 可用 下次启动再试', exc_info=True)
            finally:
                self._ocr_v6_downloading = False

        threading.Thread(target=download_task, daemon=True, name='ocr_v6_download').start()

    def after_app_shutdown(self) -> None:
        """
        App关闭后进行的操作 关闭一切可能资源操作
        @return:
        """
        self.btn_listener.stop()
        if self.controller is not None:
            self.controller.cleanup_after_app_shutdown()
        self.one_dragon_config.clear_temp_instance_indices()
        ContextEventBus.after_app_shutdown(self)
        OneDragonEnvContext.after_app_shutdown(self)
        from one_dragon.base.conditional_operation.operator import ConditionalOperator
        ConditionalOperator.after_app_shutdown()
        from one_dragon.base.conditional_operation.operation_executor import (
            OperationExecutor,
        )
        OperationExecutor.after_app_shutdown()
        from one_dragon.base.conditional_operation.state_record_service import (
            StateRecordService,
        )
        StateRecordService.after_app_shutdown()
        from one_dragon.utils import gpu_executor
        gpu_executor.shutdown(wait=False)
        from one_dragon.base.operation.application_base import Application
        Application.after_app_shutdown()
        self.run_context.after_app_shutdown()
        self.push_service.after_app_shutdown()
        self.overlay_debug_bus.clear()