reason_first_program/stub.py

"""
Stub definitions and the @reason_first decorator.

A stub is a function skeleton with:
  - A decorator @reason_first(spec="...") that declares the informal specification
  - An optional #> comment inside the body with additional constraints
  - Surrounding context: imports, type hints, downstream uses

The decorator captures all of this metadata to define the valid program space.
"""

from __future__ import annotations

import ast
import inspect
import textwrap
import hashlib
from dataclasses import dataclass, field
from typing import Any, Callable, Optional
from pathlib import Path


@dataclass
class StubConstraints:
    """Formal constraints derived from the stub's surrounding context."""

    imports: list[str] = field(default_factory=list)
    type_hints: dict[str, str] = field(default_factory=dict)
    return_type: Optional[str] = None
    inline_specs: list[str] = field(default_factory=list)  # #> comments
    downstream_uses: list[str] = field(default_factory=list)
    decorator_spec: str = ""
    preconditions: list[str] = field(default_factory=list)
    postconditions: list[str] = field(default_factory=list)

    def to_prompt_context(self) -> str:
        """Render constraints as natural language for LLM prompting."""
        parts = []
        if self.decorator_spec:
            parts.append(f"Specification: {self.decorator_spec}")
        if self.inline_specs:
            parts.append(f"Inline constraints: {'; '.join(self.inline_specs)}")
        if self.type_hints:
            hints = ", ".join(f"{k}: {v}" for k, v in self.type_hints.items())
            parts.append(f"Type hints: {hints}")
        if self.return_type:
            parts.append(f"Return type: {self.return_type}")
        if self.imports:
            parts.append(f"Available imports: {', '.join(self.imports)}")
        if self.preconditions:
            parts.append(f"Preconditions: {'; '.join(self.preconditions)}")
        if self.postconditions:
            parts.append(f"Postconditions: {'; '.join(self.postconditions)}")
        return "\n".join(parts)


@dataclass
class Stub:
    """
    A program stub: a function skeleton that defines a valid program space.

    The stub captures:
      - name: function name
      - source: the full source code of the decorated function
      - constraints: formal constraints extracted from context
      - module_source: the full module source (for imports, downstream uses)
      - test_inputs: optional test inputs for execution-based validation
    """

    name: str
    source: str
    signature: str
    constraints: StubConstraints
    module_source: Optional[str] = None
    test_inputs: Optional[list[dict[str, Any]]] = None
    test_outputs: Optional[list[Any]] = None
    property_tests: Optional[list[Callable]] = None

    @property
    def stub_id(self) -> str:
        """Unique identifier for this stub."""
        content = f"{self.name}:{self.signature}:{self.constraints.decorator_spec}"
        return hashlib.sha256(content.encode()).hexdigest()[:12]

    def to_completion_prompt(self, style: str = "direct") -> str:
        """
        Generate an LLM prompt for completing this stub.

        Args:
            style: 'direct' for straightforward completion,
                   'diverse' for encouraging algorithmic diversity (SFS scattering),
                   'concept_guided' for concept-directed generation
        """
        if style == "direct":
            return self._direct_prompt()
        elif style == "diverse":
            return self._diverse_prompt()
        elif style == "concept_guided":
            return self._concept_guided_prompt()
        else:
            raise ValueError(f"Unknown prompt style: {style}")

    def _direct_prompt(self) -> str:
        return (
            f"Complete the following Python function. "
            f"Only output the function body.\n\n"
            f"{self.constraints.to_prompt_context()}\n\n"
            f"```python\n{self.source}\n```"
        )

    def _diverse_prompt(self) -> str:
        """SFS-inspired scattering: ask for diverse algorithmic directions first."""
        return (
            f"Consider this Python function stub:\n\n"
            f"```python\n{self.source}\n```\n\n"
            f"{self.constraints.to_prompt_context()}\n\n"
            f"First, list 3-5 fundamentally different algorithmic approaches "
            f"to implement this function (e.g., iterative vs recursive, "
            f"different data structures, different time/space tradeoffs).\n\n"
            f"Then implement ONE of these approaches. Choose a different approach "
            f"than you normally would. Only output the function body."
        )

    def _concept_guided_prompt(self) -> str:
        return (
            f"Complete this function stub:\n\n"
            f"```python\n{self.source}\n```\n\n"
            f"{self.constraints.to_prompt_context()}\n\n"
            f"Implement the function. Only output the function body."
        )


