# This script generated by AI and they'll be refactored later

import inspect
import typing
from typing import get_type_hints, Callable, Literal
import re
import json
import yaml

def python_type_to_json_type(py_type):
    """Convert Python type to JSON schema type string.
    
    Args:
        py_type: Python type annotation (e.g., int, str, list[str])
        
    Returns:
        str: JSON schema type string
        
    Example:
        >>> python_type_to_json_type(int)
        'integer'
        >>> python_type_to_json_type(list[str])
        'array'
    """
    origin = typing.get_origin(py_type) or py_type

    if origin in [int]:
        return "integer"
    elif origin in [float]:
        return "number"
    elif origin in [bool]:
        return "boolean"
    elif origin in [str]:
        return "string"
    elif origin in [list, tuple]:
        return "array"
    elif origin in [dict]:
        return "object"
    else:
        return "string"  # fallback

def convert_to_tool(func) -> dict:
    """Convert a Python function to a tool definition with JSON schema.
    
    Args:
        func: Python function to convert
        
    Returns:
        dict: Tool definition with name, description, and parameters schema
        
    Example:
        >>> def greet(name: str, age: int = 25) -> str:
        ...     '''Greet a person'''
        ...     return f"Hello {name}, age {age}"
        >>> tool = convert_to_tool(greet)
        >>> tool['name']
        'greet'
        >>> tool['parameters']['required']
        ['name']
    """
    sig = inspect.signature(func)
    type_hints = get_type_hints(func)

    parameters = {
        "type": "object",
        "properties": {},
        "required": []
    }

    for name, param in sig.parameters.items():
        param_type = type_hints.get(name, str)
        param_schema = {"type": python_type_to_json_type(param_type)}
        parameters["properties"][name] = param_schema
        if param.default is param.empty:
            parameters["required"].append(name)

    tool = {
        "name": func.__name__,
        "description": func.__doc__ or "",
        "parameters": parameters
    }

    return tool

def register_tools(tools: list):
    """Register multiple functions as tools with their definitions.
    
    Args:
        tools: List of Python functions to register
        
    Returns:
        dict: Dictionary mapping tool names to their definitions and functions
        
    Example:
        >>> def add(a: int, b: int) -> int:
        ...     return a + b
        >>> def multiply(x: float, y: float) -> float:
        ...     return x * y
        >>> registered = register_tools([add, multiply])
        >>> 'add' in registered
        True
        >>> registered['add']['tool'](2, 3)
        5
    """
    _tools = {}
    for tool in tools:
        tool_definition = convert_to_tool(tool)
        name = tool_definition.get("name", None)
        _tools[name] = {
            "definition": tool_definition,
            "tool": tool
        }
    return _tools

def list_tools(funcs: list[Callable]) -> str:
    """Generate a formatted list of tools for tool selection prompts.
    
    Extracts descriptions from function docstrings between <|start|> and <|end|> markers.
    
    Args:
        funcs: List of callable functions
        
    Returns:
        str: Formatted string listing tool names and descriptions
        
    Example:
        >>> def calc(x: int) -> int:
        ...     '''<|start|>Calculate something<|end|>'''
        ...     return x * 2
        >>> def search(query: str) -> str:
        ...     '''<|start|>Search for information<|end|>'''
        ...     return f"Results for {query}"
        >>> print(list_tools([calc, search]))
        \t- calc: Calculate something
        \t- search: Search for information
    """
    prompt = []
    for func in funcs:
        name = func.__name__
        definition = str(func.__doc__).strip()
        match = re.search(r'<\|start\|>(.*?)<\|end\|>', definition, re.DOTALL)
        text = match.group(1).strip() if match else ""
        text = re.sub(r'\s+', ' ', text).strip()
        prompt.append(f"\t- {name}: {text}")
    return "\n".join(prompt)

