"""Streaming utilities for agent and team execution"""

from typing import Dict, Any, Callable, Optional, List
import structlog

logger = structlog.get_logger()


class StreamingHelper:
    """
    Helper for handling streaming from Agno Agent/Team executions.

    Provides utilities for:
    - Publishing events to Control Plane
    - Tracking run_id from streaming chunks
    - Collecting response content
    - Publishing tool execution events
    - Handling member message streaming
    - Tracking tool IDs for proper start/complete matching
    - Splitting assistant messages into pre-tool and post-tool phases
    """

    def __init__(self, control_plane_client, execution_id: str):
        self.control_plane = control_plane_client
        self.execution_id = execution_id
        self.run_id_published = False
        self.response_content = []
        self.member_message_ids = {}  # Track message_id per member
        self.active_streaming_member = None  # Track which member is streaming
        self.tool_execution_ids = {}  # Track tool IDs for matching start/complete events
        self.tool_messages = []  # Track tool messages for session persistence

        # NEW: Track message phases for proper assistant message splitting
        self.pre_tool_content = []  # Content before first tool use
        self.post_tool_content = []  # Content after tools complete
        self.tool_phase = "pre"  # Current phase: "pre", "during", or "post"
        self.first_tool_timestamp = None  # Timestamp when first tool started
        self.tools_complete_timestamp = None  # Timestamp when all tools completed
        self.has_any_tools = False  # Track if any tools were executed

    def handle_run_id(self, chunk: Any, on_run_id: Optional[Callable[[str], None]] = None) -> None:
        """
        Capture and publish run_id from first streaming chunk.

        Args:
            chunk: Streaming chunk from Agno
            on_run_id: Optional callback when run_id is captured
        """
        if not self.run_id_published and hasattr(chunk, 'run_id') and chunk.run_id:
            run_id = chunk.run_id

            logger.info("run_id_captured", execution_id=self.execution_id[:8], run_id=run_id[:16])

            # Publish to Control Plane for UI
            self.control_plane.publish_event(
                execution_id=self.execution_id,
                event_type="run_started",
                data={
                    "run_id": run_id,
                    "execution_id": self.execution_id,
                    "cancellable": True,
                }
            )

            self.run_id_published = True

            # Call callback if provided (for cancellation manager)
            if on_run_id:
                on_run_id(run_id)

    async def handle_content_chunk(
        self,
        chunk: Any,
        message_id: str,
        print_to_console: bool = True
    ) -> Optional[str]:
        """
        Handle content chunk from streaming response.

        Tracks content in different phases (pre-tool, during-tool, post-tool)
        to enable proper message splitting around tool usage.

        Args:
            chunk: Streaming chunk
            message_id: Unique message ID for this turn
            print_to_console: Whether to print to stdout

        Returns:
            Content string if present, None otherwise
        """
        # Check for both 'response' (RuntimeExecutionResult) and 'content' (legacy/Agno)
        content = None

        # DEBUG: Log what attributes the chunk has
        print(f"[DEBUG] StreamingHelper.handle_content_chunk: chunk type = {type(chunk).__name__}")
        print(f"[DEBUG] StreamingHelper.handle_content_chunk: has 'response' = {hasattr(chunk, 'response')}")
        print(f"[DEBUG] StreamingHelper.handle_content_chunk: has 'content' = {hasattr(chunk, 'content')}")

        if hasattr(chunk, 'response') and chunk.response:
            content = str(chunk.response)
            print(f"[DEBUG] StreamingHelper.handle_content_chunk: extracted from 'response': {repr(content[:100])}")
        elif hasattr(chunk, 'content') and chunk.content:
            content = str(chunk.content)
            print(f"[DEBUG] StreamingHelper.handle_content_chunk: extracted from 'content': {repr(content[:100])}")
        else:
            print(f"[DEBUG] StreamingHelper.handle_content_chunk: NO CONTENT FOUND!")

        if content:
            # Track content in appropriate phase
            if self.tool_phase == "pre":
                self.pre_tool_content.append(content)
            elif self.tool_phase == "post":
                self.post_tool_content.append(content)
            # Note: During "during" phase, we don't collect assistant content
            # as tools are executing

            self.response_content.append(content)

            if print_to_console:
                print(content, end='', flush=True)

            # Stream to Control Plane for real-time UI updates (NON-BLOCKING)
            await self.control_plane.publish_event_async(
                execution_id=self.execution_id,
                event_type="message_chunk",
                data={
                    "role": "assistant",
                    "content": content,
                    "is_chunk": True,
                    "message_id": message_id,
                    "phase": self.tool_phase,  # NEW: Include current phase
                }
            )

            return content

        return None

    def get_full_response(self) -> str:
        """Get the complete response accumulated from all chunks."""
        return ''.join(self.response_content)

    def handle_member_content_chunk(
        self,
        member_name: str,
        content: str,
        print_to_console: bool = True
    ) -> str:
        """
        Handle content chunk from a team member.

        Args:
            member_name: Name of the team member
            content: Content string
            print_to_console: Whether to print to stdout

        Returns:
            The member's message_id
        """
        import time

        # Generate unique message ID for this member if not exists
        if member_name not in self.member_message_ids:
            self.member_message_ids[member_name] = f"{self.execution_id}_{member_name}_{int(time.time() * 1000000)}"

            # Print member name header once when they start
            if print_to_console:
                print(f"\n[{member_name}] ", end='', flush=True)

        # If switching to a different member, mark the previous one as complete
        if self.active_streaming_member and self.active_streaming_member != member_name:
            self.publish_member_complete(self.active_streaming_member)

        # Track that this member is now actively streaming
        self.active_streaming_member = member_name

        # Print content without repeated member name prefix
        if print_to_console:
            print(content, end='', flush=True)

        # Stream member chunk to Control Plane
        message_id = self.member_message_ids[member_name]
        self.control_plane.publish_event(
            execution_id=self.execution_id,
            event_type="member_message_chunk",
            data={
                "role": "assistant",
                "content": content,
                "is_chunk": True,
                "message_id": message_id,
                "source": "team_member",
                "member_name": member_name,
            }
        )

        return message_id

    def publish_member_complete(self, member_name: str) -> None:
        """
        Publish member_message_complete event and clear the message_id.

        Args:
            member_name: Name of the member to mark as complete
        """
        if member_name in self.member_message_ids:
            self.control_plane.publish_event(
                execution_id=self.execution_id,
                event_type="member_message_complete",
                data={
                    "message_id": self.member_message_ids[member_name],
                    "member_name": member_name,
                    "source": "team_member",
                }
            )

            # CRITICAL: Clear the message_id for this member after completing
            # This ensures a NEW message_id is generated for their next turn
            # Without this, all turns from the same member would edit the same message!
            del self.member_message_ids[member_name]
            logger.info("member_message_id_cleared", member_name=member_name, execution_id=self.execution_id[:8])

    def finalize_streaming(self) -> None:
        """
        Finalize streaming by marking any active member as complete.
        Call this when streaming ends.
        """
        if self.active_streaming_member:
            self.publish_member_complete(self.active_streaming_member)
            self.active_streaming_member = None

    def get_tool_messages(self) -> List[Dict[str, Any]]:
        """
        Get all tool messages collected during streaming for session persistence.

        Returns:
            List of tool message dicts with role='system', tool metadata, and timestamps
        """
        return self.tool_messages

    def publish_tool_start(
        self,
        tool_name: str,
        tool_execution_id: str,
        tool_args: Optional[Dict[str, Any]] = None,
        source: str = "agent",
        member_name: Optional[str] = None
    ) -> str:
        """
        Publish tool execution start event.

        Also transitions message phase from "pre" to "during" on first tool use.

        Args:
            tool_name: Name of the tool
            tool_execution_id: Unique ID for this tool execution
            tool_args: Tool arguments
            source: "agent" or "team_member" or "team_leader"  or "team"
            member_name: Name of member (if tool is from a member)

        Returns:
            message_id for this tool execution
        """
        import time
        from datetime import datetime, timezone

        # Mark transition to "during" phase on first tool
        if self.tool_phase == "pre":
            self.tool_phase = "during"
            self.first_tool_timestamp = datetime.now(timezone.utc).isoformat()
            self.has_any_tools = True
            logger.info(
                "phase_transition_to_during",
                execution_id=self.execution_id[:8],
                tool_name=tool_name,
                pre_tool_content_length=len(''.join(self.pre_tool_content))
            )

        message_id = f"{self.execution_id}_tool_{tool_execution_id}"
        is_member_tool = member_name is not None
        parent_message_id = self.member_message_ids.get(member_name) if is_member_tool else None

        # Store tool info for matching with completion event
        tool_key = f"{member_name or 'leader'}_{tool_name}_{int(time.time())}"
        self.tool_execution_ids[tool_key] = {
            "tool_execution_id": tool_execution_id,
            "message_id": message_id,
            "tool_name": tool_name,
            "member_name": member_name,
            "parent_message_id": parent_message_id,
            "tool_args": tool_args,  # Store args for persistence
        }

        event_type = "member_tool_started" if is_member_tool else "tool_started"

        self.control_plane.publish_event(
            execution_id=self.execution_id,
            event_type=event_type,
            data={
                "tool_name": tool_name,
                "tool_execution_id": tool_execution_id,
                "message_id": message_id,
                "tool_arguments": tool_args,
                "source": "team_member" if is_member_tool else "team_leader",
                "member_name": member_name,
                "parent_message_id": parent_message_id,
                "message": f"🔧 Executing tool: {tool_name}",
            }
        )

        return message_id

    def publish_tool_complete(
        self,
        tool_name: str,
        tool_execution_id: str,
        status: str = "success",
        output: Optional[str] = None,
        error: Optional[str] = None,
        source: str = "agent",
        member_name: Optional[str] = None
    ) -> None:
        """
        Publish tool execution completion event.

        Args:
            tool_name: Name of the tool
            tool_execution_id: Unique ID for this tool execution
            status: "success" or "failed"
            output: Tool output (if successful)
            error: Error message (if failed)
            source: "agent" or "team_member" or "team_leader" or "team"
            member_name: Name of member (if tool is from a member)
        """
        import time

        # Find the stored tool info from the start event
        tool_key_pattern = f"{member_name or 'leader'}_{tool_name}"
        matching_tool = None
        for key, tool_info in list(self.tool_execution_ids.items()):
            if key.startswith(tool_key_pattern):
                matching_tool = tool_info
                # Remove from tracking dict
                del self.tool_execution_ids[key]
                break

        if matching_tool:
            message_id = matching_tool["message_id"]
            parent_message_id = matching_tool["parent_message_id"]
            # Use the stored tool_execution_id from the start event
            tool_execution_id = matching_tool["tool_execution_id"]
            tool_args = matching_tool.get("tool_args")  # Get stored args
        else:
            # Fallback if start event wasn't captured
            message_id = f"{self.execution_id}_tool_{tool_execution_id}"
            parent_message_id = self.member_message_ids.get(member_name) if member_name else None
            tool_args = None
            logger.warning("tool_completion_without_start", tool_name=tool_name, member_name=member_name)

        is_member_tool = member_name is not None
        event_type = "member_tool_completed" if is_member_tool else "tool_completed"

        tool_data = {
            "tool_name": tool_name,
            "tool_execution_id": tool_execution_id,  # Now uses the stored ID from start event
            "message_id": message_id,
            "status": status,
            "tool_output": output[:50000] if output else None,  # Increased from 1000 to 50000 to preserve Metabase URLs and other important data
            "tool_error": error,
            "source": "team_member" if is_member_tool else "team_leader",
            "member_name": member_name,
            "parent_message_id": parent_message_id,
            "message": f"{'✅' if status == 'success' else '❌'} Tool {status}: {tool_name}",
        }

        self.control_plane.publish_event(
            execution_id=self.execution_id,
            event_type=event_type,
            data=tool_data
        )

        # Store tool message for session persistence
        # NEW: Tool messages are now role="user" with tool_result content blocks
        # This aligns with Claude API best practices where tool results come from user
        from datetime import datetime, timezone

        tool_result_content = output[:50000] if output and status == "success" else (error or "")

        self.tool_messages.append({
            "role": "user",  # CHANGED: Tool results are user role (Claude API format)
            "content": [
                {
                    "type": "tool_result",
                    "tool_use_id": tool_execution_id,
                    "content": tool_result_content,
                    "is_error": status != "success",
                }
            ],
            "message_type": "tool_result",  # NEW: Mark as tool result message
            "tool_name": tool_name,
            "tool_execution_id": tool_execution_id,
            "tool_input": tool_args,  # Frontend expects "tool_input" not "tool_arguments"
            "tool_output": output[:50000] if output else None,
            "tool_error": error,
            "tool_status": status,  # Frontend expects "tool_status" not "status"
            "member_name": member_name,
            "message_id": message_id,
            "parent_message_id": parent_message_id,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        })

        logger.info(
            "tool_result_message_created",
            execution_id=self.execution_id[:8],
            tool_name=tool_name,
            status=status,
            has_output=bool(output),
            has_error=bool(error)
        )

    def transition_to_post_tool_phase(self):
        """
        Transition from "during" phase to "post" phase.

        Call this after all tools have completed to start collecting
        post-tool assistant content.
        """
        if self.tool_phase == "during":
            from datetime import datetime, timezone
            self.tool_phase = "post"
            self.tools_complete_timestamp = datetime.now(timezone.utc).isoformat()
            logger.info(
                "phase_transition_to_post",
                execution_id=self.execution_id[:8],
                tools_completed=len(self.tool_messages)
            )

    def finalize_streaming(self) -> None:
        """
        Finalize streaming by marking any active member as complete
        and transitioning to post-tool phase if needed.

        Call this when streaming ends.
        """
        # Transition to post phase if we had tools
        if self.has_any_tools and self.tool_phase == "during":
            self.transition_to_post_tool_phase()

        # Handle team member completion
        if self.active_streaming_member:
            self.publish_member_complete(self.active_streaming_member)
            self.active_streaming_member = None

    def get_assistant_message_parts(self) -> List[Dict[str, Any]]:
        """
        Get assistant messages split into pre-tool and post-tool parts.

        Returns:
            List of assistant message dicts. May contain:
            - Pre-tool message (if content exists before first tool)
            - Post-tool message (if content exists after tools complete)
            - Single message (if no tools were used)
        """
        from datetime import datetime, timezone

        assistant_parts = []

        # If no tools were used, return single message with all content
        if not self.has_any_tools:
            full_content = ''.join(self.response_content)
            if full_content:
                assistant_parts.append({
                    "content": full_content,
                    "phase": "complete",
                    "timestamp": datetime.now(timezone.utc).isoformat(),
                })
            return assistant_parts

        # If tools were used, split into pre and post parts
        pre_content = ''.join(self.pre_tool_content)
        if pre_content:
            assistant_parts.append({
                "content": pre_content,
                "phase": "pre",
                "timestamp": self.first_tool_timestamp or datetime.now(timezone.utc).isoformat(),
            })

        post_content = ''.join(self.post_tool_content)
        if post_content:
            assistant_parts.append({
                "content": post_content,
                "phase": "post",
                "timestamp": self.tools_complete_timestamp or datetime.now(timezone.utc).isoformat(),
            })

        return assistant_parts

    def build_structured_messages(
        self,
        execution_id: str,
        turn_number: int,
        user_message: Dict[str, Any],
    ) -> List[Dict[str, Any]]:
        """
        Build structured message list with proper interleaving.

        Creates proper message flow:
        1. User message
        2. Assistant message (pre-tool) if tools were used
        3. Tool result messages (role="user")
        4. Assistant message (post-tool) if tools were used
        OR
        1. User message
        2. Assistant message (complete) if no tools

        Args:
            execution_id: Execution ID
            turn_number: Turn number in conversation
            user_message: User message dict (already constructed)

        Returns:
            List of messages in proper order with message_ids assigned
        """
        messages = [user_message]

        assistant_parts = self.get_assistant_message_parts()

        if not self.has_any_tools:
            # Simple case: no tools, single assistant message
            if assistant_parts:
                part = assistant_parts[0]
                messages.append({
                    "role": "assistant",
                    "content": part["content"],
                    "timestamp": part["timestamp"],
                    "message_id": f"{execution_id}_assistant_{turn_number}",
                })
        else:
            # Complex case: tools were used, split assistant messages
            for i, part in enumerate(assistant_parts):
                phase_suffix = f"_{part['phase']}" if part['phase'] in ['pre', 'post'] else ""
                messages.append({
                    "role": "assistant",
                    "content": part["content"],
                    "timestamp": part["timestamp"],
                    "message_id": f"{execution_id}_assistant_{turn_number}{phase_suffix}",
                    "phase": part["phase"],
                })

                # After pre-tool message, insert tool result messages
                if part["phase"] == "pre":
                    messages.extend(self.tool_messages)

            # If we only have pre-tool content (no post), still add tool messages at end
            if assistant_parts and assistant_parts[-1]["phase"] == "pre":
                messages.extend(self.tool_messages)
            # If we have no assistant parts but have tools, add tool messages
            elif not assistant_parts and self.tool_messages:
                messages.extend(self.tool_messages)

        logger.info(
            "structured_messages_built",
            execution_id=execution_id[:8],
            turn_number=turn_number,
            total_messages=len(messages),
            has_tools=self.has_any_tools,
            assistant_parts=len(assistant_parts),
            tool_messages=len(self.tool_messages)
        )

        return messages


