"""
ApeTool - LangChain BaseTool adapter for Ape tasks.

Wraps Ape tasks as LangChain tools, enabling deterministic execution
within LangChain agent workflows.
"""

from typing import Optional, Type, Any
from pathlib import Path

try:
    from langchain_core.tools import BaseTool
except ImportError:
    try:
        from langchain.tools import BaseTool
    except ImportError as e:
        raise ImportError(
            "ape-langchain requires langchain or langchain-core to be installed. "
            "Install it with: pip install langchain"
        ) from e

from pydantic import BaseModel, Field, ConfigDict

from ape_langchain.core import (
    compile_ape,
    validate_ape,
    execute_ape_function,
    ApeModule,
    ApeCompileError,
    ApeValidationError,
    ApeExecutionError,
)
from ape_langchain.task import ApeTask


class ApeTool(BaseTool):
    """
    A LangChain tool backed by an Ape module function.
    
    This adapter allows Ape tasks to be used seamlessly in LangChain
    agent workflows, with automatic validation and deterministic execution.
    
    Key features:
    - Compiles and validates Ape code at initialization
    - Validates inputs before execution
    - Provides clear error messages
    - Fully compatible with LangChain agents
    
    Attributes:
        name: Tool name (defaults to function name)
        description: Tool description for LLM
        ape_file: Path to the .ape source file
        function: Name of the function to execute
        
    Example:
        >>> tool = ApeTool(
        ...     ape_file="pricing.ape",
        ...     function="calculate_total_price",
        ...     description="Calculate price with tax"
        ... )
        >>> result = tool._run(base_price="99.99", tax_rate="0.21")
    """
    
    # LangChain BaseTool required fields
    name: str = ""
    description: str = ""
    
    # Ape-specific fields
    ape_file: str
    function: str
    
    # Internal state (not Pydantic fields)
    _module: Optional[ApeModule] = None
    _task: Optional[ApeTask] = None
    
    model_config = ConfigDict(arbitrary_types_allowed=True)
    
    def __init__(
        self,
        ape_file: str,
        function: str,
        description: Optional[str] = None,
        **kwargs: Any
    ):
        """
        Initialize ApeTool with an Ape source file and function name.
        
        Args:
            ape_file: Path to .ape source file
            function: Name of the function to wrap
            description: Optional description for LLM (auto-generated if None)
            **kwargs: Additional arguments for BaseTool
            
        Raises:
            ApeCompileError: If Ape compilation fails
            ApeValidationError: If validation fails
            KeyError: If function doesn't exist in module
        """
        # Compile and validate Ape module
        try:
            module = compile_ape(ape_file)
            validate_ape(module)
        except (ApeCompileError, ApeValidationError) as e:
            raise RuntimeError(f"Failed to initialize ApeTool: {e}") from e
        
        # Extract task metadata
        try:
            task = ApeTask.from_module(module, function)
        except KeyError as e:
            available = ", ".join(module.list_functions())
            raise KeyError(
                f"Function '{function}' not found in {ape_file}. "
                f"Available functions: {available}"
            ) from e
        
        # Set name and description
        tool_name = kwargs.pop("name", function)
        tool_description = description or task.description or f"Execute Ape task: {function}"
        
        # Initialize BaseTool
        super().__init__(
            name=tool_name,
            description=tool_description,
            ape_file=ape_file,
            function=function,
            **kwargs
        )
        
        # Store internal state
        self._module = module
        self._task = task
    
    def _run(
        self,
        **kwargs: Any
    ) -> str:
        """
        Execute the Ape function with validated inputs.
        
        This is called by LangChain when the agent invokes the tool.
        
        Args:
            **kwargs: Arguments to pass to the Ape function
            
        Returns:
            String representation of the result
            
        Raises:
            TypeError: If input validation fails
            ApeExecutionError: If execution fails
        """
        if self._task is None or self._module is None:
            raise RuntimeError("ApeTool not properly initialized")
        
        # Validate inputs
        try:
            self._task.validate_inputs(**kwargs)
        except TypeError as e:
            return f"Error: {e}"
        
        # Execute Ape function
        try:
            result = execute_ape_function(self._module, self.function, **kwargs)
            return str(result)
        except ApeExecutionError as e:
            return f"Execution error: {e}"
        except Exception as e:
            return f"Unexpected error: {e}"
    
    async def _arun(
        self,
        **kwargs: Any
    ) -> str:
        """
        Async execution - not yet implemented.
        
        Args:
            **kwargs: Arguments to pass to the Ape function
            
        Raises:
            NotImplementedError: Async execution not yet supported
        """
        raise NotImplementedError("ApeTool does not yet support async execution")
    
    def get_json_schema(self) -> dict:
        """
        Get JSON Schema representation of this tool.
        
        Useful for OpenAI function calling and similar APIs.
        
        Returns:
            JSON Schema dict
        """
        if self._task is None:
            raise RuntimeError("ApeTool not properly initialized")
        
        return self._task.to_json_schema()


__all__ = ["ApeTool"]