class StubRegistry:
    """Global registry of all stubs defined with @reason_first."""

    _stubs: dict[str, Stub] = {}

    @classmethod
    def register(cls, stub: Stub) -> None:
        cls._stubs[stub.stub_id] = stub

    @classmethod
    def get(cls, stub_id: str) -> Optional[Stub]:
        return cls._stubs.get(stub_id)

    @classmethod
    def list_all(cls) -> list[Stub]:
        return list(cls._stubs.values())

    @classmethod
    def clear(cls) -> None:
        cls._stubs.clear()


def _extract_inline_specs(source: str) -> list[str]:
    """Extract #> comments from source code."""
    specs = []
    for line in source.split("\n"):
        stripped = line.strip()
        if stripped.startswith("#>"):
            specs.append(stripped[2:].strip())
    return specs


def _extract_imports(module_source: str) -> list[str]:
    """Extract import statements from module source."""
    imports = []
    try:
        tree = ast.parse(module_source)
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    imports.append(alias.name)
            elif isinstance(node, ast.ImportFrom):
                module = node.module or ""
                for alias in node.names:
                    imports.append(f"{module}.{alias.name}")
    except SyntaxError:
        pass
    return imports


def _extract_type_hints(func: Callable) -> tuple[dict[str, str], Optional[str]]:
    """Extract type hints from function signature."""
    hints = {}
    return_type = None
    try:
        sig = inspect.signature(func)
        for name, param in sig.parameters.items():
            if param.annotation != inspect.Parameter.empty:
                hints[name] = str(param.annotation)
            if sig.return_annotation != inspect.Signature.empty:
                return_type = str(sig.return_annotation)
    except (ValueError, TypeError):
        pass
    return hints, return_type


def reason_first(
    spec: str = "",
    preconditions: Optional[list[str]] = None,
    postconditions: Optional[list[str]] = None,
    test_inputs: Optional[list[dict[str, Any]]] = None,
    test_outputs: Optional[list[Any]] = None,
):
    """
    Decorator that marks a function as a reason-first program stub.

    Args:
        spec: Natural language specification of what the function should do
        preconditions: List of precondition descriptions
        postconditions: List of postcondition descriptions
        test_inputs: Optional list of test input dicts for validation
        test_outputs: Optional list of expected outputs

    Example:
        @reason_first(
            spec="Sort items by priority, breaking ties by recency",
            postconditions=["output is sorted", "all input items present in output"]
        )
        def process_queue(items: list[Item]) -> list[Item]:
            #> stable sort; O(n log n); must preserve Item identity
            ...
    """

    def decorator(func: Callable) -> Callable:
        source = inspect.getsource(func)
        source = textwrap.dedent(source)
        signature = str(inspect.signature(func))
        inline_specs = _extract_inline_specs(source)

        # Try to get module source for import/context extraction
        module_source = None
        try:
            module_file = inspect.getfile(func)
            if module_file:
                module_source = Path(module_file).read_text()
        except (TypeError, OSError):
            pass

        imports = _extract_imports(module_source) if module_source else []
        type_hints, return_type = _extract_type_hints(func)

        constraints = StubConstraints(
            imports=imports,
            type_hints=type_hints,
            return_type=return_type,
            inline_specs=inline_specs,
            decorator_spec=spec,
            preconditions=preconditions or [],
            postconditions=postconditions or [],
        )

        stub = Stub(
            name=func.__name__,
            source=source,
            signature=signature,
            constraints=constraints,
            module_source=module_source,
            test_inputs=test_inputs,
            test_outputs=test_outputs,
        )

        StubRegistry.register(stub)

        # Attach stub metadata to the function
        func._rfp_stub = stub
        return func

    return decorator