# FILE: pocketflow_a2a_agent/__main__.py
import click
import logging
import os

# Import from the common code you copied
from common.server import A2AServer
from common.types import AgentCard, AgentCapabilities, AgentSkill, MissingAPIKeyError

# Import your custom TaskManager (which now imports from your original files)
from task_manager import PocketFlowTaskManager

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@click.command()
@click.option("--host", "host", default="localhost")
@click.option("--port", "port", default=10003) # Use a different port from other agents
def main(host, port):
    """Starts the PocketFlow A2A Agent server."""
    try:
        # Check for necessary API keys (add others if needed)
        if not os.getenv("OPENAI_API_KEY"):
            raise MissingAPIKeyError("OPENAI_API_KEY environment variable not set.")

        # --- Define the Agent Card ---
        capabilities = AgentCapabilities(
            streaming=False, # This simple implementation is synchronous
            pushNotifications=False,
            stateTransitionHistory=False # PocketFlow state isn't exposed via A2A history
        )
        skill = AgentSkill(
            id="web_research_qa",
            name="Web Research and Answering",
            description="Answers questions using web search results when necessary.",
            tags=["research", "qa", "web search"],
            examples=[
                "Who won the Nobel Prize in Physics 2024?",
                "What is quantum computing?",
                "Summarize the latest news about AI.",
            ],
            # Input/Output modes defined in the TaskManager
            inputModes=PocketFlowTaskManager.SUPPORTED_CONTENT_TYPES,
            outputModes=PocketFlowTaskManager.SUPPORTED_CONTENT_TYPES,
        )
        agent_card = AgentCard(
            name="PocketFlow Research Agent (A2A Wrapped)",
            description="A simple research agent based on PocketFlow, made accessible via A2A.",
            url=f"http://{host}:{port}/", # The endpoint A2A clients will use
            version="0.1.0-a2a",
            capabilities=capabilities,
            skills=[skill],
            # Assuming no specific provider or auth for this example
            provider=None,
            authentication=None,
            defaultInputModes=PocketFlowTaskManager.SUPPORTED_CONTENT_TYPES,
            defaultOutputModes=PocketFlowTaskManager.SUPPORTED_CONTENT_TYPES,
        )

        # --- Initialize and Start Server ---
        task_manager = PocketFlowTaskManager() # Instantiate your custom manager
        server = A2AServer(
            agent_card=agent_card,
            task_manager=task_manager,
            host=host,
            port=port,
        )

        logger.info(f"Starting PocketFlow A2A server on http://{host}:{port}")
        server.start()

    except MissingAPIKeyError as e:
        logger.error(f"Configuration Error: {e}")
        exit(1)
    except Exception as e:
        logger.error(f"An error occurred during server startup: {e}", exc_info=True)
        exit(1)


if __name__ == "__main__":
    main()