from __future__ import annotations

import asyncio
import json
import logging
import sys
from collections import defaultdict
from pathlib import Path
from threading import Thread
from typing import TYPE_CHECKING, Any, Callable, Literal, Protocol

from ulauncher import cli, paths
from ulauncher.gi import GLib
from ulauncher.internals import log_wire
from ulauncher.modes.extensions import ext_exceptions
from ulauncher.modes.extensions.extension_dependencies import ExtensionDependencies
from ulauncher.modes.extensions.extension_record import ExtensionRecord, PreviewExtensionRecord
from ulauncher.modes.extensions.extension_registry import ExtensionRegistry
from ulauncher.modes.extensions.extension_runtime import DEBUGPY_HOST, DEBUGPY_PORT, ExtensionRuntime
from ulauncher.utils import scheduling
from ulauncher.utils.eventbus import EventBus

if TYPE_CHECKING:
    from ulauncher.internals import ipc

logger = logging.getLogger(__name__)
events = EventBus("extensions", skip_if_not_bound=True)


class ExtensionServiceListener(Protocol):
    """What the service reports back about the extensions it runs. Implemented by ExtensionMode,
    which registers itself via ExtensionService.activate."""

    def started(self, ext_id: str) -> None: ...

    def errored(self, ext_id: str) -> None: ...

    def invalidate_cache(self) -> None: ...

    def handle_message(self, ext_id: str, message: ipc.ExtensionMessage) -> None: ...


