# main.py
import os
import uvicorn

from a2a.server.apps import A2AStarletteApplication
from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.tasks import InMemoryTaskStore

# Typed models expected by the server
from a2a.types import AgentCard, AgentCapabilities, AgentSkill


# ---- Minimal agent executor --------------------------------------------------
class HelloWorldAgentExecutor:
    async def call(self, message: str, **kwargs):
        # Return whatever your A2A client expects (text is fine for a demo)
        return {"text": f"Hello from A2A on HF Spaces! You said: {message}"}


def build_app():
    # At least one skill is recommended
    hello_skill = AgentSkill(
        id="hello",
        name="Say Hello",
        description="Return a friendly greeting",
        tags=["hello", "demo"],
        examples=["hi", "hello there", "yo"],
    )

    # Optional public URL (HF sets SPACE_HOST, but this is just metadata)
    public_url = os.getenv("SPACE_HOST")
    if public_url:
        public_url = f"https://{public_url}/"
    else:
        public_url = ""

    agent_card = AgentCard(
        name="HF A2A Agent",
        description="A2A demo running on Hugging Face Spaces",
        version="1.0.0",
        url=public_url,
        protocol_version="0.3.0",
        default_input_modes=["text/plain", "application/json"],
        default_output_modes=["text/plain", "application/json"],
        capabilities=AgentCapabilities(streaming=False),
        skills=[hello_skill],
        # Keep it simple; no extended card
        supports_authenticated_extended_card=False,
    )

    app = A2AStarletteApplication(
        agent_card=agent_card,
        http_handler=DefaultRequestHandler(
            agent_executor=HelloWorldAgentExecutor(),
            task_store=InMemoryTaskStore(),
        ),
    ).build()

    # ---- Add GET endpoints so HF/browser doesn't 405 on "/" ------------------
    from starlette.responses import HTMLResponse, PlainTextResponse

    async def landing(_):
        return HTMLResponse(
            "<h3>HF A2A Agent</h3>"
            "<p>Health: <a href=\"/health\">/health</a></p>"
            "<p>Send a message (POST JSON) to <code>/message/send</code>:</p>"
            "<pre>{\"message\": \"hello\"}</pre>"
        )

    async def health(_):
        return PlainTextResponse("ok")

    app.router.add_route("/", landing, methods=["GET"])
    app.router.add_route("/health", health, methods=["GET"])

    return app


if __name__ == "__main__":
    port = int(os.getenv("PORT", "7860"))  # HF provides PORT
    uvicorn.run(build_app(), host="0.0.0.0", port=port)