"""
Z-Arch Extension Registry
-----------------------------------------------
Discovers, loads, and dispatches Z-Arch extensions registered via Python entry points.

The ExtensionRegistry is responsible for:
- Discovering installed extension packages
- Determining which extensions claim specific extension blocks
- Dispatching lifecycle events to extension hooks

Execution locality is enforced by Z-Arch core policy. Hooks classified as remote
are never executed in this process. Instead, their invocation is forwarded to an
isolated extension runner environment.

Copyright © 2026 RAM Cloud Code LLC
Licensed under the Apache License, Version 2.0.
"""
from __future__ import annotations
from typing import Any, Collection, Dict, List
from importlib.metadata import distributions
from pathlib import Path
import os
import sys
import json

from zarch.extensions.base import ZArchExtension
from zarch.extensions.context import ZArchContext
from zarch.extensions.policy import HOOK_EXECUTION_POLICY

APP_DATA_DIR = Path.home() / ".zarch" if os.getenv("ZARCH_CTX") not in ["control-plane", "ext-runner"] else Path("/tmp/.zarch")

class ExtensionRegistry:
    def __init__(self) -> None:
        self.loaded_extensions: List[ZArchExtension] = []
        self._discover_installed_extensions()

    def _discover_installed_extensions(self) -> None:
        plugin_dir = (APP_DATA_DIR / "plugins").resolve()
        plugin_dir.mkdir(parents=True, exist_ok=True)

        if str(plugin_dir) not in sys.path:
            sys.path.insert(0, str(plugin_dir))

        for dist in distributions(path=[str(plugin_dir)]):
            for ep in dist.entry_points:
                if ep.group != "zarch.extensions":
                    continue
                extension_class = ep.load()
                self.loaded_extensions.append(extension_class())

    def get_claiming_extensions(
        self,
        extension_name: str,
        extension_block: Dict[str, Any]
    ) -> List[ZArchExtension]:
        claiming_list = []
        for extension in self.loaded_extensions:
            try:
                if extension.claim(extension_name, extension_block):
                    claiming_list.append(extension)
            except Exception:
                continue
        return claiming_list

    def dispatch_lifecycle_event(
        self,
        event_name: str,
        project_context: ZArchContext,
        config: Dict[str, Any],
        *,
        force_local: bool = False,
        extension_names: Collection[str] | None = None,
        event_data: Dict[str, Any] | None = None,
    ) -> None:
        """
        Dispatch a lifecycle event to extensions according to Z-Arch execution policy.

        Lifecycle hooks are executed in one of two modes:
        - Local: hooks are executed directly in the current process
        - Remote: hooks are not executed here and are instead forwarded to an isolated
        extension runner environment

        The execution mode is determined solely by Z-Arch core policy and cannot be
        influenced by extensions. When dispatched remotely, extension hook methods are
        never called in this process.

        The force_local flag is used internally by the extension runner to ensure hooks
        execute locally within the isolated environment and do not recurse into another
        remote dispatch.
        """

        execution = "local" if force_local else HOOK_EXECUTION_POLICY.get(event_name, "local")
        extensions_map_raw: Dict[str, Dict[str, Any]] = config.get("extensions", {}) or {}
        if not isinstance(extensions_map_raw, dict):
            extensions_map_raw = {}

        if extension_names:
            selected = set(extension_names)
            extensions_map = {
                name: block
                for name, block in extensions_map_raw.items()
                if name in selected
            }
        else:
            extensions_map = extensions_map_raw

        previous_event_data = project_context.get_event_data()
        project_context._set_event_data(event_data)
        try:
            if execution == "remote":
                self._invoke_extension_runner(
                    project_context,
                    event_name,
                    extensions_map,
                    event_data=event_data,
                )
                return

            for extension_name, extension_block in extensions_map.items():
                if not isinstance(extension_block, dict):
                    continue
                claiming_extensions = self.get_claiming_extensions(extension_name, extension_block)
                config_for_extension = extension_block.get("config", {})

                for extension_instance in claiming_extensions:
                    hook = getattr(extension_instance, event_name, None)
                    if callable(hook):
                        hook(project_context, config_for_extension)
        finally:
            project_context._set_event_data(previous_event_data)

    # Remote execution
    # ----------------------------------------------------------
    def _invoke_extension_runner(
        self,
        project_context: ZArchContext,
        event_name: str,
        extensions_map: Dict[str, Dict[str, Any]],
        *,
        event_data: Dict[str, Any] | None = None,
    ) -> None:
        """
        Invoke the Z-Arch extension runner in an isolated execution environment.

        This method is used exclusively for lifecycle hooks classified as remote by
        Z-Arch policy. It launches a dedicated execution environment and passes a
        serialized payload describing the lifecycle event and extension configuration.

        The runner environment contains the full Z-Arch CLI and runtime but executes
        outside of sensitive control-plane processes.
        """
        cp_region = os.getenv("ZARCH_CP_REGION")
        if not cp_region:
            raise RuntimeError("Missing ZARCH_CP_REGION")

        payload =  {
                "event_name": event_name,
                "project_id": project_context.id,
                "region": project_context.region,
                "project_root_path": str(project_context.project_root_path),
                "non_interactive": True,
                "extensions": extensions_map,
                "event_data": event_data,
                "github_repo": os.getenv("ZARCH_GITHUB_REPO"),
                "github_branch": os.getenv("ZARCH_GITHUB_BRANCH"),
            }
            
        project_context.gcloud([
            "run",
            "jobs",
            "execute",
            "zarch-extension-runner",
            "--region", cp_region,
            f"ZARCH_EXTENSION_PAYLOAD={json.dumps(payload)}",
            "--wait",
        ])
