"""Function calling utilities for structured tool use.

Enables LLMs to output structured function calls and handle results.
"""

import json
import re
from typing import Any, Callable, Dict, List, Optional, Union
from dataclasses import dataclass, field


# Function call markers
FUNCTION_START = "<function>"
FUNCTION_END = "</function>"
RESULT_START = "<function_result>"
RESULT_END = "</function_result>"


@dataclass
class FunctionCall:
    """Represents a parsed function call from model output."""
    name: str
    arguments: Dict[str, Any]
    raw: str = ""


@dataclass 
class FunctionDefinition:
    """Definition of a callable function for the model."""
    name: str
    description: str
    parameters: Dict[str, Any]
    handler: Optional[Callable] = None
    
    def to_prompt_format(self) -> str:
        """Format for including in system prompt."""
        params_str = json.dumps(self.parameters, indent=2)
        return f"""Function: {self.name}
Description: {self.description}
Parameters: {params_str}"""


class FunctionRegistry:
    """Registry of available functions for the model to call.
    
    Example:
        >>> registry = FunctionRegistry()
        >>> 
        >>> @registry.register("get_weather", "Get current weather", {"city": "string"})
        >>> def get_weather(city: str) -> dict:
        ...     return {"temp": 25, "condition": "sunny"}
        >>> 
        >>> # Use in prompt
        >>> system_prompt = registry.get_system_prompt()
    """
    
    def __init__(self):
        self.functions: Dict[str, FunctionDefinition] = {}
    
    def register(
        self,
        name: str,
        description: str,
        parameters: Dict[str, Any],
    ) -> Callable:
        """Decorator to register a function.
        
        Args:
            name: Function name for the model.
            description: What the function does.
            parameters: Parameter schema.
            
        Returns:
            Decorated function.
        """
        def decorator(func: Callable) -> Callable:
            self.functions[name] = FunctionDefinition(
                name=name,
                description=description,
                parameters=parameters,
                handler=func,
            )
            return func
        return decorator
    
    def add(
        self,
        name: str,
        handler: Callable,
        description: str,
        parameters: Dict[str, Any],
    ):
        """Add a function to the registry.
        
        Args:
            name: Function name.
            handler: The callable to execute.
            description: What the function does.
            parameters: Parameter schema.
        """
        self.functions[name] = FunctionDefinition(
            name=name,
            description=description,
            parameters=parameters,
            handler=handler,
        )
    
    def get_system_prompt(self) -> str:
        """Generate system prompt section describing available functions."""
        if not self.functions:
            return ""
        
        lines = ["You have access to the following functions:", ""]
        for func in self.functions.values():
            lines.append(func.to_prompt_format())
            lines.append("")
        
        lines.extend([
            "To call a function, output:",
            f'{FUNCTION_START}{{"name": "function_name", "arguments": {{}}}}{FUNCTION_END}',
            "",
            "Wait for the result before continuing.",
        ])
        
        return "\n".join(lines)
    
    def execute(self, call: FunctionCall) -> Any:
        """Execute a function call.
        
        Args:
            call: The parsed function call.
            
        Returns:
            Function result.
            
        Raises:
            KeyError: If function not found.
            Exception: If function execution fails.
        """
        if call.name not in self.functions:
            raise KeyError(f"Unknown function: {call.name}")
        
        func_def = self.functions[call.name]
        if func_def.handler is None:
            raise ValueError(f"No handler for function: {call.name}")
        
        return func_def.handler(**call.arguments)


def parse_function_call(output: str) -> Optional[FunctionCall]:
    """Extract function call from model output.
    
    Args:
        output: Model's generated text.
        
    Returns:
        Parsed FunctionCall or None if not found.
        
    Example:
        >>> output = '<function>{"name": "get_weather", "arguments": {"city": "Hanoi"}}</function>'
        >>> call = parse_function_call(output)
        >>> print(call.name)  # "get_weather"
        >>> print(call.arguments)  # {"city": "Hanoi"}
    """
    pattern = rf'{re.escape(FUNCTION_START)}(.+?){re.escape(FUNCTION_END)}'
    match = re.search(pattern, output, re.DOTALL)
    
    if not match:
        return None
    
    try:
        raw = match.group(1).strip()
        data = json.loads(raw)
        return FunctionCall(
            name=data.get("name", ""),
            arguments=data.get("arguments", {}),
            raw=raw,
        )
    except json.JSONDecodeError:
        return None


