operation.py

from abc import ABC, abstractmethod
from enum import Enum
import json
import logging
from typing import Type, Dict, Optional, Any, List, ClassVar, Union
from pydantic import BaseModel, Field, ConfigDict
from pydantic_settings import BaseSettings
from rdflib.term import Node
from rdflib import URIRef, Literal, BNode, Graph
from rdflib.namespace import XSD
from rdflib.query import Result
from web_algebra.exceptions import (
    InvalidFormError,
    UnknownOperationError,
    VariableNotFoundError,
)


# JSON-LD keyword set used to recognise a dict as RDF data (a JSON-LD
# document/fragment) rather than generic JSON to recurse into. Presence of any
# of these is a strong, unambiguous signal — they are JSON-LD reserved terms
# with no legitimate meaning in non-RDF JSON.
_JSONLD_KEYS = ("@context", "@graph", "@id", "@type")


class OperationKind(str, Enum):
    """What an operation does to the world.

    The ordering is by blast radius, so a composite form (a `ForEach` whose body
    `PATCH`es) takes the kind of the worst thing it can do — `max()` over the
    kinds of the operations it contains.
    """

    READ = "read"
    WRITE = "write"
    DESTRUCTIVE = "destructive"

    @property
    def severity(self) -> int:
        return _KIND_SEVERITY[self]

    def __lt__(self, other: "OperationKind") -> bool:
        return self.severity < other.severity


_KIND_SEVERITY = {
    OperationKind.READ: 0,
    OperationKind.WRITE: 1,
    OperationKind.DESTRUCTIVE: 2,
}


