ui/chat_interface.py

"""
Chat interface implementation using Gradio.
"""
import gradio as gr
import re
import asyncio
import logging
import json
from agents import Runner
from agent_system.agents import orchestrator_agent, process_agent_response
from agent_system.auth import AuthState
from agent_system.video_handler import handle_video_with_stages, handle_script_to_video
from agent_system.context import ConversationContext, Artifact, ArtifactType
from typing import List, Dict, Any, Optional
from datetime import datetime

logger = logging.getLogger(__name__)


class ChatState:
    """Maintains conversation state across messages."""
    def __init__(self):
        self.context = ConversationContext()
        logger.info("Created new ChatState with fresh ConversationContext")

# Custom chat function for Gradio interface
async def chat_with_agents(message: str, history: List[Dict[str, Any]], state: Optional[ChatState] = None):
    """
    Process user messages through the agent system.

    Args:
        message: The current user message
        history: List of previous messages in role/content format
        state: Conversation state (created per session)

    Yields:
        Updated message history for streaming updates
    """
    # Initialize state if not provided (first message)
    if state is None:
        state = ChatState()
        logger.info("Created new ChatState for conversation")
    
    # Create unique conversation ID
    conversation_id = f"conv_{len(history)}"

    # NEW: Use context-based auth state instead of history scanning
    is_authenticated = state.context.is_authenticated()
    user_email = state.context.get_authenticated_user()
    
    logger.info(f"User authenticated: {is_authenticated}, email: {user_email}")

    # Add initial greeting for new chats (though we now set it in the UI)
    # This is kept for backward compatibility
    if len(history) == 0:
        initial_history = []
        initial_history.append({
            "role": "assistant",
            "content": "Welcome to AI Agent Chat! Please enter your email address to continue."
        })
        yield initial_history
        return

    # Authentication is now handled by the orchestrator and auth_agent
    # No need for special email handling here

    # Build input list from history
    inputs = []
    for msg in history:
        inputs.append({
            "content": msg["content"],
            "role": msg["role"]
        })

    # Add current message
    inputs.append({"content": message, "role": "user"})

    # Process ALL requests through orchestrator (it will handle authentication)
    try:
        # Prepare orchestrator input with context
        # Format the message with context information
        context_summary = state.context.get_context_summary()
        
        # Prepare full conversation history for orchestrator
        orchestrator_input = []
        
        # Add conversation history
        for msg in history:
            orchestrator_input.append({
                "role": msg["role"],
                "content": msg["content"]
            })
        
        # Add current user message with context
        orchestrator_message = f"""User message: {message}

Current context:
- Available artifacts: {len(context_summary['available_artifacts'])}
- Last artifacts by type: {context_summary.get('last_artifacts', {})}

Available artifacts details:
{json.dumps(context_summary['available_artifacts'], indent=2) if context_summary['available_artifacts'] else 'None'}

Please process this request considering the available context and conversation history."""
        
        orchestrator_input.append({
            "role": "user",
            "content": orchestrator_message
        })
        
        logger.info(f"Sending to orchestrator - authenticated: {is_authenticated}, artifacts available: {len(state.context.artifacts)}")
        
        # DETAILED LOGGING: Check message format before sending
        logger.info(f"=== ORCHESTRATOR INPUT VALIDATION ===")
        logger.info(f"Input type: {type(orchestrator_input)}")
        logger.info(f"Input length: {len(orchestrator_input)}")
        
        for i, msg in enumerate(orchestrator_input):
            logger.info(f"Message {i}: type={type(msg)}, keys={list(msg.keys()) if isinstance(msg, dict) else 'NOT_DICT'}")
            if isinstance(msg, dict):
                logger.info(f"  role: {repr(msg.get('role'))} (type: {type(msg.get('role'))})")
                logger.info(f"  content: {repr(msg.get('content')[:100])}... (type: {type(msg.get('content'))})")
            else:
                logger.error(f"  INVALID MESSAGE FORMAT: {repr(msg)}")
        
        logger.info(f"=== CALLING ORCHESTRATOR ===")
        
        try:
            # Run orchestrator with context (it will check auth internally)
            result = Runner.run_streamed(
                orchestrator_agent,
                input=orchestrator_input,
                context=state.context
            )
            logger.info("✅ Orchestrator call succeeded")
        except Exception as e:
            logger.error(f"❌ ORCHESTRATOR CALL FAILED: {e}")
            logger.error(f"Exception type: {type(e)}")
            import traceback
            logger.error(f"Traceback: {traceback.format_exc()}")
            raise
        
        # Create a copy of history and add the user message
        new_history = history.copy()
        new_history.append({"role": "user", "content": message})
        
        # Process orchestrator response and create artifacts
        logger.info(f"=== STARTING ORCHESTRATOR RESPONSE PROCESSING ===")
        async for event in process_orchestrator_response_with_artifacts(result, state, new_history):
 
            yield event
            
    except Exception as e:
        logger.error(f"Orchestration error: {e}", exc_info=True)
        error_response = f"I encountered an error: {str(e)}"
        new_history = history.copy()
        new_history.append({"role": "user", "content": message})
        new_history.append({
            "role": "assistant", 
            "content": error_response
        })
        # Clean the history format for Gradio messages type
        clean_history = []
        for msg in new_history:
            clean_history.append({
                "role": msg["role"],
                "content": msg["content"]
            })
        yield clean_history


