"""
ApeTask - Representation of an Ape task with JSON Schema conversion.

This module provides introspection and metadata for Ape tasks,
enabling integration with LLM function calling APIs (OpenAI, Anthropic, etc.)
"""

from dataclasses import dataclass
from typing import Dict, Optional, Any

from ape_langchain.core import ApeModule


@dataclass
class ApeTask:
    """
    Representation of a single callable task in an Ape module.
    
    Encodes:
    - Task name
    - Input parameters (names and types)
    - Output type
    - Optional description
    
    Attributes:
        name: Function/task name
        inputs: Mapping of parameter name to type string
        output: Output type string (optional)
        description: Human-readable description (optional)
    """
    
    name: str
    inputs: Dict[str, str]
    output: Optional[str] = None
    description: Optional[str] = None
    
    @classmethod
    def from_module(cls, module: ApeModule, function_name: str) -> "ApeTask":
        """
        Extract task metadata from a compiled Ape module.
        
        Uses the ApeModule.get_function_signature() API to introspect
        the function and extract its parameters and return type.
        
        Args:
            module: Compiled ApeModule
            function_name: Name of the function to extract
            
        Returns:
            ApeTask with metadata about the function
            
        Raises:
            KeyError: If function doesn't exist in module
            
        Example:
            >>> from ape import compile
            >>> module = compile("calculator.ape")
            >>> task = ApeTask.from_module(module, "add")
            >>> print(task.inputs)
            {'a': 'str', 'b': 'str'}
        """
        signature = module.get_function_signature(function_name)
        
        return cls(
            name=signature.name,
            inputs=signature.inputs,
            output=signature.output,
            description=signature.description
        )
    
    def to_json_schema(self) -> Dict[str, Any]:
        """
        Convert task metadata to JSON Schema format.
        
        This enables integration with:
        - OpenAI function calling
        - Anthropic tool use
        - LangChain tool schemas
        - Any system that uses JSON Schema
        
        Returns:
            JSON Schema dict describing inputs and outputs
            
        Note:
            In v0.1.0, all types are mapped to 'string' by default.
            Future versions will support rich type mapping:
            - Ape Integer -> JSON Schema integer
            - Ape List[String] -> JSON Schema array of strings
            - Ape custom entities -> JSON Schema objects
            
        Example:
            >>> task = ApeTask(name="add", inputs={"a": "str", "b": "str"})
            >>> schema = task.to_json_schema()
            >>> print(schema)
            {
                'type': 'function',
                'function': {
                    'name': 'add',
                    'parameters': {
                        'type': 'object',
                        'properties': {
                            'a': {'type': 'string'},
                            'b': {'type': 'string'}
                        },
                        'required': ['a', 'b']
                    }
                }
            }
        """
        # Map Ape types to JSON Schema types
        # v0.1.0: Simple mapping (all strings)
        # v0.2.0: Will add int, float, bool, array, object
        type_mapping = {
            'str': 'string',
            'int': 'integer',
            'float': 'number',
            'bool': 'boolean',
            'Any': 'string',  # Default fallback
        }
        
        properties = {}
        required = []
        
        for param_name, param_type in self.inputs.items():
            # Map Ape type to JSON Schema type
            json_type = type_mapping.get(param_type, 'string')
            properties[param_name] = {'type': json_type}
            required.append(param_name)
        
        schema = {
            'type': 'function',
            'function': {
                'name': self.name,
                'description': self.description or f"Execute Ape task: {self.name}",
                'parameters': {
                    'type': 'object',
                    'properties': properties,
                    'required': required
                }
            }
        }
        
        return schema
    
    def validate_inputs(self, **kwargs) -> None:
        """
        Validate that provided arguments match expected inputs.
        
        Checks:
        1. All required parameters are present
        2. No unexpected parameters are provided
        
        Args:
            **kwargs: Arguments to validate
            
        Raises:
            TypeError: If validation fails
            
        Note:
            Type checking is currently basic. Future versions will
            perform deep type validation using Ape's type system.
        """
        provided = set(kwargs.keys())
        required = set(self.inputs.keys())
        
        # Check for missing parameters
        missing = required - provided
        if missing:
            raise TypeError(
                f"Missing required parameters for '{self.name}': {', '.join(sorted(missing))}"
            )
        
        # Check for unexpected parameters
        unexpected = provided - required
        if unexpected:
            raise TypeError(
                f"Unexpected parameters for '{self.name}': {', '.join(sorted(unexpected))}"
            )
    
    def __repr__(self) -> str:
        params = ', '.join(f"{k}: {v}" for k, v in self.inputs.items())
        output_str = f" -> {self.output}" if self.output else ""
        return f"ApeTask({self.name}({params}){output_str})"


__all__ = ["ApeTask"]