def format_function_result(result: Any, name: str = "") -> str:
    """Format function result for model input.
    
    Args:
        result: The function result.
        name: Optional function name for context.
        
    Returns:
        Formatted result string.
    """
    if isinstance(result, (dict, list)):
        result_str = json.dumps(result, ensure_ascii=False)
    else:
        result_str = str(result)
    
    return f"{RESULT_START}{result_str}{RESULT_END}"


def create_function_call_training_example(
    instruction: str,
    function_name: str,
    arguments: Dict[str, Any],
    result: Any,
    final_response: str,
) -> Dict[str, str]:
    """Create a training example for function calling.
    
    Args:
        instruction: User's instruction.
        function_name: Name of function to call.
        arguments: Function arguments.
        result: Function result.
        final_response: Final response after function call.
        
    Returns:
        Dict with "instruction" and "output" keys for training.
        
    Example:
        >>> example = create_function_call_training_example(
        ...     instruction="What's the weather in Hanoi?",
        ...     function_name="get_weather",
        ...     arguments={"city": "Hanoi"},
        ...     result={"temp": 28, "condition": "sunny"},
        ...     final_response="The weather in Hanoi is 28°C and sunny."
        ... )
    """
    function_call = f'{FUNCTION_START}{{"name": "{function_name}", "arguments": {json.dumps(arguments)}}}{FUNCTION_END}'
    result_str = format_function_result(result)
    
    output = f"{function_call}\n{result_str}\n{final_response}"
    
    return {
        "instruction": instruction,
        "output": output,
    }


class FunctionCallingAgent:
    """Agent that handles function calling loop with a model.
    
    Example:
        >>> agent = FunctionCallingAgent(model, tokenizer, registry)
        >>> response = agent.chat("What's the weather in Hanoi?")
    """
    
    def __init__(
        self,
        model,
        tokenizer,
        registry: FunctionRegistry,
        max_function_calls: int = 5,
        system_prompt: str = "",
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.registry = registry
        self.max_function_calls = max_function_calls
        self.base_system_prompt = system_prompt
    
    def get_full_system_prompt(self) -> str:
        """Get complete system prompt with function definitions."""
        parts = []
        if self.base_system_prompt:
            parts.append(self.base_system_prompt)
        parts.append(self.registry.get_system_prompt())
        return "\n\n".join(parts)
    
    def chat(
        self,
        user_input: str,
        max_length: int = 512,
        **generate_kwargs,
    ) -> str:
        """Process user input with function calling support.
        
        Args:
            user_input: User's message.
            max_length: Max generation length.
            **generate_kwargs: Additional args for model.generate().
            
        Returns:
            Final response after any function calls.
        """
        system = self.get_full_system_prompt()
        
        # Build conversation
        messages = []
        if system:
            messages.append(f"<|im_start|>system\n{system}<|im_end|>")
        messages.append(f"<|im_start|>user\n{user_input}<|im_end|>")
        messages.append("<|im_start|>assistant\n")
        
        prompt = "\n".join(messages)
        
        for _ in range(self.max_function_calls + 1):
            # Generate response
            output = self.model.generate(
                prompt,
                self.tokenizer,
                max_length=max_length,
                **generate_kwargs,
            )
            
            # Extract assistant's response
            if "<|im_start|>assistant\n" in output:
                response = output.split("<|im_start|>assistant\n")[-1]
                if "<|im_end|>" in response:
                    response = response.split("<|im_end|}")[0]
            else:
                response = output
            
            # Check for function call
            call = parse_function_call(response)
            if call is None:
                # No function call, return final response
                return response.strip()
            
            # Execute function
            try:
                result = self.registry.execute(call)
                result_str = format_function_result(result, call.name)
            except Exception as e:
                result_str = format_function_result({"error": str(e)})
            
            # Add to conversation and continue
            prompt = output + f"\n{result_str}\n"
        
        return response.strip()
