# Copyright (c) 2026 Huawei Technologies Co., Ltd.
# All Rights Reserved.
#
# SPDX-License-Identifier: Apache-2.0

"""ExtensionSender -- one-shot pre-positioning facade over A2ATransport.

Single responsibility: send Authorization-T / Notification-T (and any
other one-shot extension) messages to agents BEFORE the workflow starts.
Bypasses Task-T prompt generation and the Negotiation-T auto-loop, and
does not emit events through the global EventCallback (the returned
result is the callback).

Kept separate from WorkflowEngineClient so a caller that only wants to
pre-position is not forced to hold a workflow-machinery facade.
"""

from typing import Dict, Any, Optional
from loguru import logger

from a2at_engine.client.a2a_transport import A2ATransport
from a2at_engine.client.extensions import A2ATExtension
from a2at_engine.core.models import SendMessageResult


class ExtensionSender:
    """One-shot extension message sender built on a shared A2ATransport."""

    def __init__(self, transport: A2ATransport):
        self._transport = transport

    @property
    def transport(self) -> A2ATransport:
        return self._transport

    def get_a2at_client(self):
        return self._transport.get_a2at_client()

    async def send_extension_message(
        self,
        agent_name: str,
        instruction: str,
        natural_language_input: str,
        extension: A2ATExtension,
    ) -> SendMessageResult:
        """Send a one-shot extension message for pre-positioning.

        The metadata value is generated by the A2A-T SDK (LLM + prompt
        template) from the natural-language input; if the SDK cannot
        generate, the input text is used as-is.
        """
        agent_card = self._transport.get_card(agent_name)
        if not agent_card:
            raise RuntimeError(f"Agent not found: {agent_name}")
        metadata_value = self._generate_extension_prompt(extension, natural_language_input)
        if not metadata_value:
            metadata_value = natural_language_input
            logger.info(f"[ExtensionSender] SDK prompt generation unavailable for {agent_name} ({extension.display_name}), using input as metadata")
        logger.info(f"[ExtensionSender] sendExtensionMessage to {agent_name}: extension={extension.display_name}, metadataValue={len(metadata_value)} chars")
        metadata = {extension.uri: metadata_value}
        client = self._transport.create_a2a_client(agent_card)
        send_req = self._transport.build_send_request(instruction, None, metadata)
        response_text, last_task, last_meta, task_state = (
            await self._transport.consume_stream(client, send_req)
        )
        if response_text is None and last_task is not None:
            response_text = str(last_task)
        result = SendMessageResult(
            text=response_text or "", task=last_task,
            metadata=last_meta, task_state=task_state,
        )
        logger.info(f"[ExtensionSender] Extension response from {agent_name}: state={result.task_state}")
        return result

    async def send_authorization(
        self, agent_name: str, instruction: str, natural_language_input: str,
    ) -> SendMessageResult:
        """Convenience for Authorization-T pre-positioning."""
        return await self.send_extension_message(
            agent_name, instruction, natural_language_input, A2ATExtension.AUTHORIZATION_T)

    async def send_notification(
        self, agent_name: str, instruction: str, natural_language_input: str,
    ) -> SendMessageResult:
        """Convenience for Notification-T pre-positioning."""
        return await self.send_extension_message(
            agent_name, instruction, natural_language_input, A2ATExtension.NOTIFICATION_T)

    # ------------------------------------------------------------------
    # Extension prompt generation dispatch
    # ------------------------------------------------------------------

    def _generate_extension_prompt(self, extension, natural_language_input):
        """Dispatch to the SDK extension-specific prompt generation."""
        if extension == A2ATExtension.TASK_T:
            return self.generate_prompt_text(natural_language_input)
        if extension == A2ATExtension.NEGOTIATION_T:
            return self.generate_negotiation_prompt(natural_language_input)
        if extension == A2ATExtension.AUTHORIZATION_T:
            return self.generate_authorization_prompt(natural_language_input)
        if extension == A2ATExtension.NOTIFICATION_T:
            return self.generate_notification_prompt(natural_language_input)
        return ''

    def generate_prompt_text(self, natural_language_input: str) -> str:
        """Generate structured Task-T prompt text from natural-language input."""
        a2at_client = self._transport.get_a2at_client()
        if not a2at_client:
            return ''
        try:
            prompt_result = a2at_client.generate_task_prompt(natural_language_input)
            if hasattr(prompt_result, 'success') and prompt_result.success:
                text = getattr(prompt_result, 'prompt_text', None)
                if text:
                    return text
            else:
                failure = getattr(prompt_result, 'failure', None)
                if failure:
                    logger.warning('[ExtensionSender] SDK Task-T prompt generation failed: ' + str(getattr(failure, 'message', '')))
        except Exception as e:
            logger.warning('[ExtensionSender] SDK Task-T prompt generation error: ' + str(e))
        return ''

    def generate_negotiation_prompt(self, natural_language_input: str) -> str:
        """Generate Negotiation-T prompt text. Reserved for SDK support."""
        if not self._transport.get_a2at_client():
            return ''
        return ''

    def generate_authorization_prompt(self, natural_language_input: str) -> str:
        """Generate Authorization-T prompt text. Reserved for SDK support."""
        if not self._transport.get_a2at_client():
            return ''
        return ''

    def generate_notification_prompt(self, natural_language_input: str) -> str:
        """Generate Notification-T prompt text. Reserved for SDK support."""
        if not self._transport.get_a2at_client():
            return ''
        return ''