def create_tool_hook(control_plane_client, execution_id: str):
    """
    Create a tool hook function for Agno Agent/Team.

    This hook is called before and after each tool execution
    to publish real-time updates to the Control Plane.

    Args:
        control_plane_client: Control Plane client instance
        execution_id: Execution ID

    Returns:
        Hook function compatible with Agno tool_hooks
    """
    import time

    def tool_hook(tool_name: str, tool_args: dict, result: Any = None, error: Exception = None):
        """Tool hook for real-time updates"""
        tool_execution_id = f"{tool_name}_{int(time.time() * 1000000)}"

        if error is None and result is None:
            # Tool starting
            control_plane_client.publish_event(
                execution_id=execution_id,
                event_type="tool_started",
                data={
                    "tool_name": tool_name,
                    "tool_execution_id": tool_execution_id,
                    "tool_arguments": tool_args,
                    "message": f"🔧 Starting: {tool_name}",
                }
            )
        else:
            # Tool completed
            status = "failed" if error else "success"
            control_plane_client.publish_event(
                execution_id=execution_id,
                event_type="tool_completed",
                data={
                    "tool_name": tool_name,
                    "tool_execution_id": tool_execution_id,
                    "status": status,
                    "tool_output": str(result)[:1000] if result else None,
                    "tool_error": str(error) if error else None,
                    "message": f"{'✅' if status == 'success' else '❌'} {status}: {tool_name}",
                }
            )

    return tool_hook
