"""Schema registry for managing extraction schemas."""

from dataclasses import dataclass, field
from typing import Any

from maeris_mcp.types.protocol import ExtractionInstructions, OperationType


@dataclass
class Schema:
    """Represents a registered extraction schema."""

    id: str
    name: str
    description: str
    json_schema: dict[str, Any]
    applicable_operations: list[OperationType]
    instructions: ExtractionInstructions


@dataclass
class SchemaRegistry:
    """Manages all extraction schemas."""

    _schemas: dict[str, Schema] = field(default_factory=dict)

    def __post_init__(self) -> None:
        """Register built-in schemas after initialization."""
        self._register_builtin_schemas()

    def get(self, schema_id: str) -> Schema | None:
        """Get a schema by ID."""
        return self._schemas.get(schema_id)

    def get_for_operation(self, op: OperationType) -> list[Schema]:
        """Get all schemas applicable to an operation."""
        return [
            schema
            for schema in self._schemas.values()
            if op in schema.applicable_operations
        ]

    def register(self, schema: Schema) -> None:
        """Register a new schema."""
        self._schemas[schema.id] = schema

    def _register_builtin_schemas(self) -> None:
        """Register the default extraction schemas."""
        from maeris_mcp.schemas.api import api_extraction_schema
        from maeris_mcp.schemas.component import component_extraction_schema
        from maeris_mcp.schemas.project import project_extraction_schema
        from maeris_mcp.schemas.search import search_schema

        self.register(component_extraction_schema())
        self.register(project_extraction_schema())
        self.register(search_schema())
        self.register(api_extraction_schema())
