extensions.py

from __future__ import annotations

import hashlib
import importlib.util
import os
import sys
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple


EXTENSION_API_VERSION = 1


class PrefixExtensionError(Exception):
    pass


@dataclass(frozen=True)
class ExtensionMetadata:
    name: str
    version: str = "0.0.0"
    requires_api: int = EXTENSION_API_VERSION


# ---- Types ----

TypeCondition = Callable[["TypeContext", Any], int]
TypeToStr = Callable[["TypeContext", Any], str]
TypeEquals = Callable[["TypeContext", Any, Any], bool]
TypeDefault = Callable[["TypeContext"], Any]


@dataclass(frozen=True)
class TypeSpec:
    name: str
    condition_int: TypeCondition
    to_str: TypeToStr
    equals: Optional[TypeEquals] = None
    default_value: Optional[TypeDefault] = None
    printable: bool = True


@dataclass
class TypeRegistry:
    _types: Dict[str, TypeSpec] = field(default_factory=dict)
    _sealed: set[str] = field(default_factory=set)

    def seal(self, name: str) -> None:
        self._sealed.add(name)

    def register(self, spec: TypeSpec, *, seal: bool = False) -> None:
        name = spec.name
        if not name or not isinstance(name, str):
            raise PrefixExtensionError("Type name must be a non-empty string")
        if name in self._types:
            raise PrefixExtensionError(f"Type '{name}' is already defined")
        self._types[name] = spec
        if seal:
            self._sealed.add(name)

    def ensure_new(self, name: str) -> None:
        if name in self._sealed:
            raise PrefixExtensionError(f"Type '{name}' is sealed and cannot be modified")
        if name in self._types:
            raise PrefixExtensionError(f"Type '{name}' already exists")

    def has(self, name: str) -> bool:
        return name in self._types

    def get(self, name: str) -> TypeSpec:
        try:
            return self._types[name]
        except KeyError:
            raise PrefixExtensionError(f"Unknown type '{name}'")

    def get_optional(self, name: str) -> Optional[TypeSpec]:
        return self._types.get(name)

    def names(self) -> set[str]:
        return set(self._types.keys())


@dataclass(frozen=True)
class StepContext:
    step_index: int
    rule: str
    location: Any  # SourceLocation | None
    extra: Optional[Dict[str, Any]]


@dataclass
class HookRegistry:
    # event -> list[(priority, handler, ext_name)]
    _events: Dict[str, List[Tuple[int, Callable[..., None], str]]] = field(default_factory=dict)
    # list[(every_n, handler, ext_name, name)]
    _step_rules: List[Tuple[int, Callable[[Any, StepContext], None], str, str]] = field(default_factory=list)
    # list[(name, runner, ext_name)]
    _repls: List[Tuple[str, Callable[["ReplContext"], int], str]] = field(default_factory=list)

    def on_event(self, event: str, handler: Callable[..., None], *, priority: int, ext_name: str) -> None:
        self._events.setdefault(event, []).append((priority, handler, ext_name))
        self._events[event].sort(key=lambda t: t[0], reverse=True)

    def emit(self, event: str, *args: Any, **kwargs: Any) -> None:
        for _priority, handler, _ext in self._events.get(event, []):
            handler(*args, **kwargs)

    def add_step_rule(self, *, name: str, every_n: int, handler: Callable[[Any, StepContext], None], ext_name: str) -> None:
        if every_n <= 0:
            raise PrefixExtensionError("every_n_steps must be >= 1")
        self._step_rules.append((every_n, handler, ext_name, name))

    def after_step(self, interpreter: Any, ctx: StepContext) -> None:
        for every_n, handler, _ext, _name in self._step_rules:
            if ctx.step_index % every_n == 0:
                handler(interpreter, ctx)

    def register_repl(self, *, name: str, runner: Callable[["ReplContext"], int], ext_name: str) -> None:
        if not name:
            raise PrefixExtensionError("REPL name must be non-empty")
        self._repls.append((name, runner, ext_name))

    def pick_repl(self) -> Optional[Tuple[str, Callable[["ReplContext"], int], str]]:
        if not self._repls:
            return None
        if len(self._repls) == 1:
            return self._repls[0]
        names = ", ".join(sorted({n for n, _, _ in self._repls}))
        raise PrefixExtensionError(f"Multiple REPL providers registered ({names}); explicit selection not implemented")


