"""
Tool registry for managing tool schemas.
"""

from typing import Any, Dict, List, Optional, Tuple, Type, Union

from pydantic import BaseModel
from jsonschema import validators as jsonschema_validators
from jsonschema import ValidationError as JsonSchemaValidationError
import json
import sys

sys.tracebacklimit = (
    0  # Suppress traceback for validation errors to improve readability
)


class ToolRegistryError(Exception):
    """Custom exception for tool registry errors."""

    pass


class ToolRegistry:
    """
    Global registry for tool schemas.

    This stores tool definitions indexed by conversation/test context,
    allowing decorators to look up tool metadata during execution.
    """

    def __init__(self):
        self._tool_specs: Dict[str, List[Dict[str, Any]]] = {}

    def register_specs(
        self, tool_specs: List[Dict[str, Any]], context_key: str = "global"
    ) -> None:
        """Register tool specs for a context."""
        self._tool_specs[context_key] = tool_specs

    def get_tool_specs(self, context_key: str = "global") -> List[Dict[str, Any]]:
        """Get registered tool specs for a context."""
        return self._tool_specs.get(context_key, [])


# Global registry instance
_registry = ToolRegistry()


def _normalize_schema(
    schema: Optional[Union[Dict[str, Any], Type[BaseModel], BaseModel]],
) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
    if schema is None:
        return None, None
    if isinstance(schema, dict):
        return schema, "json_schema"
    if isinstance(schema, BaseModel):
        return schema.model_json_schema(), "pydantic"
    if isinstance(schema, type) and issubclass(schema, BaseModel):
        return schema.model_json_schema(), "pydantic"
    raise ValueError("Tool schema must be a JSON schema dict or a Pydantic BaseModel")


def _require_schema_description(
    schema: Optional[Dict[str, Any]], tool_name: str
) -> str:
    if not isinstance(schema, dict):
        raise ValueError(
            f"Tool '{tool_name}' parameters must include a JSON schema with a top-level description"
        )
    description = schema.get("description")
    if not isinstance(description, str) or not description.strip():
        raise ValueError(
            f"Tool '{tool_name}' parameters schema must include a non-empty top-level description"
        )
    return description.strip()


def _validate_tool_shape(tool: Dict[str, Any]) -> None:
    if tool.get("type") != "function":
        raise ValueError(
            f"Only 'function' type tools are supported, got {tool.get('type')}"
        )
    if "function" not in tool:
        raise ValueError("Tool must have a 'function' field")
    if "name" not in tool["function"]:
        raise ValueError("Tool function must have a 'name' field")


def _make_validator_for_schema(schema: Dict[str, Any]):
    validator_cls = jsonschema_validators.validator_for(schema)
    validator_cls.check_schema(schema)
    return validator_cls(
        schema,
        format_checker=validator_cls.FORMAT_CHECKER,
    )


def format_json_schema_validation_error(e: JsonSchemaValidationError) -> str:
    # fields avaiable are
    # ValidationError.message
    # ValidationError.validator
    # ValidationError.validator_value
    # ValidationError.schema
    # ValidationError.relative_schema_path
    # ValidationError.absolute_schema_path
    # ValidationError.schema_path
    # ValidationError.relative_path
    # ValidationError.absolute_path
    # ValidationError.json_path
    # ValidationError.path
    # ValidationError.instance
    # ValidationError.context
    # ValidationError.cause
    # ValidationError.parent
    # This function formats a JSON Schema validation error into succinct message that includes the path to the error and
    # the relevant part of the schema. This is useful for debugging when validating tool input/output against the defined schemas.
    schema_str = json.dumps(e.schema, indent=2)
    if len(schema_str) > 150:
        schema_str = schema_str[:150] + "\n... (truncated)"
    message = f"""JSON Schema validation error at path {'/'.join(str(p) for p in e.absolute_path)}: 
{e.message}

Schema path: {'/'.join(str(p) for p in e.absolute_schema_path)}

Failed schema: {schema_str}

Value that failed validation: {json.dumps(e.instance, indent=2)}

Please ensure that the tool input/output matches the defined schema for the tool and that the schema is valid.
"""
    return message


def _validate_examples(
    examples: Optional[List[Dict[str, Any]]],
    tool_name: str,
    input_schema: Optional[Dict[str, Any]],
    output_schema: Optional[Dict[str, Any]],
) -> None:
    if examples is None:
        raise ValueError(f"Tool '{tool_name}' must include an 'examples' field")
    if not isinstance(examples, list):
        raise ValueError(f"Tool '{tool_name}' examples must be a list of objects")
    output_validator = _make_validator_for_schema(output_schema)
    input_validator = _make_validator_for_schema(input_schema)
    for index, example in enumerate(examples):
        if not isinstance(example, dict):
            raise ValueError(f"Tool '{tool_name}' examples[{index}] must be a dict")
        if "input" not in example or "output" not in example:
            raise ValueError(
                f"Tool '{tool_name}' examples[{index}] must include input/output"
            )
        try:
            input_validator.validate(example["input"])
        except JsonSchemaValidationError as e:
            formatted_message = format_json_schema_validation_error(e)
            raise ToolRegistryError(
                f"Tool '{tool_name}' example {index} input validation error: {formatted_message}"
            ) from None
        try:
            output_validator.validate(example["output"])
        except JsonSchemaValidationError as e:
            formatted_message = format_json_schema_validation_error(e)
            raise ToolRegistryError(
                f"Tool '{tool_name}' example {index} output validation error: {formatted_message}"
            ) from None


