Berry MCP Server

"""
Tool output schemas and capability metadata for enhanced MCP integration
"""

from dataclasses import dataclass, field
from typing import Any, Optional, Union


@dataclass
class ToolOutputSchema:
    """Schema definition for tool output"""

    type: str = "object"  # JSON Schema type
    properties: dict[str, Any] = field(default_factory=dict)
    required: list[str] = field(default_factory=list)
    description: str = ""
    examples: list[Any] = field(default_factory=list)

    def add_property(
        self,
        name: str,
        prop_type: str,
        description: str = "",
        required: bool = False,
        enum: list[Any] | None = None,
        format: str | None = None,
    ) -> None:
        """Add a property to the schema"""
        prop_def: dict[str, Any] = {"type": prop_type, "description": description}

        if enum:
            prop_def["enum"] = enum
        if format:
            prop_def["format"] = format

        self.properties[name] = prop_def

        if required:
            self.required.append(name)

    def to_json_schema(self) -> dict[str, Any]:
        """Convert to JSON Schema format"""
        schema = {
            "type": self.type,
            "properties": self.properties,
            "description": self.description,
        }

        if self.required:
            schema["required"] = self.required

        if self.examples:
            schema["examples"] = self.examples

        return schema


@dataclass
class CapabilityMetadata:
    """Metadata for enhanced capability discovery"""

    name: str
    description: str
    version: str = "1.0.0"
    category: str = "general"
    tags: list[str] = field(default_factory=list)
    dependencies: list[str] = field(default_factory=list)
    output_schema: ToolOutputSchema | None = None
    supports_streaming: bool = False
    supports_cancellation: bool = False
    requires_authentication: bool = False
    estimated_duration: str | None = None  # "fast", "medium", "slow"
    resource_requirements: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary representation"""
        result = {
            "name": self.name,
            "description": self.description,
            "version": self.version,
            "category": self.category,
            "tags": self.tags,
            "dependencies": self.dependencies,
            "supports_streaming": self.supports_streaming,
            "supports_cancellation": self.supports_cancellation,
            "requires_authentication": self.requires_authentication,
        }

        if self.output_schema:
            result["output_schema"] = self.output_schema.to_json_schema()

        if self.estimated_duration:
            result["estimated_duration"] = self.estimated_duration

        if self.resource_requirements:
            result["resource_requirements"] = self.resource_requirements

        return result


class SchemaBuilder:
    """Builder class for creating tool output schemas"""

    @staticmethod
    def create_object_schema(description: str = "") -> ToolOutputSchema:
        """Create an object schema"""
        return ToolOutputSchema(type="object", description=description)

    @staticmethod
    def create_array_schema(item_type: str, description: str = "") -> ToolOutputSchema:
        """Create an array schema"""
        return ToolOutputSchema(
            type="array",
            description=description,
            properties={"items": {"type": item_type}},
        )

    @staticmethod
    def create_string_schema(
        description: str = "", enum: list[str] | None = None
    ) -> ToolOutputSchema:
        """Create a string schema"""
        schema = ToolOutputSchema(type="string", description=description)
        if enum:
            schema.properties = {"enum": enum}
        return schema

    @staticmethod
    def create_file_result_schema() -> ToolOutputSchema:
        """Create schema for file operation results"""
        schema = SchemaBuilder.create_object_schema("File operation result")
        schema.add_property(
            "success", "boolean", "Whether the operation succeeded", required=True
        )
        schema.add_property("path", "string", "File path")
        schema.add_property("size", "integer", "File size in bytes")
        schema.add_property(
            "modified", "string", "Last modified timestamp", format="date-time"
        )
        schema.add_property("error", "string", "Error message if operation failed")
        return schema

    @staticmethod
    def create_search_result_schema() -> ToolOutputSchema:
        """Create schema for search results"""
        schema = SchemaBuilder.create_object_schema("Search results")
        schema.add_property("query", "string", "The search query", required=True)
        schema.add_property(
            "total_results", "integer", "Total number of results", required=True
        )
        schema.add_property(
            "results", "array", "Array of search results", required=True
        )
        schema.add_property("page", "integer", "Current page number")
        schema.add_property("per_page", "integer", "Results per page")
        return schema

    @staticmethod
    def create_api_response_schema() -> ToolOutputSchema:
        """Create schema for API responses"""
        schema = SchemaBuilder.create_object_schema("API response")
        schema.add_property("status_code", "integer", "HTTP status code", required=True)
        schema.add_property("data", "object", "Response data")
        schema.add_property("headers", "object", "Response headers")
        schema.add_property("error", "string", "Error message if request failed")
        return schema


class CapabilityBuilder:
    """Builder class for creating capability metadata"""

    @staticmethod
    def create_file_tool_capability(
        name: str, description: str, supports_streaming: bool = False
    ) -> CapabilityMetadata:
        """Create capability metadata for file tools"""
        return CapabilityMetadata(
            name=name,
            description=description,
            category="file_operations",
            tags=["files", "io"],
            output_schema=SchemaBuilder.create_file_result_schema(),
            supports_streaming=supports_streaming,
            supports_cancellation=True,
            estimated_duration="fast",
        )

    @staticmethod
    def create_search_tool_capability(
        name: str, description: str, requires_auth: bool = False
    ) -> CapabilityMetadata:
        """Create capability metadata for search tools"""
        return CapabilityMetadata(
            name=name,
            description=description,
            category="search",
            tags=["search", "query"],
            output_schema=SchemaBuilder.create_search_result_schema(),
            supports_streaming=True,
            supports_cancellation=True,
            requires_authentication=requires_auth,
            estimated_duration="medium",
        )

    @staticmethod
    def create_api_tool_capability(
        name: str, description: str, dependencies: list[str] | None = None
    ) -> CapabilityMetadata:
        """Create capability metadata for API tools"""
        return CapabilityMetadata(
            name=name,
            description=description,
            category="api",
            tags=["api", "http"],
            dependencies=dependencies or [],
            output_schema=SchemaBuilder.create_api_response_schema(),
            supports_streaming=False,
            supports_cancellation=True,
            requires_authentication=True,
            estimated_duration="medium",
            resource_requirements={"network": True, "rate_limit": "standard"},
        )