Component handlers

class ComponentHandlerRegistry:
    """Registry for component handlers with O(1) lookup.

    Stores component handlers keyed by component name, providing
    registration, lookup, and existence checking operations.

    Attributes:
        _handlers: Internal dict mapping component names to handlers.

    Examples:
        Create and use a registry:

        ```python
        registry = ComponentHandlerRegistry()
        registry.register("instruction", InstructionHandler())
        handler = registry.get("instruction")
        ```

    See Also:
        - [`get_handler()`][gepa_adk.adapters.component_handlers.get_handler]:
            Convenience function for default registry.

    Note:
        A default registry instance is available as `component_handlers`
        module variable, with convenience functions `get_handler()` and
        `register_handler()`.
    """

    def __init__(self) -> None:
        """Initialize empty registry.

        Examples:
            ```python
            registry = ComponentHandlerRegistry()
            assert not registry.has("instruction")
            ```

        Note:
            Creates an empty internal dict for handler storage.
        """
        self._handlers: dict[str, ComponentHandler] = {}

    def register(self, name: str, handler: ComponentHandler) -> None:
        """Register a handler for a component name.

        Args:
            name: Component name (e.g., "instruction", "output_schema").
                Must be a non-empty string.
            handler: Handler implementing ComponentHandler protocol.

        Raises:
            ValueError: If name is empty or None.
            TypeError: If handler doesn't implement ComponentHandler protocol.

        Examples:
            ```python
            registry.register("instruction", InstructionHandler())
            registry.register("output_schema", OutputSchemaHandler())
            ```

        Note:
            Overwrites existing handler if name already registered.
            Logs a debug message on replacement.
        """
        if not name:
            raise ValueError("Component name must be a non-empty string")

        if not isinstance(handler, ComponentHandler):
            raise TypeError(
                f"Handler does not implement ComponentHandler protocol: {type(handler)}"
            )

        if name in self._handlers:
            logger.debug(
                "component_handler.register.replace",
                name=name,
                old_handler=type(self._handlers[name]).__name__,
                new_handler=type(handler).__name__,
            )

        self._handlers[name] = handler
        logger.debug(
            "component_handler.register.success",
            name=name,
            handler=type(handler).__name__,
        )

    def get(self, name: str) -> ComponentHandler:
        """Retrieve handler for component name.

        Args:
            name: Component name to look up.

        Returns:
            The registered ComponentHandler.

        Raises:
            ValueError: If name is empty or None.
            KeyError: If no handler registered for name.

        Examples:
            ```python
            handler = registry.get("instruction")
            original = handler.apply(agent, "New value")
            ```
        """
        if not name:
            raise ValueError("Component name must be a non-empty string")

        if name not in self._handlers:
            raise KeyError(f"No handler registered for component: {name}")

        return self._handlers[name]

    def has(self, name: str) -> bool:
        """Check if handler exists for component name.

        Args:
            name: Component name to check.

        Returns:
            True if handler registered, False otherwise.

        Examples:
            ```python
            if registry.has("instruction"):
                handler = registry.get("instruction")
            ```

        Note:
            Outputs False for empty/None names (no ValueError).
            This allows safe checking without exception handling.
        """
        if not name:
            return False
        return name in self._handlers

    def names(self) -> list[str]:
        """Return sorted list of registered handler names.

        Returns:
            Sorted list of component names with registered handlers.

        Examples:
            ```python
            available = registry.names()
            # ["generate_content_config", "instruction", "output_schema"]
            ```

        Note:
            Output is sorted alphabetically for consistent error messages
            and validation feedback.
        """
        return sorted(self._handlers.keys())