@dataclass(frozen=True)
class TypeContext:
    interpreter: Any
    location: Any  # SourceLocation | None


@dataclass(frozen=True)
class ReplContext:
    verbose: bool
    services: "RuntimeServices"
    # factories / I/O
    make_interpreter: Callable[[str, str], Any]


@dataclass
class RuntimeServices:
    metadata: List[ExtensionMetadata] = field(default_factory=list)
    type_registry: TypeRegistry = field(default_factory=TypeRegistry)
    hook_registry: HookRegistry = field(default_factory=HookRegistry)
    # operators are registered directly into interpreter.builtins at attach time
    operators: List[Tuple[str, int, Optional[int], Callable[..., Any], str]] = field(default_factory=list)


class ExtensionAPI:
    def __init__(self, *, services: RuntimeServices, ext_name: str, asmodule: bool = False) -> None:
        self._services = services
        self._ext_name = ext_name
        self._asmodule = bool(asmodule)

    # ---- metadata ----
    def metadata(self, *, name: str, version: str = "0.0.0", requires_api: int = EXTENSION_API_VERSION) -> None:
        self._services.metadata.append(ExtensionMetadata(name=name, version=version, requires_api=requires_api))

    # ---- operators ----
    def register_operator(
        self,
        name: str,
        min_args: int,
        max_args: Optional[int],
        impl: Callable[..., Any],
        *,
        doc: str = "",
    ) -> None:
        if not name:
            raise PrefixExtensionError("Operator name must be non-empty")
        qualified = name
        if self._asmodule:
            prefix = str(self._ext_name)
            if not qualified.startswith(prefix + "."):
                qualified = f"{prefix}.{qualified}"
        self._services.operators.append((qualified, int(min_args), None if max_args is None else int(max_args), impl, doc))

    def operator(self, name: str, min_args: int, max_args: Optional[int] = None, *, doc: str = "") -> Callable[[Callable[..., Any]], Callable[..., Any]]:
        def deco(fn: Callable[..., Any]) -> Callable[..., Any]:
            self.register_operator(name, min_args, max_args, fn, doc=doc)
            return fn

        return deco

    # ---- types ----
    def register_type(
        self,
        name: str,
        *,
        condition_int: TypeCondition,
        to_str: TypeToStr,
        equals: Optional[TypeEquals] = None,
        default_value: Optional[TypeDefault] = None,
        printable: bool = True,
    ) -> None:
        self._services.type_registry.ensure_new(name)
        self._services.type_registry.register(
            TypeSpec(
                name=name,
                condition_int=condition_int,
                to_str=to_str,
                equals=equals,
                default_value=default_value,
                printable=printable,
            )
        )

    def type(self, name: str, *, printable: bool = True) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
        def deco(fn: Callable[..., Any]) -> Callable[..., Any]:
            raise PrefixExtensionError("Use register_type(...) explicitly; decorator form not supported")

        return deco

    # ---- hooks ----
    def on_event(self, event: str, handler: Optional[Callable[..., None]] = None, *, priority: int = 0) -> Callable[..., None] | Callable[[Callable[..., None]], Callable[..., None]]:
        if handler is None:
            def deco(fn: Callable[..., None]) -> Callable[..., None]:
                self._services.hook_registry.on_event(event, fn, priority=priority, ext_name=self._ext_name)
                return fn
            return deco
        self._services.hook_registry.on_event(event, handler, priority=priority, ext_name=self._ext_name)
        return handler

    def every_n_steps(self, every_n: int, handler: Optional[Callable[[Any, StepContext], None]] = None, *, name: str = "") -> Callable[[Any, StepContext], None] | Callable[[Callable[[Any, StepContext], None]], Callable[[Any, StepContext], None]]:
        if handler is None:
            def deco(fn: Callable[[Any, StepContext], None]) -> Callable[[Any, StepContext], None]:
                self._services.hook_registry.add_step_rule(name=name or fn.__name__, every_n=every_n, handler=fn, ext_name=self._ext_name)
                return fn
            return deco
        self._services.hook_registry.add_step_rule(name=name or handler.__name__, every_n=every_n, handler=handler, ext_name=self._ext_name)
        return handler

    # ---- repl ----
    def register_repl(self, *, name: str, runner: Callable[[ReplContext], int]) -> None:
        self._services.hook_registry.register_repl(name=name, runner=runner, ext_name=self._ext_name)


