app.py

import os
import re
import shutil
import traceback
import gradio as gr
from pathlib import Path
from histopath.agent import A1
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Get passcode from environment
PASSCODE = os.getenv("GRADIO_PASSWORD")

# Initialize agent (will be created after passcode validation)
agent = None


def check_for_output_files():
    """Check for all files in the output directory and return their paths."""
    output_dir = Path("./output")
    if not output_dir.exists():
        return [], []
    
    image_extensions = {".png", ".jpg", ".jpeg", ".svg", ".tif", ".tiff"}
    data_extensions = {".csv", ".txt", ".json", ".npy"}
    
    images = []
    data_files = []
    
    for file in output_dir.iterdir():
        if file.is_file():
            if file.suffix.lower() in image_extensions:
                images.append(str(file))
            elif file.suffix.lower() in data_extensions:
                data_files.append(str(file))
    
    return images, data_files


def preview_uploaded_file(uploaded_file):
    """Preview the uploaded file - show image or file info."""
    if uploaded_file is None:
        return None, None, "No file uploaded"
    
    file_path = Path(uploaded_file.name)
    file_ext = file_path.suffix.lower()
    
    image_extensions = {".png", ".jpg", ".jpeg", ".svg", ".tif", ".tiff", ".svs"}
    
    if file_ext in image_extensions:
        # Show image preview
        return uploaded_file.name, None, f"📷 Previewing: {file_path.name}"
    else:
        # Show file info
        file_size = Path(uploaded_file.name).stat().st_size / 1024  # KB
        return None, uploaded_file.name, f"📄 File: {file_path.name} ({file_size:.1f} KB)"


def parse_agent_output(output):
    """Parse agent output to extract code blocks, observations, and regular text."""
    # Strip out the message divider bars
    output = re.sub(r'={30,}\s*(Human|Ai)\s+Message\s*={30,}', '', output)
    output = output.strip()
    
    parsed = {
        "type": "text",
        "content": output,
        "code": None,
        "observation": None,
        "thinking": None
    }
    
    # Check for code execution block
    execute_match = re.search(r'<execute>(.*?)</execute>', output, re.DOTALL)
    if execute_match:
        parsed["type"] = "code"
        parsed["code"] = execute_match.group(1).strip()
        # Extract text before the code block (thinking/explanation)
        text_before = output[:execute_match.start()].strip()
        # Remove any think tags but keep the content
        text_before = re.sub(r'<think>(.*?)</think>', r'\1', text_before, flags=re.DOTALL)
        text_before = re.sub(r'={30,}.*?={30,}', '', text_before).strip()
        parsed["thinking"] = text_before if text_before else None
        return parsed
    
    # Check for observation block
    observation_match = re.search(r'<observation>(.*?)</observation>', output, re.DOTALL)
    if observation_match:
        parsed["type"] = "observation"
        parsed["observation"] = observation_match.group(1).strip()
        # Extract text before observation if any
        text_before = output[:observation_match.start()].strip()
        text_before = re.sub(r'<think>(.*?)</think>', r'\1', text_before, flags=re.DOTALL)
        text_before = re.sub(r'={30,}.*?={30,}', '', text_before).strip()
        parsed["thinking"] = text_before if text_before else None
        return parsed
    
    # Check for solution block
    solution_match = re.search(r'<solution>(.*?)</solution>', output, re.DOTALL)
    if solution_match:
        parsed["type"] = "solution"
        parsed["content"] = solution_match.group(1).strip()
        # Get thinking before solution
        text_before = output[:solution_match.start()].strip()
        text_before = re.sub(r'<think>(.*?)</think>', r'\1', text_before, flags=re.DOTALL)
        text_before = re.sub(r'={30,}.*?={30,}', '', text_before).strip()
        parsed["thinking"] = text_before if text_before else None
        return parsed
    
    # Clean up any remaining tags for display
    cleaned = re.sub(r'<think>(.*?)</think>', r'\1', output, flags=re.DOTALL)
    cleaned = re.sub(r'={30,}.*?={30,}', '', cleaned).strip()
    parsed["content"] = cleaned
    
    return parsed