class ExtensionService(ExtensionRegistry):
    """The app's ExtensionRegistry: also owns the extension processes and handles lifecycle intents.

    Extensions only run in the app process, so only app code uses the service. The CLI does its
    disk work with a plain ExtensionRegistry and asks the app to reconcile the processes via
    D-Bus events ("extensions:reload", "extensions:stop", ...).
    """

    runtimes: dict[str, ExtensionRuntime]
    stopped_listeners: dict[str, list[Callable[[], None]]]
    listener: ExtensionServiceListener | None
    _preview_log: tuple[str, logging.Handler] | None

    def __init__(self) -> None:
        super().__init__(lifecycle=self)
        self.runtimes = {}
        self.stopped_listeners = defaultdict(list)
        self.listener = None
        self._preview_log = None
        events.set_self(self)

    def activate(self, listener: ExtensionServiceListener) -> None:
        """
        Set the listener that lifecycle events report to, and start the enabled extensions.
        Called by ExtensionMode when the app loads its modes.
        """
        self.listener = listener
        scheduling.run_when_idle(self._start_extensions)

    def _start_extensions(self) -> None:
        for ext in self.iterate():
            if ext.state.is_enabled:
                self.start_extension(ext)

    def _run_batch_job(
        self,
        extension_ids: list[str],
        jobs: list[Literal["start", "stop"]],
        callback: Callable[[], None] | None = None,
    ) -> None:
        records: list[ExtensionRecord] = []
        for ext_id in extension_ids:
            if ext := self.get(ext_id):
                records.append(ext)  # noqa: PERF401

        # run the reload in a separate thread to avoid blocking the main thread
        async def run_batch_async() -> None:
            for job in jobs:
                if job == "start":
                    for record in records:
                        if record.is_enabled:
                            self.start_extension(record)
                elif job == "stop":
                    await asyncio.gather(*[self.stop_extension(record) for record in records])

        def run_batch() -> None:
            try:
                asyncio.run(run_batch_async())
            finally:
                # callback must run unconditionally (owns the follow-up work)
                if callback:
                    callback()

        Thread(target=run_batch).start()

    def is_running(self, record: ExtensionRecord) -> bool:
        return record.id in self.runtimes

    async def toggle_enabled(self, record: ExtensionRecord, enabled: bool) -> bool:
        record.state.save(is_enabled=enabled, error_type="", error_message="")
        if enabled:
            return self.start_extension(record)
        await self.stop_extension(record)
        return False

    def start_extension(self, record: ExtensionRecord) -> bool:
        """
        Starts the extension in a subprocess
        Returns True if the extension was already running or successfully started, False otherwise.
        """
        if self.is_running(record):
            return True  # if already started, count as successful

        def exit_handler(cause: str, error_msg: str) -> None:
            listeners = self.stopped_listeners.get(record.id, [])
            for stop_listener in listeners:
                stop_listener()

            listeners.clear()

            if listener := self.listener:
                listener.invalidate_cache()

            if cause != "Stopped":
                logger.error('Extension "%s" exited with an error: %s (%s)', record.id, error_msg, cause)
                self.runtimes.pop(record.id, None)
                # A failing preview must not disable the installed extension by persisting its error.
                if not record.is_preview:
                    record.state.save(error_type=cause, error_message=error_msg)

                if listener := self.listener:
                    listener.errored(record.id)

                # Nothing asked the preview to stop, but it is over, so it still needs the teardown
                if self.preview is record:
                    self.stop_preview()

        def message_handler(message: ipc.ExtensionMessage) -> None:
            if listener := self.listener:
                listener.handle_message(record.id, message)

        try:
            record.manifest.validate()
            record.manifest.check_compatibility(verbose=True)
        except ext_exceptions.ManifestError as err:
            exit_handler("Invalid", str(err))
            return False
        except ext_exceptions.CompatibilityError as err:
            exit_handler("Incompatible", str(err))
            return False

        if not record.is_preview:
            record.state.save(error_type="", error_message="")  # clear any previous error

        ext_deps = ExtensionDependencies(record.id, record.path)
        extension_main = f"{record.path}/main.py"
        cmd = [sys.executable, extension_main]

        if isinstance(record, PreviewExtensionRecord) and record.with_debugger:
            cmd = [
                sys.executable,
                "-m",
                "debugpy",
                "--listen",
                f"{DEBUGPY_HOST}:{DEBUGPY_PORT}",
                "--wait-for-client",
                extension_main,
            ]

        prefs = {p_id: pref.value for p_id, pref in record.preferences.items()}
        triggers = {t_id: t.default_keyword for t_id, t in record.manifest.triggers.items() if t.default_keyword}
        # backwards compatible v2 preferences format (with keywords added back)
        v2_prefs = {**triggers, **prefs}
        env = {
            # Assume previews should always run verbose (works until we add a log-level flag)
            "VERBOSE": str(int(cli.get_args().verbose or record.is_preview)),
            "PYTHONPATH": ":".join(x for x in [paths.APPLICATION, ext_deps.get_dependencies_path()] if x),
            "EXTENSION_PREFERENCES": json.dumps(v2_prefs, separators=(",", ":")),
            "ULAUNCHER_EXTENSION_ID": record.id,
            "ULAUNCHER_INPUT_DEBOUNCE": str(record.manifest.input_debounce),
            # v2 extensions already get a compatibility-mode warning, so only warn v3 ones.
            "ULAUNCHER_API_DEPRECATION_WARNINGS": "1" if record.manifest.supports_current_api() else "0",
        }

        # socketpair/spawnv can fail for host-environment reasons (fd exhaustion, fork limits,
        # missing interpreter). Route these through exit_handler like a crash so the error is
        # surfaced consistently and any queued start listeners get cleaned up, rather than
        # raising out into callers.
        try:
            self.runtimes[record.id] = ExtensionRuntime(record.id, cmd, env, exit_handler, message_handler)
        except (OSError, GLib.Error) as err:
            exit_handler("FailedToStart", str(err))
            return False
        if listener := self.listener:
            listener.invalidate_cache()
            listener.started(record.id)
        return self.is_running(record)

    async def stop_extension(self, record: ExtensionRecord) -> None:
        """Stops the extension process if it is running."""
        if runtime := self.runtimes.pop(record.id, None):
            loop = asyncio.get_running_loop()
            stopped_future: asyncio.Future[None] = loop.create_future()

            def resolve() -> None:
                if not stopped_future.done():
                    stopped_future.set_result(None)

            # The listener fires on the GTK main thread. A plain set_result there would not wake
            # this loop, making every stop take the full timeout below.
            def on_stopped() -> None:
                if not loop.is_closed():
                    loop.call_soon_threadsafe(resolve)

            listeners = self.stopped_listeners[record.id]
            listeners.append(on_stopped)
            runtime.stop()

            try:
                await asyncio.wait_for(stopped_future, timeout=5.0)
            except asyncio.TimeoutError:
                logger.warning('Timed out waiting for extension "%s" to exit', record.id)
            finally:
                if on_stopped in listeners:
                    listeners.remove(on_stopped)

    def send_message(self, record: ExtensionRecord, message: ipc.Event, request_id: int | None = None) -> None:
        """
        Sends a JSON message to the extension if it is running.
        """
        if runtime := self.runtimes.get(record.id):
            runtime.send_message(message, request_id)

    def save_user_preferences(self, record: ExtensionRecord, data: Any) -> None:
        from ulauncher.internals.ipc import EventType

        old_preferences = {p_id: pref.value for p_id, pref in record.preferences.items()}
        record.persist_preferences(data)
        for p_id, new_value in data.get("preferences", {}).items():
            # Only notify about values changing for preferences declared in the manifest
            if p_id in old_preferences and new_value != old_preferences[p_id]:
                event_data: ipc.UpdatePreferencesEvent = {
                    "type": EventType.UPDATE_PREFERENCES,
                    "args": (p_id, new_value, old_preferences[p_id]),
                }
                self.send_message(record, event_data)

    @events.on
    def reload(
        self,
        extension_ids: list[str] | None = None,
    ) -> None:
        if not extension_ids:
            logger.warning("Reload message received without any extension IDs. No extensions will be restarted.")
            return

        logger.info("Reloading extension(s): %s", ", ".join(extension_ids))

        self._run_batch_job(
            extension_ids,
            ["stop", "start"],
            callback=lambda: logger.info("%s extensions (re)loaded", len(extension_ids)),
        )

    @events.on
    def stop(self, extension_ids: list[str] | None = None) -> None:
        if not extension_ids:
            logger.warning("Stop message received without any extension IDs. No extensions will be stopped.")
            return

        logger.info("Stopping extension(s): %s", ", ".join(extension_ids))

        self._run_batch_job(
            extension_ids, ["stop"], callback=lambda: logger.info("%s extensions stopped", len(extension_ids))
        )

    def _attach_preview_log(self, ext_id: str) -> None:
        """Mirror the previewed extension's records to a file for the CLI that started it to tail.

        In the extensions' wire format, so the CLI restores the records and applies its own colors,
        rather than parsing them back out of a formatted line.
        """
        self.detach_preview_log()
        handler = logging.FileHandler(paths.PREVIEW_LOG_FILE, mode="w", encoding="utf-8")
        handler.setFormatter(log_wire.WireFormatter())
        # Catches the sub-loggers too, since records propagate up from "<ext_id>.<name>"
        logging.getLogger(ext_id).addHandler(handler)
        self._preview_log = (ext_id, handler)

    def detach_preview_log(self) -> None:
        """Also called on app shutdown, so a CLI tailing the log isn't left waiting for an end that
        can no longer come."""
        if self._preview_log:
            ext_id, handler = self._preview_log
            logging.getLogger(ext_id).removeHandler(handler)
            handler.close()
            # Load-bearing: dropping the file is what tells the CLI that the preview is over
            Path(paths.PREVIEW_LOG_FILE).unlink(missing_ok=True)
            self._preview_log = None

    @events.on
    def preview_ext(self, ext_id: str, path: str, with_debugger: bool = False) -> None:
        """Run an extension from a dev path WITHOUT installing it. Triggered from the CLI via D-Bus.

        While the preview is active, the record with this id launches from the dev path; its
        dependencies are installed by the CLI before this is called.
        """

        logger.info("[preview] Previewing ext_id=%s path=%s debugger=%s", ext_id, path, with_debugger)
        preview = PreviewExtensionRecord(ext_id, path, with_debugger=with_debugger)
        self.preview = preview

        # Only relaunch if the current preview is still active
        def start_if_still_previewing() -> None:
            if self.preview is preview:
                self._attach_preview_log(ext_id)
                self._run_batch_job([ext_id], ["start"])

        def start_on_main_thread() -> None:
            # The batch job calls back from a worker thread. Previews are set on the main thread
            scheduling.run_when_idle(start_if_still_previewing)

        # Relaunch from the dev path, stopping any conflicting installed instances
        self._run_batch_job([ext_id], ["stop"], callback=start_on_main_thread)

    @events.on
    def stop_preview(self) -> None:
        """Stop the active preview and restore the installed extension. Triggered from the CLI via D-Bus"""
        if not self.preview:
            # will happen if preview is interrupted before started
            logger.debug("[preview] Ignoring stop_preview; no preview active")
            return

        preview = self.preview
        ext_id = preview.id
        logger.info("[preview] Stopping preview %s", ext_id)

        def restore_original() -> None:
            if self.preview is preview:
                self.preview = None
                self.detach_preview_log()
            elif self.preview and self.preview.id == ext_id:
                # Another preview started for this extension during the stop, and owns the log now
                return
            # Reload from the installed path, or drop the record if the extension was never installed.
            record = self.get(ext_id)
            if record and record.is_enabled:
                self._run_batch_job([ext_id], ["start"])

        def restore_on_main_thread() -> None:
            # The batch job calls back from a worker thread. Previews are set on the main thread
            scheduling.run_when_idle(restore_original)

        self._run_batch_job([ext_id], ["stop"], callback=restore_on_main_thread)


ext_service = ExtensionService()