def _extension_search_roots() -> List[str]:
    """Return directories to search for the interpreter's built-in `ext/`.

    When running as a frozen executable (e.g. PyInstaller), `__file__` points to
    the extracted bundle, but users often ship extensions alongside the exe.
    Prefer the exe directory, then the bundled extraction dir, then source dir.
    """

    roots: List[str] = []
    if getattr(sys, "frozen", False):
        roots.append(os.path.dirname(os.path.abspath(sys.executable)))
        meipass = getattr(sys, "_MEIPASS", None)
        if isinstance(meipass, str) and meipass:
            roots.append(os.path.abspath(meipass))
    roots.append(os.path.dirname(os.path.abspath(__file__)))
    # de-dup while preserving order
    out: List[str] = []
    seen: set[str] = set()
    for r in roots:
        if r not in seen:
            out.append(r)
            seen.add(r)
    return out


def _resolve_in_builtin_ext(path: str) -> Optional[str]:
    if os.path.isabs(path):
        return None
    for root in _extension_search_roots():
        candidate = os.path.join(root, "ext", path)
        if os.path.exists(candidate):
            return os.path.abspath(candidate)
    return None


def _unique_module_name(path: str) -> str:
    base = os.path.basename(path)
    digest = hashlib.sha256(os.path.abspath(path).encode("utf-8")).hexdigest()[:12]
    safe = "".join(ch if ch.isalnum() else "_" for ch in base)
    return f"prefix_ext_{safe}_{digest}"


def load_extension_module(path: str) -> Any:
    # If the given path doesn't exist, and it is a relative path, try the
    # interpreter's own `ext` directory as a fallback before failing.
    if not os.path.exists(path):
        if not os.path.isabs(path):
            resolved = _resolve_in_builtin_ext(path)
            if resolved is not None:
                path = resolved
            else:
                raise PrefixExtensionError(f"Extension not found: {path}")
        else:
            raise PrefixExtensionError(f"Extension not found: {path}")
    mod_name = _unique_module_name(path)
    spec = importlib.util.spec_from_file_location(mod_name, path)
    if spec is None or spec.loader is None:
        raise PrefixExtensionError(f"Failed to load extension module: {path}")
    module = importlib.util.module_from_spec(spec)

    # Insert the module into sys.modules before execution so that code
    # running during import (e.g. dataclass processing) can look up the
    # module by name. Clean up/restore on failure.
    prev_mod = sys.modules.get(mod_name)
    sys.modules[mod_name] = module

    # Let extensions import siblings by temporarily prepending their directory.
    ext_dir = os.path.dirname(os.path.abspath(path))
    sys.path.insert(0, ext_dir)
    try:
        try:
            spec.loader.exec_module(module)  # type: ignore[union-attr]
        except Exception as exc:
            # restore previous sys.modules state
            if prev_mod is None:
                sys.modules.pop(mod_name, None)
            else:
                sys.modules[mod_name] = prev_mod
            # Never allow a Python traceback to escape due to extension code.
            raise PrefixExtensionError(
                f"Failed to load extension module {path}: {exc.__class__.__name__}: {exc}"
            )
    finally:
        if sys.path and sys.path[0] == ext_dir:
            sys.path.pop(0)
    return module


