"""Demo Agent 3: Orchestrator Agent -- this is the one that matters for the demo.

It's an A2A *server* (external callers send it a plan_trip request) that is
also an A2A *client* (it calls the Weather Agent and Translator Agent over
A2A, using its own scoped credentials for each, and combines the results).

This is the concrete proof that a2a-passport's credential-resolution design
works: one agent, two different outbound calls, two different
audiences/tokens, resolved automatically by RemoteAgentClient -- the
executor code below never touches a token directly. It's also proof of the
federation story: the two OAuth2ClientCredentialsProvider instances below
could each point at a *different* IdP (e.g. one at the demo mock IdP, one at
a real Keycloak realm) with zero other code changes.
"""
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))

import uvicorn
from a2a.helpers import get_data_parts, new_data_part, new_task_from_user_message
from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.events import EventQueue
from a2a.server.tasks import TaskUpdater
from a2a.types import AgentCapabilities, AgentCard, AgentInterface, AgentSkill
from a2a.types.a2a_pb2 import TaskState

from a2a_passport import OAuth2ClientCredentialsProvider, RemoteAgentClient, build_agent_app
from a2a_passport.auth.verifiers import SharedSecretVerifier
from a2a_passport.testing import JWT_SECRET

PORT = 9203
AUDIENCE = "orchestrator-agent"
REQUIRED_SCOPE = "orchestrator:invoke"

IDP_TOKEN_URL = "http://127.0.0.1:9000/token"
WEATHER_URL = "http://127.0.0.1:9201"
TRANSLATOR_URL = "http://127.0.0.1:9202"

# The orchestrator's OWN outbound credentials -- used to call weather/translator.
# In production these come from a secrets manager, not a literal in source.
_ORCH_TOKEN_PROVIDER = OAuth2ClientCredentialsProvider(
    token_url=IDP_TOKEN_URL,
    client_id="orchestrator-agent",
    client_secret="orchestrator-secret",
)

_WEATHER_TO_LANGUAGE = {
    "paris": "french",
    "tokyo": "japanese",
    "bengaluru": "hindi",
}


class OrchestratorExecutor(AgentExecutor):
    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        task = context.current_task or new_task_from_user_message(context.message)
        if not context.current_task:
            await event_queue.enqueue_event(task)
        updater = TaskUpdater(event_queue=event_queue, task_id=task.id, context_id=task.context_id)
        await updater.update_status(state=TaskState.TASK_STATE_WORKING)

        incoming = get_data_parts(context.message.parts)
        request = incoming[0] if incoming else {}
        city = request.get("city", "")

        # --- Hop 1: call the Weather Agent (audience="weather-agent") ---
        async with RemoteAgentClient(
            WEATHER_URL, _ORCH_TOKEN_PROVIDER, audience="weather-agent"
        ) as weather_client:
            weather_result = await weather_client.send_data({"city": city})

        # --- Hop 2: call the Translator Agent (audience="translator-agent") ---
        target_language = _WEATHER_TO_LANGUAGE.get(str(city).strip().lower(), "french")
        async with RemoteAgentClient(
            TRANSLATOR_URL, _ORCH_TOKEN_PROVIDER, audience="translator-agent"
        ) as translator_client:
            translation_result = await translator_client.send_data(
                {"text": weather_result.get("condition"), "target_language": target_language}
            )

        combined = {
            "city": city,
            "weather": weather_result,
            "translated_condition": translation_result.get("translated_text"),
            "target_language": target_language,
        }
        await updater.add_artifact(parts=[new_data_part(data=combined)], name="trip_plan")
        await updater.update_status(state=TaskState.TASK_STATE_COMPLETED)

    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
        raise NotImplementedError("cancel is not supported by the orchestrator demo")


def build_app():
    skill = AgentSkill(
        id="plan_trip",
        name="Plan Trip",
        description="Given a city, fetches weather and returns the condition translated "
        "into that city's local language -- by calling two other agents over A2A.",
        input_modes=["application/json"],
        output_modes=["application/json"],
        tags=["demo", "orchestration"],
        examples=['{"city": "Paris"}'],
    )
    card = AgentCard(
        name="Trip Orchestrator Agent",
        description="Demo agent that chains calls to the Weather and Translator agents.",
        version="0.1.0",
        default_input_modes=["application/json"],
        default_output_modes=["application/json"],
        capabilities=AgentCapabilities(streaming=False),
        supported_interfaces=[AgentInterface(protocol_binding="JSONRPC", url=f"http://127.0.0.1:{PORT}")],
        skills=[skill],
    )
    return build_agent_app(
        card=card,
        executor=OrchestratorExecutor(),
        verifier=SharedSecretVerifier(JWT_SECRET),
        expected_audience=AUDIENCE,
        required_scope=REQUIRED_SCOPE,
        service_name="orchestrator-agent",
    )


if __name__ == "__main__":
    uvicorn.run(build_app(), host="127.0.0.1", port=PORT)