class InstructionHandler:
    """Handler for agent.instruction component.

    Manages serialization, application, and restoration of the
    agent's instruction (system prompt) during evolution.

    Examples:
        ```python
        handler = InstructionHandler()
        original = handler.serialize(agent)  # "Be helpful"
        handler.apply(agent, "Be concise")
        # ... evaluate ...
        handler.restore(agent, original)  # Back to "Be helpful"
        ```

    Note:
        All state is stored in the agent object - handler is stateless.
        No instance attributes are maintained.
    """

    def serialize(self, agent: "LlmAgent") -> str:
        """Extract instruction from agent as string.

        Args:
            agent: The LlmAgent instance.

        Returns:
            The agent's instruction as string.
            Returns empty string if instruction is None.

        Examples:
            ```python
            text = handler.serialize(agent)
            # text == "You are a helpful assistant."
            ```
        """
        instruction = agent.instruction
        if instruction is None:
            return ""
        return str(instruction)

    def apply(self, agent: "LlmAgent", value: str) -> str:
        """Apply new instruction to agent, return original.

        Args:
            agent: The LlmAgent instance to modify.
            value: The new instruction string.

        Returns:
            The original instruction value.

        Examples:
            ```python
            original = handler.apply(agent, "New instruction")
            # agent.instruction is now "New instruction"
            ```
        """
        original = self.serialize(agent)
        agent.instruction = value
        logger.debug(
            "instruction_handler.apply",
            original_preview=original[:50] if original else "",
            new_preview=value[:50] if value else "",
        )
        return original

    def restore(self, agent: "LlmAgent", original: str) -> None:
        """Restore original instruction to agent.

        Args:
            agent: The LlmAgent instance to restore.
            original: The original instruction value.

        Examples:
            ```python
            handler.restore(agent, original)
            # agent.instruction is back to original
            ```
        """
        agent.instruction = original
        logger.debug(
            "instruction_handler.restore",
            instruction_preview=original[:50] if original else "",
        )

class GenerateContentConfigHandler:
    """Handler for agent.generate_content_config component.

    Manages serialization, application, and restoration of the
    agent's LLM generation configuration during evolution.

    Examples:
        ```python
        handler = GenerateContentConfigHandler()
        original = handler.serialize(agent)  # YAML string
        handler.apply(agent, "temperature: 0.5")
        # ... evaluate ...
        handler.restore(agent, original_config)
        ```

    Note:
        All state is stored in the agent object - handler is stateless.
        On invalid config, logs warning and keeps original.
    """

    def serialize(self, agent: "LlmAgent") -> str:
        """Extract generate_content_config from agent as YAML.

        Args:
            agent: The LlmAgent instance.

        Returns:
            YAML string with parameter descriptions as comments.
            Returns empty string if generate_content_config is None.

        Examples:
            ```python
            yaml_str = handler.serialize(agent)
            # yaml_str contains YAML with temperature, top_p, etc.
            ```
        """
        config = getattr(agent, "generate_content_config", None)
        return serialize_generate_config(config)

    def apply(self, agent: "LlmAgent", value: str) -> Any:
        """Apply new generate_content_config to agent, return original.

        Args:
            agent: The LlmAgent instance to modify.
            value: YAML string defining the new config parameters.

        Returns:
            The original GenerateContentConfig (or None).

        Examples:
            ```python
            original = handler.apply(agent, "temperature: 0.5")
            # agent.generate_content_config.temperature is now 0.5
            ```

        Note:
            On deserialization or validation failure, logs warning and
            keeps original config. Never raises exceptions.
        """
        original = getattr(agent, "generate_content_config", None)

        try:
            new_config = deserialize_generate_config(value, original)

            # Validate the parsed config
            config_dict = {}
            for param in [
                "temperature",
                "top_p",
                "top_k",
                "max_output_tokens",
                "presence_penalty",
                "frequency_penalty",
            ]:
                param_value = getattr(new_config, param, None)
                if param_value is not None:
                    config_dict[param] = param_value

            errors = validate_generate_config(config_dict)
            if errors:
                logger.warning(
                    "generate_content_config_handler.apply.validation_failed",
                    errors=errors,
                    config_preview=value[:100] if value else "",
                )
                # Keep original - don't modify agent
                return original

            agent.generate_content_config = new_config
            logger.debug(
                "generate_content_config_handler.apply",
                original_temp=getattr(original, "temperature", None)
                if original
                else None,
                new_temp=getattr(new_config, "temperature", None),
            )
        except ConfigValidationError as e:
            logger.warning(
                "generate_content_config_handler.apply.failed",
                error=str(e),
                config_preview=value[:100] if value else "",
            )
            # Keep original - don't modify agent

        return original

    def restore(self, agent: "LlmAgent", original: Any) -> None:
        """Restore original generate_content_config to agent.

        Args:
            agent: The LlmAgent instance to restore.
            original: The original GenerateContentConfig (or None).

        Examples:
            ```python
            handler.restore(agent, original_config)
            # agent.generate_content_config is back to original
            ```
        """
        agent.generate_content_config = original
        logger.debug(
            "generate_content_config_handler.restore",
            has_config=original is not None,
        )