def format_message_for_display(parsed_output):
    """Format parsed output into a readable message for the chatbot."""
    msg_parts = []
    
    # Add thinking/explanation text first if present
    if parsed_output.get("thinking"):
        msg_parts.append(parsed_output["thinking"])
    
    if parsed_output["type"] == "code":
        # Add separator if there was thinking text
        if parsed_output.get("thinking"):
            msg_parts.append("\n---\n")
        
        msg_parts.append("### 💻 Executing Code\n")
        msg_parts.append(f"```python\n{parsed_output['code']}\n```")
        
    elif parsed_output["type"] == "observation":
        # Add separator if there was thinking text
        if parsed_output.get("thinking"):
            msg_parts.append("\n---\n")
        
        msg_parts.append("### 📊 Observation\n")
        msg_parts.append(f"```\n{parsed_output['observation']}\n```")
        
    elif parsed_output["type"] == "solution":
        # Add separator if there was thinking text
        if parsed_output.get("thinking"):
            msg_parts.append("\n---\n")
        
        msg_parts.append("### ✅ Solution\n")
        msg_parts.append(parsed_output['content'])
        
    else:
        # For regular text, just add the content if thinking wasn't already set
        if not parsed_output.get("thinking"):
            msg_parts.append(parsed_output["content"])
    
    return "\n\n".join(msg_parts)


def process_agent_response(prompt, uploaded_file, chatbot_history):
    """Process the agent response and update chatbot - AGGRESSIVE FIX: Minimal yields."""
    global agent
    
    if agent is None:
        chatbot_history.append({
            "role": "assistant",
            "content": "⚠️ Please enter the passcode first to initialize the agent."
        })
        yield chatbot_history, None, None, None, None, "⚠️ Agent not initialized"
        return
    
    if not prompt.strip() and uploaded_file is None:
        chatbot_history.append({
            "role": "assistant",
            "content": "⚠️ Please provide a prompt or upload a file."
        })
        yield chatbot_history, None, None, None, None, "⚠️ No input provided"
        return
    
    # Handle file upload
    file_path = None
    file_info = ""
    if uploaded_file is not None:
        try:
            # Create data directory if it doesn't exist
            data_dir = Path("./data")
            data_dir.mkdir(exist_ok=True)
            
            # Copy uploaded file to data directory
            file_name = Path(uploaded_file.name).name
            file_path = data_dir / file_name
            shutil.copy(uploaded_file.name, file_path)
            
            file_info = f"\n\n📎 **Uploaded file:** `{file_path}`\n"
            
            # Augment prompt with file path
            if prompt.strip():
                prompt = f"{prompt}\n\nUploaded file path: {file_path}"
            else:
                prompt = f"I have uploaded a file at: {file_path}. Please analyze it."
                
        except Exception as e:
            error_msg = f"❌ Error handling file upload: {str(e)}"
            chatbot_history.append({
                "role": "assistant",
                "content": error_msg
            })
            yield chatbot_history, None, None, None, None, error_msg
            return
    
    # Add user message to chat
    user_message = prompt if not file_info else f"{prompt}{file_info}"
    chatbot_history.append({"role": "user", "content": user_message})
    
    # CRITICAL FIX: Only yield once at the start to show user message
    yield chatbot_history, None, None, None, None, "🔄 Processing..."
    
    try:
        # CRITICAL FIX: Collect ALL steps without yielding
        step_count = 0
        collected_outputs = []
        
        for step in agent.go_stream(prompt):
            step_count += 1
            output = step.get("output", "")
            
            if output:
                collected_outputs.append(output)
        
        # CRITICAL FIX: Process ALL collected outputs at once
        for output in collected_outputs:
            parsed = parse_agent_output(output)
            formatted_message = format_message_for_display(parsed)
            
            # Update or append to chatbot history
            if chatbot_history and chatbot_history[-1]["role"] == "assistant":
                # Update the last assistant message
                chatbot_history[-1]["content"] = formatted_message
            else:
                # Add new assistant message
                chatbot_history.append({
                    "role": "assistant",
                    "content": formatted_message
                })
        
        # CRITICAL FIX: Check files only ONCE after all processing
        images, data = check_for_output_files()
        
        status_msg = f"✅ Complete ({step_count} steps)"
        if images:
            status_msg += f" | {len(images)} image(s)"
        if data:
            status_msg += f" | {len(data)} data file(s)"
        
        # CRITICAL FIX: Final single yield with all results
        yield chatbot_history, images, data, None, None, status_msg
        
    except Exception as e:
        error_trace = traceback.format_exc()
        error_msg = f"❌ **Error:** {str(e)}\n\n<details>\n<summary>Stack Trace</summary>\n\n```\n{error_trace}\n```\n</details>"
        
        chatbot_history.append({
            "role": "assistant",
            "content": error_msg
        })
        
        yield chatbot_history, None, None, None, None, f"❌ Error: {str(e)}"