def generate_extract_parameters_prompt(func):
    """Generate a prompt template for extracting function parameters from user input.
    
    Creates a structured prompt that can be used to extract parameters from natural
    language input and format them as YAML for function calling.
    
    Args:
        func: Python function to generate extraction prompt for
        
    Returns:
        str: Formatted prompt template with parameter extraction instructions
        
    Example:
        >>> def weather(city: str, units: str = 'celsius') -> str:
        ...     '''<|start|>Get weather for a city<|end|>'''
        ...     return f"Weather in {city}"
        >>> prompt = generate_extract_parameters_prompt(weather)
        >>> 'city: <value>' in prompt
        True
        >>> 'units: <value>' in prompt
        True
    """
    sig = inspect.signature(func)
    doc = inspect.getdoc(func) or ""
    
    # Extract description between <|start|> and <|end|>
    desc_match = re.search(r'<\|start\|>(.*?)<\|end\|>', doc, re.DOTALL)
    description = desc_match.group(1).strip() if desc_match else ""
    description = re.sub(r'\s+', ' ', description).strip()
    
    # Build parameter info
    params = []
    for name, param in sig.parameters.items():
        param_type = param.annotation.__name__ if param.annotation != inspect.Parameter.empty else "any"
        params.append(f"- {name} ({param_type})")
    
    # Build YAML template with proper formatting for lists
    yaml_lines = []
    for name, param in sig.parameters.items():
        if param.annotation != inspect.Parameter.empty:
            origin = typing.get_origin(param.annotation) or param.annotation
            if origin in [list, tuple]:
                yaml_lines.append(f"{name}: [<value>]")
            else:
                yaml_lines.append(f"{name}: <value>")
        else:
            yaml_lines.append(f"{name}: <value>")
    
    yaml_template = "\n".join(yaml_lines)
    
    prompt = f"""Extract parameters for the {func.__name__} function from user input and return them in YAML format.

Function: {func.__name__}{sig}
Description: {description}

Parameters:
{chr(10).join(params)}

Return only the YAML format:
```yaml
{yaml_template}
```
User input: {{user_input}}"""
    return prompt

def get_yaml_function_definition(func):
    """Generate a prompt template for extracting function parameters from user input.
    
    Creates a structured prompt that can be used to extract parameters from natural
    language input and format them as YAML for function calling.
    
    Args:
        func: Python function to generate extraction prompt for
        
    Returns:
        str: Formatted prompt template with parameter extraction instructions
        
    Example:
        >>> def weather(city: str, units: str = 'celsius') -> str:
        ...     '''<|start|>Get weather for a city<|end|>'''
        ...     return f"Weather in {city}"
        >>> prompt = generate_extract_parameters_prompt(weather)
        >>> 'city: <value>' in prompt
        True
        >>> 'units: <value>' in prompt
        True
    """
    sig = inspect.signature(func)
    doc = inspect.getdoc(func) or ""
    
    # Extract description between <|start|> and <|end|>
    desc_match = re.search(r'<\|start\|>(.*?)<\|end\|>', doc, re.DOTALL)
    description = desc_match.group(1).strip() if desc_match else ""
    description = re.sub(r'\s+', ' ', description).strip()
    
    # Build parameter info
    params = []
    for name, param in sig.parameters.items():
        param_type = param.annotation.__name__ if param.annotation != inspect.Parameter.empty else "any"
        params.append(f"- {name} ({param_type})")
    
    # Build YAML template with proper formatting for lists
    yaml_lines = []
    for name, param in sig.parameters.items():
        if param.annotation != inspect.Parameter.empty:
            origin = typing.get_origin(param.annotation) or param.annotation
            if origin in [list, tuple]:
                yaml_lines.append(f"{name}: [<value>]")
            elif origin == str:
                yaml_lines.append(f'{name}: "<value>"')
            else:
                yaml_lines.append(f"{name}: <value>")
        else:
            yaml_lines.append(f"{name}: <value>")
    
    yaml_template = "\n".join(yaml_lines)
    
    prompt = f"""\
Function: {func.__name__}{sig}
Description: {description}

Parameters:
{chr(10).join(params)}

Return only the YAML format:
```yaml
{yaml_template}
```
""".strip()
    return prompt