async def process_orchestrator_response(result, state: ChatState, history: List[Dict[str, Any]]):
    """Process streaming response from orchestrator and track artifacts."""
    response_text = ""
    current_agent = "orchestrator"
    
    from openai.types.responses import ResponseTextDeltaEvent, ResponseContentPartDoneEvent
    from agents import RawResponsesStreamEvent, HandoffCallItem
    
    async for event in result.stream_events():
        # Get the current agent
        current_agent = result.current_agent.name if hasattr(result, 'current_agent') else 'orchestrator'
        
        # Only process response content events
        if not isinstance(event, RawResponsesStreamEvent):
            continue
            
        data = event.data
        if isinstance(data, ResponseTextDeltaEvent):
            response_text += data.delta
            
            # Stream the response
            temp_history = history.copy()
            temp_history.append({
                "role": "assistant", 
                "content": response_text
            })
            yield temp_history
            
        elif isinstance(data, ResponseContentPartDoneEvent):
            # Final event - check if we need to extract artifacts
            # The orchestrator might have created artifacts through tool calls
            logger.info(f"Response complete from {current_agent}")
    
    # Final history with complete response
    final_history = history.copy()
    final_history.append({
        "role": "assistant",
        "content": response_text
    })
    yield final_history


async def process_orchestrator_response_with_artifacts(result, state: ChatState, new_history):
    """Process orchestrator response and create artifacts from tool results."""
    full_response = ""
    last_yield_length = 0
    yield_threshold = 30  # Yield every 30 characters for smooth streaming
    
    try:
        logger.info(f"=== PROCESSING ORCHESTRATOR STREAM EVENTS ===")
        async for event in result.stream_events():
            if hasattr(event, 'data'):
                event_data = event.data
                
                # Only collect text deltas, not function call argument deltas
                if hasattr(event_data, 'delta') and event_data.delta:
                    # Check if this is a text delta (not function call arguments)
                    event_type_name = type(event_data).__name__
                    if 'TextDelta' in event_type_name:
                        full_response += event_data.delta
                        
                        # Throttle yields for better UX - only yield every N characters or at word boundaries
                        chars_since_last_yield = len(full_response) - last_yield_length
                        should_yield = (
                            chars_since_last_yield >= yield_threshold or  # Every N characters
                            event_data.delta.endswith(' ') or           # At word boundaries
                            event_data.delta.endswith('\n') or          # At line breaks
                            event_data.delta.endswith('.') or           # At sentence ends
                            event_data.delta.endswith('!')              # At exclamations
                        )
                        
                        if should_yield:
                            # Stream the response
                            new_history_copy = new_history.copy()
                            new_history_copy.append({"role": "assistant", "content": full_response})
                            
                            # Clean the history format for Gradio messages type
                            clean_history = []
                            for msg in new_history_copy:
                                clean_history.append({
                                    "role": msg["role"],
                                    "content": msg["content"]
                                })
                            
                            yield clean_history, state
                            last_yield_length = len(full_response)
                            
                            # Add small delay for smooth streaming per Gradio best practices
                            await asyncio.sleep(0.05)
                    
                # Handle tool results for artifact creation
                elif hasattr(event.data, 'tool_calls'):
                    for tool_call in event.data.tool_calls:
                        if hasattr(tool_call, 'name') and hasattr(tool_call, 'output'):
                            create_artifact_from_tool_result(tool_call, state)
        
        # Final yield to ensure complete response is shown even if last chunk was < threshold
        if full_response and last_yield_length < len(full_response):
            new_history_copy = new_history.copy()
            new_history_copy.append({"role": "assistant", "content": full_response})
            
            # Clean the history format for Gradio messages type
            clean_history = []
            for msg in new_history_copy:
                clean_history.append({
                    "role": msg["role"],
                    "content": msg["content"]
                })
            
            logger.info(f"Final yield: complete response ({len(full_response)} chars)")
            yield clean_history, state
                            
    except Exception as e:
        logger.error(f"Response processing error: {e}")
        logger.error(f"Exception type: {type(e)}")
        import traceback
        logger.error(f"Full traceback: {traceback.format_exc()}")
        
        error_msg = full_response + f"\n\n[Error: {str(e)}]"
        new_history_copy = new_history.copy()
        new_history_copy.append({"role": "assistant", "content": error_msg})
        
        logger.info(f"=== YIELDING ERROR RESPONSE ===")
        logger.info(f"Error history length: {len(new_history_copy)}")
        logger.info(f"Error message format: {repr(new_history_copy[-1])}")
        
        # Clean the history format for Gradio messages type  
        clean_history = []
        for msg in new_history_copy:
            clean_history.append({
                "role": msg["role"],
                "content": msg["content"]
            })
        yield clean_history, state