class OutputSchemaHandler:
    """Handler for agent.output_schema component.

    Manages serialization, application, and restoration of the
    agent's output schema (Pydantic model) during evolution.

    Attributes:
        _constraints: Optional SchemaConstraints for field preservation.

    Examples:
        ```python
        handler = OutputSchemaHandler()
        handler.set_constraints(SchemaConstraints(required_fields=("score",)))
        original_schema = handler.apply(agent, new_schema_text)
        # ... evaluate ...
        handler.restore(agent, original_schema)
        ```

    Note:
        Applies serialize_pydantic_schema and deserialize_schema utilities.
        On invalid schema text, keeps original and logs warning.
        When constraints are set, validates proposed schemas before applying.
    """

    def __init__(self) -> None:
        """Initialize handler with no constraints.

        Examples:
            ```python
            handler = OutputSchemaHandler()
            assert handler._constraints is None
            ```
        """
        self._constraints: SchemaConstraints | None = None

    def set_constraints(self, constraints: SchemaConstraints | None) -> None:
        """Set schema constraints for field preservation.

        Args:
            constraints: SchemaConstraints specifying required fields and
                type preservation rules. Pass None to clear constraints.

        Examples:
            ```python
            handler.set_constraints(SchemaConstraints(required_fields=("score",)))
            handler.set_constraints(None)  # Clear constraints
            ```

        Note:
            Once set, constraints are checked during apply() - proposed schemas
            that violate constraints will be rejected and the original kept.
        """
        self._constraints = constraints
        logger.debug(
            "output_schema_handler.set_constraints",
            has_constraints=constraints is not None,
            required_fields=constraints.required_fields if constraints else None,
        )

    def serialize(self, agent: "LlmAgent") -> str:
        """Extract output schema from agent as Python source.

        Args:
            agent: The LlmAgent instance.

        Returns:
            Python source code defining the schema class.
            Returns empty string if output_schema is None.

        Examples:
            ```python
            schema_text = handler.serialize(agent)
            # schema_text contains Python class definition
            ```
        """
        output_schema = getattr(agent, "output_schema", None)
        if output_schema is None:
            return ""
        try:
            return serialize_pydantic_schema(output_schema)
        except (TypeError, OSError) as e:
            logger.warning(
                "output_schema_handler.serialize.failed",
                error=str(e),
                schema_type=type(output_schema).__name__,
            )
            return ""

    def apply(self, agent: "LlmAgent", value: str) -> Any:
        """Apply new output schema to agent, return original.

        Args:
            agent: The LlmAgent instance to modify.
            value: Python source code defining the new schema.

        Returns:
            The original output_schema (class or None).

        Examples:
            ```python
            original = handler.apply(
                agent,
                '''
            class NewSchema(BaseModel):
                result: str
            ''',
            )
            ```

        Note:
            On deserialization failure, logs warning and keeps original.
            On constraint violation, keeps original.
            Never raises exceptions - graceful degradation.
        """
        original = getattr(agent, "output_schema", None)

        try:
            new_schema = deserialize_schema(value)

            # Validate against constraints if set
            if self._constraints is not None:
                is_valid, violations = validate_schema_against_constraints(
                    new_schema, original, self._constraints
                )
                if not is_valid:
                    logger.warning(
                        "output_schema_handler.apply.constraint_violation",
                        violations=violations,
                        schema_preview=value[:100] if value else "",
                    )
                    # Keep original - constraint violation
                    return original

            agent.output_schema = new_schema
            logger.debug(
                "output_schema_handler.apply",
                original_name=original.__name__ if original else None,
                new_name=new_schema.__name__,
            )
        except SchemaValidationError as e:
            logger.warning(
                "output_schema_handler.apply.failed",
                error=str(e),
                schema_preview=value[:100] if value else "",
            )
            # Keep original - don't modify agent

        return original

    def restore(self, agent: "LlmAgent", original: Any) -> None:
        """Restore original output schema to agent.

        Args:
            agent: The LlmAgent instance to restore.
            original: The original output_schema (class or None).

        Examples:
            ```python
            handler.restore(agent, original_schema)
            # agent.output_schema is back to original
            ```
        """
        agent.output_schema = original
        logger.debug(
            "output_schema_handler.restore",
            schema_name=original.__name__ if original else None,
        )