def _normalize_tool_definition(
    tool: Dict[str, Any],
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    _validate_tool_shape(tool)
    function_def = tool["function"]
    tool_name = function_def["name"]

    input_schema, input_source = _normalize_schema(function_def.get("parameters"))
    if input_schema is None:
        input_schema = {"type": "object", "properties": {}}
        input_source = input_source or "json_schema"

    output_schema, output_source = _normalize_schema(function_def.get("returns"))
    examples = function_def.get("examples")

    _validate_examples(examples, tool_name, input_schema, output_schema)

    schema_source = (
        "pydantic"
        if input_source == "pydantic" or output_source == "pydantic"
        else "json_schema"
    )

    tool_description = _require_schema_description(input_schema, tool_name)
    tool_spec = {
        "tool_name": tool_name,
        "description": tool_description,
        "input_schema": input_schema,
        "output_schema": output_schema,
        "examples": examples,
        "schema_source": schema_source,
    }

    normalized_function = dict(function_def)
    normalized_function["description"] = tool_description
    normalized_function["parameters"] = input_schema
    normalized_function.pop("returns", None)
    normalized_function.pop("examples", None)

    normalized_tool = dict(tool)
    normalized_tool["function"] = normalized_function
    return normalized_tool, tool_spec


def register_tools(
    tools: List[Dict[str, Any]], application_id: str
) -> List[Dict[str, Any]]:
    """
    Register tools with Snowglobe.

    This function normalizes your tool definitions (including Pydantic -> JSON
    Schema conversion) and registers them internally while returning the
    normalized array for use with LLM APIs (e.g., OpenAI).

    Args:
        tools: List of tool definitions in OpenAI format

    Returns:
        The normalized tools list (ready for LLM API)

    Example:
        TOOLS = [
            {
                "type": "function",
                "function": {
                    "name": "get_weather",
                    "parameters": {
                        "type": "object",
                        "description": "Get current weather",
                        "properties": {...},
                    }
                }
            }
        ]

        def tools_defs():
            return TOOLS

        # Use TOOLS directly with OpenAI API
        TOOLS_MAP = {
            "get_weather": get_weather,
            ...}

        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=messages,
            tools=TOOLS
        )

        while True:
            response = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=messages,
                tools=TOOLS
            )
            message = response.choices[0].message

            if not message.tool_calls:
                return CompletionFunctionOutputs(response=message.content)

            messages.append(message)
            for tool_call in message.tool_calls:
                if tool_call.function.name not in TOOLS_MAP:
                    result = f"Error: Tool {tool_call.function.name} not found"
                else:
                    result = TOOLS_MAP[tool_call.function.name](**json.loads(tool_call.function.arguments))

                messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": str(result)
                })
    """
    # Validate tool schema (basic check)
    normalized_tools: List[Dict[str, Any]] = []
    tool_specs: List[Dict[str, Any]] = []
    if not application_id or not isinstance(application_id, str):
        # raise error and give instructions on how to update agent.py with def tools_defs
        raise ValueError(
            """
            Application ID is required to register tools. Please register tools with a def tools_defs function in your agent.py file that returns the list of tools, and ensure that the application_id parameter is a non-empty string.
            Example:
        def tools_defs():
            return [
                {
                    "type": "function",
                    "function": {
                        "name": "get_weather",
                        "parameters": {
                            "type": "object",
                            "description": "Get current weather",
                            "properties": {...},
                        },
                        "returns": {
                            "type": "object",
                            "description": "Current weather information",
                            "properties": {...},                        
                        }
                    }
                }
            ]
            """
        )
    for tool in tools:
        if not isinstance(tool, dict):
            raise ValueError(f"Tool must be a dictionary, got {type(tool)}")
        try:
            normalized_tool, tool_spec = _normalize_tool_definition(tool)
            normalized_tools.append(normalized_tool)
            tool_specs.append(tool_spec)
        except Exception:
            print(
                f"ERROR: Registering tool: {tool.get('function', {}).get('name', 'unknown')}"
            )
            raise

    _registry.register_specs(tool_specs, context_key=application_id)

    # Return normalized tools list for LLM APIs
    return normalized_tools


def get_registered_tool_specs(context_key: str = "global") -> List[Dict[str, Any]]:
    """Get all registered tool specs for a context."""
    return _registry.get_tool_specs(context_key)
