"""
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 base64
import os
import subprocess
import sys
import json

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

try:
    APP_DATA_DIR = Path("/tmp/.zarch") if Path("/zarch_environment").read_text().strip() in ["control-plane", "ext-runner"] else Path.home() / ".zarch"
except FileNotFoundError:
    APP_DATA_DIR = Path.home() / ".zarch"


def _plugins_base_dir() -> Path:
    return (APP_DATA_DIR / "plugins").resolve()


def _resolve_connected_repo() -> tuple[str | None, str | None]:
    # Canonical env names used by the control-plane job.
    repo = os.getenv("ZARCH_GH_REPO", "").strip()
    ref = os.getenv("ZARCH_GH_REF", "").strip()

    return (repo or None, ref or None)


def _system_python_tag() -> str | None:
    marker = Path("/zarch_environment")
    if not marker.exists():
        return None

    try:
        runtime_mode = marker.read_text().strip()
    except Exception:
        return None

    if runtime_mode not in {"control-plane", "ext-runner"}:
        return None

    try:
        out = subprocess.check_output(
            ["/usr/bin/python3", "-c", "import sys;print(f'{sys.version_info.major}{sys.version_info.minor}')"],
            text=True,
        ).strip()
    except Exception:
        return None

    if not out:
        return None
    return f"py{out}"


def _plugin_dirs_for_discovery() -> list[Path]:
    base = _plugins_base_dir()
    tags = [f"py{sys.version_info.major}{sys.version_info.minor}"]
    system_tag = _system_python_tag()
    if system_tag and system_tag not in tags:
        tags.append(system_tag)

    plugin_dirs: list[Path] = []
    for tag in tags:
        plugin_dir = (base / tag).resolve()
        plugin_dir.mkdir(parents=True, exist_ok=True)
        plugin_dirs.append(plugin_dir)

    # Backward compatibility for users that installed extensions before
    # per-interpreter plugin paths were introduced.
    if base.exists() and base not in plugin_dirs:
        plugin_dirs.append(base)
    return plugin_dirs


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

    def _discover_installed_extensions(self) -> None:
        plugin_dirs = _plugin_dirs_for_discovery()

        # Keep per-interpreter plugins ahead of legacy shared plugins.
        for plugin_dir in reversed(plugin_dirs):
            if str(plugin_dir) not in sys.path:
                sys.path.insert(0, str(plugin_dir))

        seen_entry_points: set[tuple[str, str, str]] = set()
        for plugin_dir in plugin_dirs:
            for dist in distributions(path=[str(plugin_dir)]):
                for ep in dist.entry_points:
                    if ep.group != "zarch.extensions":
                        continue

                    ep_key = (ep.group, ep.name, ep.value)
                    if ep_key in seen_entry_points:
                        continue
                    seen_entry_points.add(ep_key)

                    try:
                        extension_class = ep.load()
                        self.loaded_extensions.append(extension_class())
                    except Exception:
                        continue

    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

        if not extensions_map:
            return

        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")

        repo, branch = _resolve_connected_repo()

        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": repo,
                "github_branch": branch,
            }

        payload_json = json.dumps(payload, separators=(",", ":"), ensure_ascii=False)
        payload_b64 = base64.urlsafe_b64encode(payload_json.encode("utf-8")).decode("ascii").rstrip("=")

        out, code = project_context.gcloud([
            "run",
            "jobs",
            "execute",
            "zarch-extension-runner",
            "--region", cp_region,
            "--project", project_context.id,
            f"--update-env-vars=ZARCH_EXTENSION_PAYLOAD_B64={payload_b64}",
            "--wait",
        ])

        if code != 0:
            detail = (out or "").strip()
            if len(detail) > 800:
                detail = detail[:800] + "...(truncated)"
            raise RuntimeError(f"Failed to execute extension runner job: {detail or 'unknown error'}")
