"""
OpenAI function call execution via Ape runtime.

Handles parsing OpenAI function call arguments and executing them
through the Ape runtime with full validation.
"""

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

try:
    from ape import compile as ape_compile
    from ape import validate as ape_validate
    from ape import ApeModule, ApeCompileError, ApeValidationError, ApeExecutionError
    APE_AVAILABLE = True
except ImportError:
    APE_AVAILABLE = False
    ApeModule = None  # type: ignore
    ApeCompileError = Exception  # type: ignore
    ApeValidationError = Exception  # type: ignore
    ApeExecutionError = Exception  # type: ignore


def execute_openai_call(
    module: Any,
    function_name: str,
    arguments_json: str
) -> Any:
    """
    Execute an OpenAI function call with Ape validation.
    
    Takes the raw JSON string from OpenAI's function call response,
    validates it against the Ape task signature, and executes it
    deterministically.
    
    Args:
        module: Compiled ApeModule
        function_name: Name of the function to call
        arguments_json: JSON string of arguments from OpenAI
        
    Returns:
        Function execution result (JSON-serializable)
        
    Raises:
        json.JSONDecodeError: If arguments_json is invalid JSON
        TypeError: If arguments don't match task signature
        ApeExecutionError: If execution fails
        
    Example:
        >>> from ape import compile
        >>> module = compile("calculator.ape")
        >>> result = execute_openai_call(
        ...     module,
        ...     "add",
        ...     '{"a": "5", "b": "3"}'
        ... )
        >>> print(result)
        8
    """
    if not APE_AVAILABLE:
        raise ImportError("ape-openai requires ape-lang to be installed")
    
    # Parse JSON arguments
    try:
        arguments = json.loads(arguments_json)
    except json.JSONDecodeError as e:
        raise json.JSONDecodeError(
            f"Invalid JSON in OpenAI function arguments: {e}",
            e.doc,
            e.pos
        ) from e
    
    # Validate arguments match task signature
    try:
        signature = module.get_function_signature(function_name)
    except KeyError as e:
        available = ", ".join(module.list_functions())
        raise KeyError(
            f"Function '{function_name}' not found in module. "
            f"Available: {available}"
        ) from e
    
    # Check required parameters
    provided = set(arguments.keys())
    required = set(signature.inputs.keys())
    
    missing = required - provided
    if missing:
        raise TypeError(
            f"Missing required parameters for '{function_name}': {', '.join(sorted(missing))}"
        )
    
    unexpected = provided - required
    if unexpected:
        raise TypeError(
            f"Unexpected parameters for '{function_name}': {', '.join(sorted(unexpected))}"
        )
    
    # Execute via Ape runtime
    try:
        result = module.call(function_name, **arguments)
        return result
    except Exception as e:
        raise ApeExecutionError(
            f"Execution of '{function_name}' failed: {e}"
        ) from e


class ApeOpenAIFunction:
    """
    High-level wrapper for Ape tasks as OpenAI functions.
    
    Combines schema generation and execution in a convenient interface.
    
    Attributes:
        module: Compiled ApeModule
        function_name: Name of the function
        description: Custom description for OpenAI
        
    Example:
        >>> func = ApeOpenAIFunction.from_ape_file("calculator.ape", "add")
        >>> 
        >>> # Get OpenAI tool definition
        >>> tool = func.to_openai_tool()
        >>> 
        >>> # Execute function call
        >>> result = func.execute('{"a": "5", "b": "3"}')
    """
    
    def __init__(
        self,
        module: Any,
        function_name: str,
        description: Optional[str] = None
    ):
        """
        Initialize ApeOpenAIFunction.
        
        Args:
            module: Compiled ApeModule
            function_name: Name of the function to wrap
            description: Optional custom description
        """
        if not APE_AVAILABLE:
            raise ImportError("ape-openai requires ape-lang to be installed")
        
        self.module = module
        self.function_name = function_name
        self.description = description
        
        # Validate function exists
        try:
            self.signature = module.get_function_signature(function_name)
        except KeyError as e:
            available = ", ".join(module.list_functions())
            raise KeyError(
                f"Function '{function_name}' not found. Available: {available}"
            ) from e
    
    @classmethod
    def from_ape_file(
        cls,
        ape_file: str | Path,
        function_name: str,
        description: Optional[str] = None
    ) -> "ApeOpenAIFunction":
        """
        Create ApeOpenAIFunction from an Ape source file.
        
        Args:
            ape_file: Path to .ape source file
            function_name: Name of the function to wrap
            description: Optional custom description
            
        Returns:
            ApeOpenAIFunction instance
            
        Raises:
            ApeCompileError: If compilation fails
            KeyError: If function doesn't exist
        """
        module = ape_compile(ape_file)
        ape_validate(module)
        return cls(module, function_name, description)
    
    def to_openai_tool(self) -> dict:
        """
        Generate OpenAI tool definition.
        
        Returns:
            OpenAI tool schema dictionary
        """
        from ape_openai.schema import ape_task_to_openai_schema
        from ape_openai import ApeTask
        
        # Create ApeTask from signature
        task = ApeTask(
            name=self.signature.name,
            inputs=self.signature.inputs,
            output=self.signature.output,
            description=self.description or self.signature.description
        )
        
        return ape_task_to_openai_schema(task)
    
    def execute(self, arguments_json: str) -> Any:
        """
        Execute OpenAI function call.
        
        Args:
            arguments_json: JSON string of arguments from OpenAI
            
        Returns:
            Function execution result
            
        Raises:
            json.JSONDecodeError: If JSON is invalid
            TypeError: If arguments don't match signature
            ApeExecutionError: If execution fails
        """
        return execute_openai_call(
            self.module,
            self.function_name,
            arguments_json
        )
    
    def __repr__(self) -> str:
        return f"ApeOpenAIFunction(function='{self.function_name}')"


__all__ = [
    "execute_openai_call",
    "ApeOpenAIFunction",
]