class Operation(ABC, BaseModel):
    """
    Abstract base class for all operations with dual execution paths:
    1. execute() - Pure RDFLib function
    2. execute_json() - JSON argument processing

    Operations can optionally implement MCPTool interface for MCP client access.
    """

    registry: ClassVar[Dict[str, Type["Operation"]]] = {}

    #: What this operation does to the world, for the plan summary a client shows
    #: before approving a document. Reading is the default because most operations
    #: only compute; the ones that issue a mutating request say so themselves.
    kind: ClassVar[OperationKind] = OperationKind.READ

    settings: BaseSettings = Field(exclude=True)
    #: The iteration focus an enclosing ForEach established, empty outside one.
    #: A `default_factory`, not a literal: a bare `{}` is one dict shared by every
    #: instance ever constructed, so two executions in one process — which is what
    #: the HTTP service runs — would see each other's focus.
    context: Any = Field(default_factory=dict)

    model_config = ConfigDict(extra="allow")

    @classmethod
    def name(cls) -> str:
        return cls.__name__

    @classmethod
    @abstractmethod
    def description(cls) -> str:
        pass

    @classmethod
    @abstractmethod
    def inputSchema(cls) -> dict:
        pass

    @abstractmethod
    def execute(self, *args) -> Union[Node, Result, Graph]:
        """Pure function: RDFLib terms → RDFLib terms/Results/Graphs"""
        pass

    @abstractmethod
    def execute_json(
        self, arguments: dict, variable_stack: list = None
    ) -> Union[Node, Result, Graph]:
        """JSON execution: processes JSON args, returns RDFLib objects"""
        pass

    @classmethod
    def register(cls, operation_cls: Type["Operation"]) -> None:
        if not issubclass(operation_cls, cls):
            raise ValueError(
                f"Cannot register {operation_cls}: Must be a subclass of Operation."
            )
        cls.registry[operation_cls.name()] = operation_cls
        logging.info(f"Registered operation: {operation_cls.name()}")

    @classmethod
    def list_operations(cls) -> List[Type["Operation"]]:
        return list(cls.registry.values())

    @classmethod
    def get(cls, name: str) -> Optional[Type["Operation"]]:
        return cls.registry.get(name)

    @classmethod
    def process_json(
        cls,
        settings: BaseSettings,
        json_data: Any,
        context: dict = None,
        variable_stack: list = None,
    ) -> Any:
        """Class method for processing JSON with @op structures"""
        if context is None:
            context = {}
        if variable_stack is None:
            variable_stack = []
        if isinstance(json_data, dict):
            if "@op" in json_data:
                op_name = json_data["@op"]
                op_args = json_data.get("args", {})

                operation_cls = cls.get(op_name)
                if not operation_cls:
                    raise UnknownOperationError(f"Unknown operation: {op_name}")

                operation = operation_cls(settings=settings, context=context)
                result = operation.execute_json(op_args, variable_stack)

                # Return RDFLib objects as-is for operation chaining
                return result

            # URI reference form (formal-semantics.md §2.2 rule 2): an object
            # whose ONLY member is `@id` evaluates to a URI. This is JSON-LD's
            # node-reference syntax — a bare node reference carries no triples,
            # so reusing it as the URI form is unambiguous. Inside an RDF data
            # form this rule never fires: `_resolve_jsonld` walks those without
            # re-entering this dispatch for non-`@op` objects.
            if set(json_data.keys()) == {"@id"}:
                inner = cls.process_json(
                    settings, json_data["@id"], context, variable_stack
                )
                return URIRef(str(inner))

            # JSON-LD shape recognition — a dict carrying any JSON-LD reserved
            # key is RDF data (a JSON-LD document or fragment), not generic
            # JSON to recurse into. It may still embed `@op` operations or
            # variable references in its values (e.g. an `@id` computed from a
            # runtime binding), so resolve those in place but keep the dict
            # intact and let the consuming op parse it.
            #
            # We deliberately do NOT parse to an `rdflib.Graph` here. Every op
            # that consumes a JSON-LD body (POST, PUT, Merge, ldh Create*/Add*)
            # already parses a dict itself — and does so with the right base
            # IRI (`publicID=<target url>`), which a central parse here cannot
            # know. Parsing centrally would also freeze any unresolved `@op`
            # holes: `@id` expects a string IRI, so an op-valued `@id` would
            # silently become a blank node and the intended URI would be lost.
            if any(k in json_data for k in _JSONLD_KEYS):
                return cls._resolve_jsonld(
                    settings, json_data, context, variable_stack
                )

            # 🔁 Recurse into each value — allows nested @op inside generic
            # JSON structures (e.g. SPARQL binding objects), without
            # collapsing pure-data dicts that JSON-LD detection above
            # already handled.
            return {
                k: cls.process_json(settings, v, context, variable_stack)
                for k, v in json_data.items()
            }

        elif isinstance(json_data, list):
            # Sequence form (formal-semantics.md §3.2): elements evaluate in
            # order in a fresh variable scope, so a Variable bound in step N is
            # visible to steps N+1.. and to nested forms, and goes out of scope
            # when the sequence ends.
            variable_stack.append({})
            try:
                items: list = []
                for item in json_data:
                    cls.concatenate(
                        items,
                        cls.process_json(settings, item, context, variable_stack),
                    )
                return items
            finally:
                variable_stack.pop()

        else:
            # Convert plain values to RDFLib terms
            return cls.json_to_rdflib(json_data)

    @staticmethod
    def concatenate(items: list, value: Any) -> None:
        """Append `value` to `items` as XDM concatenation does
        (formal-semantics.md §3.1): a sequence contributes its items, Unit
        (`None`) contributes none, anything else — a `Result` included — is
        one item."""
        if value is None:
            return
        if isinstance(value, list):
            for item in value:
                Operation.concatenate(items, item)
            return
        items.append(value)

    @classmethod
    def _resolve_jsonld(
        cls,
        settings: BaseSettings,
        json_data: Any,
        context: dict = None,
        variable_stack: list = None,
    ) -> Any:
        """Resolve embedded `@op` nodes inside a JSON-LD document in place.

        Walks a JSON-LD-shaped structure evaluating only `@op` nodes (to their
        RDFLib terms — `str` subclasses for URIs/literals, so they serialise
        cleanly) while leaving every other scalar as plain JSON for the
        consuming op's JSON-LD parser to interpret. Nested dicts/lists keep
        their structure so the whole document serialises as one unit; in
        particular, fragments that merely look like JSON-LD (e.g. a
        `{"@id": "..."}` object reference) stay dicts rather than being parsed
        into standalone Graphs.
        """
        if isinstance(json_data, dict):
            if "@op" in json_data:
                return cls.process_json(settings, json_data, context, variable_stack)
            return {
                k: cls._resolve_jsonld(settings, v, context, variable_stack)
                for k, v in json_data.items()
            }
        elif isinstance(json_data, list):
            return [
                cls._resolve_jsonld(settings, item, context, variable_stack)
                for item in json_data
            ]
        else:
            return json_data

    # Variable stack management methods
    def push_variable_scope(self, variable_stack: list):
        """Create a new variable scope (like entering a new XSLT template)."""
        variable_stack.append({})

    def pop_variable_scope(self, variable_stack: list):
        """Exit the current variable scope (like leaving an XSLT template)."""
        if variable_stack:
            variable_stack.pop()

    def set_variable(self, name: str, value: Any, variable_stack: list):
        """Set a variable in the current scope."""
        if not variable_stack:
            variable_stack.append({})
        variable_stack[-1][name] = value

    def get_variable(self, name: str, variable_stack: list) -> Any:
        """Get a variable value, searching from innermost to outermost scope."""
        for scope in reversed(variable_stack):
            if name in scope:
                return scope[name]
        raise VariableNotFoundError(f"Variable '{name}' not found")

    # Conversion helpers between different formats
    @staticmethod
    def to_graph(data: Any, *, base: Optional[str] = None) -> Graph:
        """Convert a JSON-LD document into an `rdflib.Graph`.

        The single conversion point from JSON-LD transport into the rdflib
        type system. Operations resolve their arguments to plain data (via
        `process_json`) and then call this to hand `execute` a `Graph` — so
        `execute` always works in rdflib terms, never raw JSON. Input that is
        already a `Graph` (e.g. the output of an upstream op such as CONSTRUCT)
        passes through unchanged.

        Both shapes of JSON-LD document are accepted: a single node object
        (`dict`) and a top-level array of nodes (`list`) — the latter is what
        SPARQL DESCRIBE/CONSTRUCT over multiple subjects returns. rdflib's
        JSON-LD parser handles both natively.

        `base` is the document's base IRI — typically the target URL of the
        enclosing op — used to resolve any relative IRIs in the JSON-LD.
        """
        if isinstance(data, Graph):
            return data
        if isinstance(data, (dict, list)):
            graph = Graph()
            graph.parse(data=json.dumps(data), format="json-ld", publicID=base)
            return graph
        raise TypeError(
            f"Cannot convert {type(data).__name__} to Graph; "
            "expected a JSON-LD dict/list or an rdflib.Graph"
        )

    @staticmethod
    def json_to_rdflib(data) -> Node:
        """Convert JSON/binding objects to RDFLib terms"""
        if data is None:
            # formal-semantics.md §2.2: null is not a valid form.
            raise InvalidFormError("null is not a valid Web Algebra form")
        if isinstance(data, dict) and "type" in data and "value" in data:
            # SPARQL binding object - values may have been processed to RDFLib terms
            type_str = str(data["type"])  # Convert potential Literal to string
            value_str = str(data["value"])  # Convert potential Literal to string

            if type_str == "uri":
                return URIRef(value_str)
            elif type_str == "literal":
                datatype = data.get("datatype")
                lang = data.get("xml:lang")
                if datatype:
                    datatype = URIRef(str(datatype))
                if lang:
                    lang = str(lang)
                return Literal(value_str, datatype=datatype, lang=lang)
            elif type_str == "bnode":
                return BNode(value_str)
            else:
                raise ValueError(f"Unknown binding type: {type_str}")
        elif isinstance(data, (URIRef, Literal, BNode)):
            # Already RDFLib term
            return data
        elif hasattr(data, "__class__") and data.__class__.__name__ in [
            "JSONResult",
            "Result",
        ]:
            # Result should pass through unchanged
            return data
        elif isinstance(data, str):
            # Plain string → always convert to string literal
            return Literal(data, datatype=XSD.string)
        elif isinstance(data, bool):
            # bool before int — bool is an int subclass in Python
            return Literal(data, datatype=XSD.boolean)
        elif isinstance(data, int):
            return Literal(data, datatype=XSD.integer)
        elif isinstance(data, float):
            return Literal(data, datatype=XSD.double)
        else:
            # Default: convert to string literal
            return Literal(str(data), datatype=XSD.string)

    @staticmethod
    def plain_to_rdflib(value: Any) -> Node:
        """Convert plain Python values to RDFLib terms for MCP interface"""
        if isinstance(value, str):
            # Plain string → always convert to string literal
            return Literal(value, datatype=XSD.string)
        elif isinstance(value, bool):
            # bool before int — bool is an int subclass in Python
            return Literal(value, datatype=XSD.boolean)
        elif isinstance(value, int):
            return Literal(value, datatype=XSD.integer)
        elif isinstance(value, float):
            return Literal(value, datatype=XSD.double)
        else:
            return Literal(str(value), datatype=XSD.string)

    @staticmethod
    def is_string_literal(term: Any) -> bool:
        """True for a language-tag-free string literal — a SPARQL simple
        literal or its RDF 1.1 equivalent, an xsd:string literal."""
        return (
            isinstance(term, Literal)
            and term.language is None
            and (term.datatype is None or term.datatype == XSD.string)
        )

    @staticmethod
    def to_string_literal(term: Node) -> Literal:
        """Convert Literal terms to string-compatible literals, following SPARQL semantics"""
        if isinstance(term, Literal):
            # Both xsd:string and rdf:langString are string-compatible in SPARQL
            if term.datatype == XSD.string:
                return term  # Already xsd:string, return as-is
            elif term.language is not None:
                return term  # rdf:langString (datatype=None, lang=xx), return as-is (compatible with string operations)
            elif term.datatype is None and term.language is None:
                # Plain literal without datatype or language - treat as string
                return term
            else:
                # Other literal datatypes need explicit conversion
                raise TypeError(
                    f"Cannot implicitly convert {term.datatype} to string. Use Str() operation for explicit casting."
                )
        else:
            # URIRef, BNode, etc. require explicit casting
            raise TypeError(
                f"Cannot implicitly convert {type(term).__name__} to string. Use Str() operation for explicit casting."
            )

    @staticmethod
    def rdflib_to_plain(term: Node) -> Any:
        """Convert RDFLib terms to plain Python values for MCP interface"""
        if isinstance(term, URIRef):
            return str(term)
        elif isinstance(term, Literal):
            # Convert based on datatype
            if term.datatype == XSD.integer:
                return int(term)
            elif term.datatype == XSD.double or term.datatype == XSD.float:
                return float(term)
            elif term.datatype == XSD.boolean:
                return term.toPython()
            else:
                return str(term)
        elif isinstance(term, BNode):
            return str(term)
        else:
            return str(term)