def read_prex(pointer_file: str) -> List[str]:
    if not os.path.exists(pointer_file):
        raise PrefixExtensionError(f".pre file not found: {pointer_file}")
    base_dir = os.path.dirname(os.path.abspath(pointer_file))
    out: List[str] = []
    with open(pointer_file, "r", encoding="utf-8") as handle:
        for raw in handle:
            line = raw.strip()
            if not line or line.startswith("#"):
                continue
            # Allow inline comments: path # comment
            if "#" in line:
                line = line.split("#", 1)[0].strip()
                if not line:
                    continue
            # Resolve relative entries: prefer pointer-file base dir, then
            # cwd, then interpreter's ext/ directory as a final fallback.
            candidate = line
            if not os.path.isabs(candidate): # relative path
                candidate_base = os.path.abspath(os.path.join(base_dir, candidate))
                if os.path.exists(candidate_base): # found relative to pointer file
                    out.append(candidate_base)
                    continue
                candidate_cwd = os.path.abspath(candidate)
                if os.path.exists(candidate_cwd): # found relative to cwd
                    out.append(candidate_cwd)
                    continue
                resolved = _resolve_in_builtin_ext(candidate)
                if resolved is not None:
                    out.append(resolved)
                    continue
                raise PrefixExtensionError(f"Extension referenced in {pointer_file} not found: {candidate}")
            else: # absolute path
                if os.path.exists(candidate): # found as absolute path
                    out.append(os.path.abspath(candidate))
                else:
                    raise PrefixExtensionError(f"Extension referenced in {pointer_file} not found: {candidate}")
    return out


def gather_extension_paths(paths: Sequence[str]) -> List[str]:
    expanded: List[str] = []
    for p in paths:
        if p.lower().endswith(".prex"):
            # Resolve pointer file path (absolute if possible) and expand it.
            pointer = p
            if not os.path.isabs(pointer):
                pointer = os.path.abspath(pointer)
            expanded.extend(read_prex(pointer))
        else:
            # If a relative extension path is provided and doesn't exist in
            # the current working directory, try the interpreter's `ext`
            # directory as a fallback.
            if not os.path.isabs(p) and not os.path.exists(p):
                resolved = _resolve_in_builtin_ext(p)
                if resolved is not None:
                    expanded.append(resolved)
                    continue
            expanded.append(p)
    # normalize
    return [os.path.abspath(p) for p in expanded]


def build_default_services() -> RuntimeServices:
    services = RuntimeServices()
    # Reserve built-in type names so extensions cannot redefine them.
    # The interpreter will register their concrete semantics at runtime.
    services.type_registry.seal("INT")
    services.type_registry.seal("FLT")
    services.type_registry.seal("STR")
    services.type_registry.seal("TNS")
    services.type_registry.seal("MAP")
    services.type_registry.seal("FUNC")
    return services


def load_runtime_services(paths: Sequence[str]) -> RuntimeServices:
    services = build_default_services()
    resolved = gather_extension_paths(paths)
    for path in resolved:
        module = load_extension_module(path)
        api_version = getattr(module, "PREFIX_EXTENSION_API_VERSION", EXTENSION_API_VERSION)
        if api_version != EXTENSION_API_VERSION:
            raise PrefixExtensionError(
                f"Extension {path} requires API {api_version}, host supports {EXTENSION_API_VERSION}"
            )
        register = getattr(module, "prefix_register", None)
        if register is None or not callable(register):
            raise PrefixExtensionError(f"Extension {path} must define callable prefix_register(ext)")
        ext_name = getattr(module, "PREFIX_EXTENSION_NAME", os.path.splitext(os.path.basename(path))[0])
        ext_asmodule = bool(getattr(module, "PREFIX_EXTENSION_ASMODULE", False))
        ext = ExtensionAPI(services=services, ext_name=str(ext_name), asmodule=ext_asmodule)
        register(ext)
    return services