"""Streamer class for streaming responses with context manager support."""

import json
import logging
from typing import Optional, Dict, Any, List, TYPE_CHECKING

from .artifact import Artifact, ArtifactType

if TYPE_CHECKING:
    from .gateway_client import GatewayClient
    from .message import Message
    from .conversation_client import ConversationClient

logger = logging.getLogger(__name__)


class Streamer:
    """Handles streaming responses back to the conversation via the gateway (async)."""

    def __init__(
        self,
        agent_name: str,
        original_message: "Message",
        gateway_client: "GatewayClient",
        conversation_client: Optional["ConversationClient"] = None,
        awaiting: bool = False,
        send_to: str = "user",
    ):
        self._agent_name = agent_name
        self._original_message = original_message
        self._gateway_client = gateway_client
        self._conversation_client = conversation_client
        self._awaiting = awaiting  # Used by context manager
        self._send_to = send_to
        self._finished = False
        self._response_schema: Optional[Dict[str, Any]] = None
        self._message_metadata: Optional[Dict[str, Any]] = None
        self._artifacts: List[Artifact] = []
        self._artifacts_sent: bool = False

    def set_response_schema(self, schema: Dict[str, Any]) -> "Streamer":
        """Set response schema (fluent API)."""
        self._response_schema = schema
        return self

    def set_message_metadata(self, metadata: Dict[str, Any]) -> "Streamer":
        """Set message metadata (fluent API)."""
        self._message_metadata = metadata
        return self

    def add_artifact(
        self,
        url: str,
        type: str,
        title: Optional[str] = None,
        description: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> "Streamer":
        """Add an artifact to the message (fluent API).

        Args:
            url: URL to the artifact (file or iframe src)
            type: Type of artifact ('document', 'image', 'iframe', 'video', 'audio', 'code', 'link', 'file')
            title: Optional display title
            description: Optional description or alt text
            metadata: Optional metadata dict (size, mimeType, dimensions, etc.)

        Returns:
            self for method chaining

        Example:
            streamer.add_artifact(
                url="https://example.com/report.pdf",
                type="document",
                title="Monthly Report",
                metadata={"size": 1024000, "mimeType": "application/pdf"}
            )
        """
        artifact = Artifact(
            url=url,
            type=ArtifactType(type),
            title=title,
            description=description,
            metadata=metadata,
            position=len(self._artifacts),
        )
        self._artifacts.append(artifact)
        return self

    def set_artifacts(self, artifacts: List[Dict[str, Any]]) -> "Streamer":
        """Set all artifacts at once (replaces any existing).

        Args:
            artifacts: List of artifact dicts with url, type, title, description, metadata

        Returns:
            self for method chaining
        """
        self._artifacts = [
            Artifact(
                url=a["url"],
                type=ArtifactType(a["type"]),
                title=a.get("title"),
                description=a.get("description"),
                metadata=a.get("metadata"),
                position=i,
            )
            for i, a in enumerate(artifacts)
        ]
        return self

    def stream(
        self, content: str, persist: bool = True, event_type: Optional[str] = None
    ) -> None:
        """Stream a content chunk (not done).

        Args:
            content: The content to stream
            persist: If True, this content will be saved to DB. If False, only streamed.
            event_type: Optional event type for the chunk (e.g., 'thinking', 'tool_call').
        """
        if self._finished:
            raise RuntimeError("Streamer already finished")
        self._send(content, done=False, persist=persist, event_type=event_type)

    def write(self, content: str, persist: bool = True) -> None:
        """Alias for stream() - write a content chunk.

        Args:
            content: The content to write
            persist: If True, this content will be saved to DB. If False, only streamed.
        """
        self.stream(content, persist=persist)

    async def finish(self, awaiting: Optional[bool] = None) -> None:
        """Finish streaming with empty done=True chunk.

        Args:
            awaiting: If True, sets awaiting_route to this agent's name.
                     If None, uses the awaiting value from constructor.
        """
        if self._finished:
            return  # Idempotent

        self._send("", done=True, persist=True)
        self._finished = True

        # Use constructor's awaiting if not explicitly provided
        should_await = awaiting if awaiting is not None else self._awaiting

        # Store pending_response_schema if schema was set
        if self._response_schema and self._conversation_client:
            await self._set_pending_response_schema()

        if should_await and self._conversation_client:
            await self._set_awaiting_route()

    def _send(
        self,
        content: str,
        done: bool,
        persist: bool = True,
        event_type: Optional[str] = None,
    ) -> None:
        """Send message to router via gateway (sync - uses gRPC)."""
        message_body = {"content": content, "done": done, "persist": persist}
        if event_type:
            message_body["eventType"] = event_type
        headers = {
            "conversationId": self._original_message.conversation_id,
            "userId": self._original_message.user_id,
            "from_": self._agent_name,
            "to_": self._send_to,
            "nextRoute": self._send_to,
        }
        if self._response_schema:
            headers["responseSchema"] = json.dumps(self._response_schema)
        if self._message_metadata:
            headers["messageMetadata"] = json.dumps(self._message_metadata)

        # Send artifacts with first chunk (only once)
        if self._artifacts and not self._artifacts_sent:
            headers["artifacts"] = json.dumps([a.to_dict() for a in self._artifacts])
            self._artifacts_sent = True

        # Produce via gateway gRPC (synchronous operation)
        self._gateway_client.produce(
            topic="router.inbox",
            key=self._original_message.conversation_id,
            headers=headers,
            body=message_body,
        )

    async def _set_pending_response_schema(self) -> None:
        """Store response schema on conversation for form validation."""
        try:
            await self._conversation_client.update_conversation(
                self._original_message.conversation_id,
                {"pending_response_schema": self._response_schema},
            )
        except Exception as e:
            logger.error(f"Failed to set pending_response_schema: {e}")

    async def _set_awaiting_route(self) -> None:
        """Set awaiting_route to this agent's name."""
        try:
            await self._conversation_client.update_conversation(
                self._original_message.conversation_id,
                {"awaiting_route": self._agent_name},
            )
        except Exception as e:
            logger.error(f"Failed to set awaiting_route: {e}")

    async def __aenter__(self) -> "Streamer":
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        if not self._finished:
            await self.finish(awaiting=self._awaiting)
        return False  # Don't suppress exceptions