def validate_parameters(params_dict, func):
    """Validate extracted parameters against function signature and types.
    
    Checks for missing required parameters and validates parameter types,
    with special handling for list types.
    
    Args:
        params_dict: Dictionary of parameter names to values
        func: Python function to validate against
        
    Returns:
        tuple: (is_valid: bool, error_message: str or None)
        
    Example:
        >>> def add(a: int, b: int) -> int:
        ...     return a + b
        >>> validate_parameters({'a': 1, 'b': 2}, add)
        (True, None)
        >>> validate_parameters({'a': 1}, add)
        (False, "Missing required parameters: b")
        >>> validate_parameters({'a': 'invalid', 'b': 2}, add)
        (True, None)  # Type checking is limited
    """
    sig = inspect.signature(func)
    
    # Check required parameters
    missing = []
    for name, param in sig.parameters.items():
        if param.default is param.empty and name not in params_dict:
            missing.append(name)
    
    if missing:
        return False, f"Missing required parameters: {', '.join(missing)}"
    
    # Check parameter types
    type_hints = get_type_hints(func)
    for name, value in params_dict.items():
        if name in type_hints:
            expected_type = type_hints[name]
            
            # Handle list[float] specifically
            if hasattr(expected_type, '__origin__') and expected_type.__origin__ is list:
                if not isinstance(value, list):
                    return False, f"Parameter '{name}' should be a list"
                
                # Check list elements
                if hasattr(expected_type, '__args__') and expected_type.__args__:
                    element_type = expected_type.__args__[0]
                    if element_type is float:
                        for i, item in enumerate(value):
                            if not isinstance(item, (int, float)):
                                return False, f"Parameter '{name}[{i}]' should be a number, got {type(item).__name__}"
    return True, None

def parse_codeblock_to_dict(output_str: str, codeblock: Literal['yaml', 'json'] = 'yaml'):
    """Parse YAML or JSON from markdown code blocks into Python dictionary.
    
    Extracts content from markdown code blocks and parses it as YAML or JSON.
    
    Args:
        output_str: String containing markdown with code blocks
        codeblock: Format to parse ('yaml' or 'json')
        
    Returns:
        dict: Parsed dictionary from the code block content
        
    Example:
        >>> text = '''Here's the data:
        ... ```yaml
        ... name: John
        ... age: 30
        ... ```'''
        >>> result = parse_codeblock_to_dict(text, 'yaml')
        >>> result['name']
        'John'
        >>> result['age']
        30
    """
    _output_str = output_str.split("```"+codeblock)
    if len(_output_str)>1:
        _output_str = _output_str[1]
    else:
        _output_str = _output_str[-1]
    _output_str = _output_str.split("```")[0].strip()
    if codeblock=='json':
        return json.loads(_output_str)
    return yaml.safe_load(_output_str)

def get_yaml_schema_definition(cls, as_array=False):
    schema = cls.model_json_schema()
    
    params = []
    yaml_lines = []
    
    for field_name, field_info in schema['properties'].items():
        field_type = field_info['type']
        # params.append(f"- {field_name} ({field_type})")
        description = field_info.get('description', '')
        params.append(f"- {field_name} ({field_type}) # {description}")
                
        
        if field_type == 'string':
            yaml_lines.append(f'{field_name}: "<value>"')
        else:
            yaml_lines.append(f"{field_name}: <value>")
    def lambda_indent(lines):
        _lines = []
        for enum, l in enumerate(lines):
            if enum == 0:
                _lines.append("- {l}".format(l=l))
            else:
                _lines.append("  {l}".format(l=l))
        return "\n".join(_lines)
    
    if as_array:
        yaml_template = "\n".join([lambda_indent(yaml_lines) for i in range(2)])
    else:
        yaml_template = "\n".join(yaml_lines)
    
    description = cls.__doc__ or f"Extract {cls.__name__} information"    
    prompt = f"""\
Model: {cls.__name__}
Description: {description}

Parameters:
{chr(10).join(params)}

Return only the YAML format:
```yaml
{yaml_template}
```
""".strip()
    return prompt