"""
A Demo Agent Extension for CLAIA.
"""

# External dependencies
import logging
import pluggy
from typing import Type

# Internal dependencies
from claia.lib import BaseAgent, Process
from claia.hooks import AgentHooks, AgentInfo


########################################################################
#                              CONSTANTS                               #
########################################################################
SYSTEM_PROMPT = """
You are Claia, a helpful assistant.
"""


########################################################################
#                            INITIALIZATION                            #
########################################################################
hookimpl = pluggy.HookimplMarker("claia_agents")


########################################################################
#                            BOB AGENT CLASS                           #
########################################################################
class DemoAgent(BaseAgent):
  """
  Demo agent class
  """

  @classmethod
  def process_request(cls, process, registry=None, **kwargs) -> object:
    """
    Process a request using the demo agent's system prompt.

    Args:
        process: The process to execute
        registry: Registry instance to use for model operations
        **kwargs: Additional keyword arguments

    Returns:
        The updated process with results or error information
    """
    try:
      # Get the model ID from the validated parameters
      model_id = process.parameters["model_id"]

      # Set the demo agent's system prompt if it's different
      if process.conversation.prompt.get("system", "") != SYSTEM_PROMPT:
        process.conversation.change_prompt(SYSTEM_PROMPT)

      # Run the model with the conversation using the provided model registry
      result = registry.run(model_id, process.conversation, **kwargs)

      if result.is_error():
        raise ValueError(f"Demo agent ran into a problem: {result.get_message()}")

      process.mark_completed(result.data)

    except Exception as e:
      logging.exception(f"Demo agent encountered an error for {process.id}: {str(e)}")
      process.mark_failed(str(e))

    return process


########################################################################
#                            PLUGIN HOOKS                              #
########################################################################
class DemoAgentPlugin:
  """Plugin implementation for the demo agent."""

  @hookimpl
  def get_agent_class(self, agent_name: str) -> Type[BaseAgent]:
    """Get the agent class for the demo agent."""
    if agent_name.lower() == "demo":
      return DemoAgent
    return None

  @hookimpl
  def get_agent_info(self) -> AgentInfo:
    """Get information about the demo agent."""
    return AgentInfo(
      name="demo",
      description="A demo agent",
      agent_class=DemoAgent
    )