def create_artifact_from_tool_result(tool_call, state: ChatState):
    """Create artifacts from successful tool calls."""
    try:
        tool_name = tool_call.name
        tool_output = tool_call.output
        
        if tool_name == "search_web_tool" and tool_output and "Search failed" not in tool_output:
            # Create search artifact
            artifact = Artifact(
                type=ArtifactType.SEARCH_RESULTS,
                content={
                    "results": [tool_output],
                    "full_response": tool_output,
                    "timestamp": datetime.now().isoformat()
                },
                created_by="search_web_tool",
                metadata={"summary": f"Web search results - {len(tool_output)} chars"}
            )
            state.context.add_artifact(artifact)
            # logger.info(f"Created search artifact: {artifact.id}")
            
        elif tool_name == "create_or_modify_script_tool" and tool_output and "Failed to create" not in tool_output:
            # Extract title from script output
            lines = tool_output.split('\n')
            title = "Generated Script"
            script_content = tool_output
            
            for line in lines:
                if line.startswith("Script created:"):
                    title = line.replace("Script created:", "").strip()
                    break
            
            # Create script artifact
            artifact = Artifact(
                type=ArtifactType.SCRIPT,
                content={
                    "title": title,
                    "script": script_content,
                    "created_at": datetime.now().isoformat()
                },
                created_by="create_or_modify_script_tool",
                metadata={"summary": f"Script: {title}"}
            )
            state.context.add_artifact(artifact)
            # logger.info(f"Created script artifact: {artifact.id}")
            
        elif tool_name == "create_video_tool" and tool_output and "Failed to create" not in tool_output:
            # Extract video URL from output
            video_url = ""
            title = "Generated Video"
            
            for line in tool_output.split('\n'):
                if "Video URL:" in line:
                    video_url = line.split("Video URL:", 1)[1].strip()
                elif line.startswith("Title:"):
                    title = line.replace("Title:", "").strip()
            
            # Create video artifact
            artifact = Artifact(
                type=ArtifactType.VIDEO,
                content={
                    "title": title,
                    "video_url": video_url,
                    "created_at": datetime.now().isoformat()
                },
                created_by="create_video_tool", 
                metadata={"summary": f"Video: {title} - {video_url}"}
            )
            state.context.add_artifact(artifact)
            # logger.info(f"Created video artifact: {artifact.id}")
            
    except Exception as e:
        logger.error(f"Artifact creation error: {e}")