def clear_chat():
    """Clear the chat history and outputs."""
    return [], None, None, None, None, "Ready"


def validate_passcode(input_passcode):
    """Validate the passcode and initialize the agent."""
    global agent
    
    if input_passcode == PASSCODE:
        try:
            # Initialize the agent
            agent = A1(
                path="./data",
                llm="claude-sonnet-4-20250514",
                use_tool_retriever=True,
                timeout_seconds=600
            )
            return (
                gr.update(visible=False),  # Hide passcode section
                gr.update(visible=True),   # Show main interface
                "✅ Access granted! Agent initialized successfully."
            )
        except Exception as e:
            error_msg = f"❌ Failed to initialize agent: {str(e)}"
            return (
                gr.update(visible=True),
                gr.update(visible=False),
                error_msg
            )
    else:
        return (
            gr.update(visible=True),
            gr.update(visible=False),
            "❌ Invalid passcode. Please try again."
        )


# Custom theme
custom_theme = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="purple",
    neutral_hue="slate",
    font=["Inter", "system-ui", "sans-serif"],
    text_size="md",
).set(
    button_primary_background_fill="*primary_500",
    button_primary_background_fill_hover="*primary_600",
    block_label_text_weight="600",
    block_title_text_weight="600",
)

with gr.Blocks(title="HistoPath Agent", theme=custom_theme, css="""
    .gradio-container {
        max-width: 100% !important;
    }
    .main-header {
        text-align: center;
        padding: 1.5rem 0;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        border-radius: 8px;
        margin-bottom: 1.5rem;
    }
    .main-header h1 {
        margin: 0;
        font-size: 2.2rem;
        font-weight: 700;
    }
    .main-header p {
        margin: 0.5rem 0 0 0;
        opacity: 0.95;
        font-size: 1.1rem;
    }
    .file-upload-box .wrap {
        min-width: 0 !important;
    }
    .file-upload-box .file-name {
        word-break: break-word !important;
        white-space: normal !important;
        overflow-wrap: break-word !important;
    }
    .tab-nav {
        margin-bottom: 0.5rem;
    }
    /* Better styling for code and observation blocks */
    .message.bot pre {
        background-color: #f6f8fa !important;
        border: 1px solid #d0d7de !important;
        border-radius: 6px !important;
        padding: 12px !important;
        margin: 8px 0 !important;
    }
    .message.bot h3 {
        margin-top: 12px !important;
        margin-bottom: 8px !important;
        font-weight: 600 !important;
    }
    .message.bot hr {
        border: none !important;
        border-top: 2px solid #e1e4e8 !important;
        margin: 16px 0 !important;
    }
""") as demo:
    
    # Header
    gr.HTML("""
        <div class="main-header">
            <h1>🔬 HistoPath Agent</h1>
            <p>AI-Powered Histopathology Analysis Assistant</p>
        </div>
    """)
    
    # Passcode section
    with gr.Group(visible=True) as passcode_section:
        gr.Markdown("### 🔐 Authentication Required")
        
        with gr.Row():
            passcode_input = gr.Textbox(
                label="Passcode",
                type="password",
                placeholder="Enter your passcode...",
                scale=3
            )
            passcode_btn = gr.Button("🔓 Unlock", variant="primary", scale=1, size="lg")
        
        passcode_status = gr.Textbox(
            label="Status",
            interactive=False,
            lines=2
        )
    
    # Main interface (hidden initially)
    with gr.Group(visible=False) as main_interface:
        with gr.Row(equal_height=True):
            # Left column - Chat interface
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(
                    label="💬 Conversation",
                    height=550,
                    show_label=True,
                    render_markdown=True,
                )
                
                # Input area
                with gr.Row():
                    with gr.Column(scale=7):
                        prompt_input = gr.Textbox(
                            label="Your Query",
                            placeholder="E.g., 'Caption the uploaded whole slide image' or 'Segment cells using instanseg model'",
                            lines=2,
                            max_lines=5,
                            show_label=False,
                        )
                    with gr.Column(scale=3):
                        file_upload = gr.File(
                            label="📎 Upload File",
                            file_types=[".svs", ".png", ".jpg", ".jpeg", ".tif", ".tiff", ".csv", ".txt", ".json", ".npy"],
                            height=75,
                            elem_classes="file-upload-box",
                        )
                
                with gr.Row():
                    submit_btn = gr.Button("🚀 Submit", variant="primary", scale=3, size="lg")
                    clear_btn = gr.Button("🗑️ Clear", scale=1, size="lg", variant="secondary")
                
                status_text = gr.Textbox(
                    label="Status",
                    interactive=False,
                    value="Ready",
                    show_label=False,
                    container=False,
                )
            
            # Right column - Outputs
            with gr.Column(scale=2):
                with gr.Tabs():
                    with gr.Tab("📥 Input"):
                        with gr.Column():
                            input_image_preview = gr.Image(
                                label="Input Image",
                                height=400,
                                show_label=False,
                                container=True,
                            )
                            input_file_preview = gr.File(
                                label="Input File",
                                interactive=False,
                                height=100,
                                show_label=False,
                                container=True,
                            )
                            input_status = gr.Textbox(
                                value="Upload a file to preview",
                                show_label=False,
                                interactive=False,
                                container=False,
                            )
                    
                    with gr.Tab("🖼️ Images"):
                        output_gallery = gr.Gallery(
                            label="Generated Visualizations",
                            columns=1,
                            height=600,
                            object_fit="contain",
                            show_label=False,
                            show_download_button=True,
                        )
                    
                    with gr.Tab("📄 Data"):
                        data_files = gr.File(
                            label="Generated Data Files",
                            file_count="multiple",
                            interactive=False,
                            height=600,
                            show_label=False,
                        )
    
    # Event handlers
    passcode_btn.click(
        fn=validate_passcode,
        inputs=[passcode_input],
        outputs=[passcode_section, main_interface, passcode_status]
    )
    
    # File upload preview
    file_upload.change(
        fn=preview_uploaded_file,
        inputs=[file_upload],
        outputs=[input_image_preview, input_file_preview, input_status]
    )
    
    submit_btn.click(
        fn=process_agent_response,
        inputs=[prompt_input, file_upload, chatbot],
        outputs=[chatbot, output_gallery, data_files, input_image_preview, input_file_preview, status_text]
    )
    
    clear_btn.click(
        fn=clear_chat,
        outputs=[chatbot, output_gallery, data_files, input_image_preview, input_file_preview, status_text]
    )
    
    # Allow enter key to submit
    prompt_input.submit(
        fn=process_agent_response,
        inputs=[prompt_input, file_upload, chatbot],
        outputs=[chatbot, output_gallery, data_files, input_image_preview, input_file_preview, status_text]
    )


if __name__ == "__main__":
    # Create necessary directories
    Path("./data").mkdir(exist_ok=True)
    Path("./output").mkdir(exist_ok=True)
    
    print("=" * 60)
    print("🔬 HistoPath Agent - Gradio Interface")
    print("=" * 60)
    print("Starting server...")
    print("=" * 60)
    
    # Launch the app
    demo.launch()