"""
ApeChain - High-level wrapper for direct Ape task execution.

Provides a simple interface for executing Ape tasks without
full LangChain agent orchestration.
"""

from typing import Any, Optional
from pathlib import Path

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


class ApeChain:
    """
    High-level wrapper for executing a single Ape function as a 'chain'-like object.
    
    This provides a simpler alternative to full LangChain agents when you just
    want to execute Ape tasks directly with deterministic validation.
    
    Attributes:
        ape_file: Path to the .ape source file
        function: Name of the function to execute
        
    Example:
        >>> chain = ApeChain.from_ape_file("pricing.ape", "calculate_total_price")
        >>> result = chain.run(base_price="99.99", tax_rate="0.21")
        >>> print(result)
        121.00
    """
    
    def __init__(self, ape_file: str, function: str):
        """
        Initialize ApeChain with an Ape source file and function name.
        
        Args:
            ape_file: Path to .ape source file
            function: Name of the function to execute
            
        Raises:
            ApeCompileError: If compilation fails
            ApeValidationError: If validation fails
            KeyError: If function doesn't exist
        """
        self.ape_file = ape_file
        self.function = function
        
        # Compile and validate module
        self._module = compile_ape(ape_file)
        validate_ape(self._module)
        
        # Extract task metadata
        try:
            self._task = ApeTask.from_module(self._module, function)
        except KeyError as e:
            available = ", ".join(self._module.list_functions())
            raise KeyError(
                f"Function '{function}' not found in {ape_file}. "
                f"Available functions: {available}"
            ) from e
    
    @classmethod
    def from_ape_file(cls, ape_file: str, function: str) -> "ApeChain":
        """
        Create an ApeChain from an Ape source file.
        
        This is the recommended way to create ApeChain instances.
        
        Args:
            ape_file: Path to .ape source file
            function: Name of the function to execute
            
        Returns:
            Configured ApeChain instance
            
        Example:
            >>> chain = ApeChain.from_ape_file("calculator.ape", "add")
            >>> result = chain.run(a="5", b="3")
        """
        return cls(ape_file, function)
    
    def run(self, **kwargs: Any) -> Any:
        """
        Execute the Ape function with the given keyword arguments.
        
        This validates inputs and executes the function deterministically.
        
        Args:
            **kwargs: Arguments to pass to the Ape function
            
        Returns:
            The function's return value
            
        Raises:
            TypeError: If input validation fails
            ApeExecutionError: If execution fails
            
        Example:
            >>> chain = ApeChain.from_ape_file("pricing.ape", "calculate_total_price")
            >>> result = chain.run(base_price="99.99", tax_rate="0.21")
        """
        # Validate inputs
        self._task.validate_inputs(**kwargs)
        
        # Execute
        try:
            return execute_ape_function(self._module, self.function, **kwargs)
        except ApeExecutionError as e:
            raise ApeExecutionError(
                f"Execution of '{self.function}' failed: {e}"
            ) from e
    
    def get_task(self) -> ApeTask:
        """
        Get the underlying ApeTask metadata.
        
        Useful for introspection and debugging.
        
        Returns:
            ApeTask instance with function metadata
        """
        return self._task
    
    def get_json_schema(self) -> dict:
        """
        Get JSON Schema representation of the wrapped function.
        
        Useful for generating documentation or UI forms.
        
        Returns:
            JSON Schema dict
        """
        return self._task.to_json_schema()
    
    def __repr__(self) -> str:
        return f"ApeChain(file='{self.ape_file}', function='{self.function}')"
    
    def __call__(self, **kwargs: Any) -> Any:
        """
        Make ApeChain callable directly.
        
        This allows using ApeChain instances as functions:
        >>> chain = ApeChain.from_ape_file("calculator.ape", "add")
        >>> result = chain(a="5", b="3")
        
        Args:
            **kwargs: Arguments to pass to the Ape function
            
        Returns:
            The function's return value
        """
        return self.run(**kwargs)


__all__ = ["ApeChain"]