# Create the interface
def create_chat_interface():
    """Create and configure the Gradio chat interface."""
    
    # Note: Using file= syntax for cross-platform compatibility
    # CSS for layout and button styling only
    custom_css = """
    .agent-label {
        font-size: 0.8em;
        padding: 2px 8px;
        border-radius: 4px;
        margin-bottom: 5px;
        display: inline-block;
        background: #E6F3FF;
        color: #1E90FF;
    }

    .heygen-video-embed {
        margin: 10px 0;
        width: 100%;
        max-width: 640px;
    }

    .video-link {
        margin-top: 5px;
        font-size: 0.9em;
    }
    
    /* Desktop: Make send button taller to match textbox */
    .send-btn {
        min-height: 60px !important;
        height: 60px !important;
    }
    
    /* Mobile: Stack button below textbox */
    @media (max-width: 768px) {
        .input-row {
            flex-direction: column !important;
            gap: 8px !important;
        }
        
        .input-row > * {
            width: 100% !important;
            flex: none !important;
        }
        
        .send-btn {
            min-height: 44px !important;
            height: 44px !important;
            width: 100% !important;
        }
    }
    """

    with gr.Blocks(css=custom_css, theme="soft") as demo:
        # Add state component
        chat_state = gr.State(ChatState())
        
        # Header with logo and text properly centered
        with gr.Row(elem_classes="header-row"):
            with gr.Column(scale=1, min_width=120):
                        # Text with proper vertical centering
                gr.HTML("<h2 style='margin: 0; padding: 15px 0; color: #1F2937; font-size: 1.5rem; font-weight: 600; display: flex; align-items: center; height: 45px;'>Chat to create videos</h2>")
            
        
        # Commented out login text
        # gr.HTML("<p style='text-align: center'>Please login with your email to enjoy jokes and poems!</p>")

        # Initialize with welcome message
        initial_message = [
            {"role": "assistant", "content": "Welcome! Please enter your email address to continue."}
        ]

        chatbot = gr.Chatbot(
            height=500,
            type="messages",
            show_copy_button=True,
            bubble_full_width=False,
            render_markdown=True,  # Enable markdown rendering
            sanitize_html=False,   # Allow HTML to be rendered
            value=initial_message,  # Set initial welcome message
            show_label=False
        )
        
        # Create input row with classes for styling
        with gr.Row(elem_classes="input-row"):
            msg = gr.Textbox(
                placeholder="Type your message here...",
                scale=9,
                show_label=False,
                lines=1,
                max_lines=4
            )
            submit = gr.Button("Send", variant="primary", scale=1, elem_classes="send-btn")
        
        clear = gr.Button("Clear Chat")
        
        # # Add debug info section at bottom (collapsed by default)
        # with gr.Accordion("Debug Info", open=False):
        #     context_display = gr.JSON(
        #         value={},
        #         label="Active Artifacts",
        #         elem_id="context-display"
        #     )
        
        # Update handlers to pass state
        async def respond(message, chat_history, state):
            """Handle message and update context display."""
            if not message:
                yield chat_history, "", state
                return
            
            logger.info(f"User message: {message}")
            
            # Process message
            try:
                async for updated_history in chat_with_agents(message, chat_history, state):
                    
                    # Ensure clean message format for Gradio
                    if isinstance(updated_history, tuple):
                        clean_history, _ = updated_history
                    else:
                        clean_history = updated_history
                    
                    # Double-check that messages only have role and content keys
                    final_clean_history = []
                    for msg in clean_history:
                        final_clean_history.append({
                            "role": msg["role"],
                            "content": msg["content"]
                        })
                    
                    yield final_clean_history, "", state
                    
            except Exception as e:
                logger.error(f"=== FRONTEND ERROR ===")
                logger.error(f"Frontend respond error: {e}")
                logger.error(f"Exception type: {type(e)}")
                import traceback
                logger.error(f"Full traceback: {traceback.format_exc()}")
                
                # Return error state with clean format
                error_history = chat_history.copy()
                error_history.append({"role": "assistant", "content": f"Frontend error: {str(e)}"})
                
                # Clean the error history format
                clean_error_history = []
                for msg in error_history:
                    clean_error_history.append({
                        "role": msg["role"],
                        "content": msg["content"]
                    })
                yield clean_error_history, "", state
        
        # Set up event handlers with state
        msg.submit(
            respond,
            [msg, chatbot, chat_state],
            [chatbot, msg, chat_state],
            queue=True
        )
        
        submit.click(
            respond,
            [msg, chatbot, chat_state],
            [chatbot, msg, chat_state],
            queue=True
        )
        
        # Clear preserves state but clears history
        def clear_chat(state):
            # Keep the state but clear chat history
            return [], state
        
        clear.click(clear_chat, [chat_state], [chatbot, chat_state], queue=False)
        
    
    return demo