src/streamlit_app.py

import streamlit as st
import json
from datetime import datetime
import uuid
import os
import time
import base64
from PIL import Image
import io

from main import call_agent as agent_invoker
from main import DemoState

def call_agent(query):
    """
    Mock function - replace this with your actual call_agent function
    """
    # This is a mock response for demonstration
    mock_response = [
        {"id": 1, "task": "Fetch DEM data for the specified region", "tool": "DEMFetcher"},
        {"id": 2, "task": "Extract drainage networks from available data", "tool": "DrainageExtractor"},
        {"id": 3, "task": "Analyze hydrological flow patterns", "tool": "HydrologyAnalyzer"},
        {"id": 4, "task": "Generate flood risk assessment maps", "tool": "LLM Reasoning"}
    ]
    state =  agent_invoker(query)
    # Mock state object
    class MockState:
        def __init__(self):
            self.query = state.get('query', query)
            self.response = state.get('response', mock_response)
            self.output_files_path = state.get('output_files_path', [f"outputs/geospatial_plan_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:6]}.json"])
    return MockState()






def get_image_base64(image_path):
    """Convert local image to base64 string for embedding in HTML"""
    try:
        if os.path.exists(image_path):
            with open(image_path, "rb") as img_file:
                return base64.b64encode(img_file.read()).decode()
        else:
            # Return a placeholder SVG if image doesn't exist
            placeholder_svg = f"""
            <svg width="800" height="600" xmlns="http://www.w3.org/2000/svg">
                <rect width="100%" height="100%" fill="#f0f0f0"/>
                <text x="50%" y="50%" text-anchor="middle" dy=".3em" font-family="Arial, sans-serif" font-size="24" fill="#666">
                    Image not found: {os.path.basename(image_path)}
                </text>
            </svg>
            """
            return base64.b64encode(placeholder_svg.encode()).decode()
    except Exception as e:
        # Return error placeholder
        error_svg = f"""
        <svg width="800" height="600" xmlns="http://www.w3.org/2000/svg">
            <rect width="100%" height="100%" fill="#ffebee"/>
            <text x="50%" y="50%" text-anchor="middle" dy=".3em" font-family="Arial, sans-serif" font-size="20" fill="#c62828">
                Error loading image: {str(e)}
            </text>
        </svg>
        """
        return base64.b64encode(error_svg.encode()).decode()

def create_placeholder_image(text, width=800, height=600, bg_color="#2196F3", text_color="#FFFFFF"):
    """Create a placeholder image with text"""
    try:
        img = Image.new('RGB', (width, height), bg_color)
        # Note: This creates a simple colored rectangle. For text rendering, you'd need PIL's ImageDraw
        # For now, we'll return a base64 encoded simple image
        buffer = io.BytesIO()
        img.save(buffer, format='PNG')
        return base64.b64encode(buffer.getvalue()).decode()
    except:
        # Fallback to SVG
        svg = f"""
        <svg width="{width}" height="{height}" xmlns="http://www.w3.org/2000/svg">
            <rect width="100%" height="100%" fill="{bg_color}"/>
            <text x="50%" y="50%" text-anchor="middle" dy=".3em" font-family="Arial, sans-serif" font-size="24" fill="{text_color}">
                {text}
            </text>
        </svg>
        """
        return base64.b64encode(svg.encode()).decode()

# Sample analysis data with local image paths
# Update these paths to point to your actual local images
SAMPLE_ANALYSES = [
    {
        "image_path": "output/Chennai_dem_filled.png",  # Update this path
        "caption": "Digital elevation model for CHENNAI,tamil nadu , India",
        "query": "Create a flood risk assessment map for Chennai during monsoon season"
    },
    {
        "image_path": "output/chennai_impervious.png",  # Update this path
        "caption": "Impervious Structure in chennai",
        "query": "Create a flood risk assessment map for Chennai during monsoon season"
    },
    {
        "image_path": "output/chennai_flow_direction.png",  # Update this path
        "caption": "Chennai Water flow direction",
        "query": "Create a flood risk assessment map for Chennai during monsoon season"
    },
    {
        "image_path": "output/chennai_streams.png",  # Update this path
        "caption": "Chennai Stream Flow",
       "query": "Create a flood risk assessment map for Chennai during monsoon season"
    },
    {
        "image_path": "output/chennai_raods.png",  # Update this path
        "caption": "Roads Pattern of Chennai",
       "query": "Create a flood risk assessment map for Chennai during monsoon season"
    }
]

# Page configuration
st.set_page_config(
    page_title="Geospatial AI Task Planner",
    page_icon="🌍",
    layout="wide",
    initial_sidebar_state="collapsed"
)

# Custom CSS for modern design (same as before)
st.markdown("""
<style>
    /* Hide default streamlit elements */
    .stAppHeader {display: none;}
    .stDeployButton {display: none;}
    #MainMenu {visibility: hidden;}
    .stAppDeployButton {display: none;}
    
    /* Main container styling */
    .main {
        padding: 0 !important;
    }
    
    /* Header styling */
    .header-container {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        padding: 20px;
        color: white;
        box-shadow: 0 4px 20px rgba(0,0,0,0.1);
        margin-bottom: 20px;
        border-radius: 0 0 15px 15px;
    }
    
    .header-content {
        display: flex;
        justify-content: space-between;
        align-items: center;
        max-width: 1200px;
        margin: 0 auto;
    }
    
    .logo {
        width: 80px;
        height: 80px;
        border-radius: 50%;
        object-fit: cover;
        border: 3px solid white;
        box-shadow: 0 4px 15px rgba(0,0,0,0.2);
    }
    
    .header-title {
        text-align: center;
        margin: 0;
        font-size: 2.5rem;
        font-weight: 700;
        text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
        background: linear-gradient(45deg, #fff, #f0f0f0);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        background-clip: text;
    }
    
    .header-subtitle {
        text-align: center;
        margin: 5px 0 0 0;
        font-size: 1.2rem;
        opacity: 0.9;
        font-weight: 300;
    }
    
    /* Column styling */
    .image-column {
        background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
        padding: 20px;
        border-radius: 15px;
        margin-right: 10px;
        min-height: 70vh;
        display: flex;
        flex-direction: column;
        align-items: center;
        justify-content: flex-start;
    }
    
    .chat-column {
        background: white;
        padding: 30px;
        border-radius: 15px;
        margin-left: 10px;
        min-height: 70vh;
        box-shadow: 0 4px 20px rgba(0,0,0,0.1);
        border: 1px solid #e9ecef;
    }
    
    .sample-analysis-header {
        text-align: center;
        margin-bottom: 15px;
        width: 100%;
    }
    
    .sample-analysis-title {
        font-size: 1.8rem;
        font-weight: 700;
        color: #2c3e50;
        margin: 0;
        text-shadow: 1px 1px 2px rgba(0,0,0,0.1);
    }
    
    .sample-analysis-subtitle {
        font-size: 1.1rem;
        color: #6c757d;
        margin: 5px 0 0 0;
        font-weight: 400;
        font-style: italic;
    }
    
    .query-display {
        background: linear-gradient(135deg, #e3f2fd 0%, #bbdefb 100%);
        padding: 15px;
        border-radius: 10px;
        margin-bottom: 20px;
        border-left: 4px solid #2196f3;
        width: 100%;
        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
    }
    
    .query-label {
        font-weight: 600;
        color: #1976d2;
        margin-bottom: 5px;
        font-size: 0.9rem;
    }
    
    .query-text {
        color: #333;
        font-size: 0.95rem;
        line-height: 1.4;
        margin: 0;
    }
    
    .slideshow-container {
        position: relative;
        width: 100%;
        max-width: 100%;
        margin: 0 auto;
        flex-grow: 1;
        display: flex;
        align-items: center;
        justify-content: center;
    }
    
    .slide {
        display: none;
        width: 100%;
        text-align: center;
        animation: fadeIn 0.5s ease-in-out;
    }
    
    .slide.active {
        display: block;
    }
    
    .slide-image {
        width: 100%;
        max-width: 600px;
        height: 400px;
        object-fit: cover;
        border-radius: 15px;
        box-shadow: 0 10px 30px rgba(0,0,0,0.2);
        transition: transform 0.3s ease;
    }
    
    .slide-image:hover {
        transform: scale(1.02);
    }
    
    .slide-caption {
        background: rgba(0,0,0,0.8);
        color: white;
        padding: 15px;
        border-radius: 0 0 15px 15px;
        margin-top: -4px;
        font-size: 1.1rem;
        font-weight: 500;
        text-shadow: 1px 1px 2px rgba(0,0,0,0.5);
    }
    
    .slide-indicators {
        text-align: center;
        margin-top: 20px;
    }
    
    .indicator {
        display: inline-block;
        width: 12px;
        height: 12px;
        border-radius: 50%;
        background: rgba(0,0,0,0.3);
        margin: 0 5px;
        cursor: pointer;
        transition: background 0.3s ease;
    }
    
    .indicator.active {
        background: #667eea;
    }
    
    .slide-navigation {
        position: absolute;
        top: 50%;
        transform: translateY(-50%);
        background: rgba(0,0,0,0.5);
        color: white;
        border: none;
        padding: 10px 15px;
        border-radius: 50%;
        cursor: pointer;
        font-size: 18px;
        transition: background 0.3s ease;
    }
    
    .slide-navigation:hover {
        background: rgba(0,0,0,0.8);
    }
    
    .slide-navigation.prev {
        left: 10px;
    }
    
    .slide-navigation.next {
        right: 10px;
    }
    
    @keyframes fadeIn {
        from { opacity: 0; }
        to { opacity: 1; }
    }
    
    .chat-header {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        padding: 20px;
        border-radius: 15px;
        margin-bottom: 20px;
        text-align: center;
        font-size: 1.3rem;
        font-weight: 600;
        box-shadow: 0 4px 15px rgba(0,0,0,0.1);
    }
    
    .chat-messages {
        max-height: 400px;
        overflow-y: auto;
        padding: 20px;
        background: #f8f9fa;
        border-radius: 15px;
        margin-bottom: 20px;
        border: 1px solid #e9ecef;
    }
    
    .user-message {
        background: linear-gradient(135deg, #4facfe 0%, #00f2fe 100%);
        color: white;
        padding: 15px 20px;
        border-radius: 25px 25px 5px 25px;
        margin: 15px 0;
        margin-left: 15%;
        text-align: right;
        box-shadow: 0 4px 15px rgba(79, 172, 254, 0.3);
        animation: slideInRight 0.3s ease;
    }
    
    .ai-message {
        background: white;
        color: #333;
        padding: 20px;
        border-radius: 25px 25px 25px 5px;
        margin: 15px 0;
        margin-right: 15%;
        border: 1px solid #e9ecef;
        box-shadow: 0 4px 15px rgba(0,0,0,0.05);
        animation: slideInLeft 0.3s ease;
    }
    
    .task-plan-container {
        border: 1px solid #e9ecef;
        border-radius: 15px;
        padding: 15px;
        margin-top: 15px;
        background-color: #f0f4f8;
    }

    .task-item {
        background: linear-gradient(135deg, #e6f7ff 0%, #cceeff 100%);
        padding: 15px;
        border-radius: 12px;
        margin: 10px 0;
        border-left: 4px solid #28a745;
        transition: transform 0.2s ease;
        box-shadow: 0 2px 5px rgba(0,0,0,0.08);
    }
    
    .task-item:hover {
        transform: translateX(5px);
    }
    
    .tool-badge {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        padding: 5px 12px;
        border-radius: 20px;
        font-size: 12px;
        font-weight: 600;
        display: inline-block;
        margin-top: 8px;
        box-shadow: 0 2px 8px rgba(0,0,0,0.1);
    }
    
    .input-section {
        background: white;
        padding: 20px;
        border-radius: 15px;
        border: 1px solid #e9ecef;
        box-shadow: 0 4px 15px rgba(0,0,0,0.05);
        margin-top: 20px;
    }
    
    .json-container {
        background: #f8f9fa;
        border: 1px solid #e9ecef;
        border-radius: 10px;
        padding: 15px;
        margin: 10px 0;
        font-family: 'Courier New', monospace;
        font-size: 12px;
        max-height: 300px;
        overflow-y: auto;
    }
    
    .download-section {
        background: #e8f5e8;
        border: 1px solid #28a745;
        border-radius: 10px;
        padding: 15px;
        margin: 10px 0;
        text-align: center;
    }
    
    .stTextArea textarea {
        border-radius: 12px !important;
        border: 2px solid #e9ecef !important;
        transition: border-color 0.3s ease !important;
    }
    
    .stTextArea textarea:focus {
        border-color: #667eea !important;
        box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
    }
    
    .stButton button {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
        color: white !important;
        border: none !important;
        border-radius: 25px !important;
        padding: 12px 30px !important;
        font-weight: 600 !important;
        font-size: 16px !important;
        transition: all 0.3s ease !important;
        box-shadow: 0 4px 15px rgba(102, 126, 234, 0.3) !important;
    }
    
    .stButton button:hover {
        transform: translateY(-2px) !important;
        box-shadow: 0 6px 20px rgba(102, 126, 234, 0.4) !important;
    }
    
    .stDownloadButton button {
        background: linear-gradient(135deg, #28a745 0%, #20c997 100%) !important;
        color: white !important;
        border: none !important;
        border-radius: 20px !important;
        padding: 10px 25px !important;
        font-weight: 600 !important;
        font-size: 14px !important;
        transition: all 0.3s ease !important;
        box-shadow: 0 4px 15px rgba(40, 167, 69, 0.3) !important;
    }
    
    .stDownloadButton button:hover {
        transform: translateY(-2px) !important;
        box-shadow: 0 6px 20px rgba(40, 167, 69, 0.4) !important;
    }
    
    @keyframes slideInRight {
        from { transform: translateX(100%); opacity: 0; }
        to { transform: translateX(0); opacity: 1; }
    }
    
    @keyframes slideInLeft {
        from { transform: translateX(-100%); opacity: 0; }
        to { transform: translateX(0); opacity: 1; }
    }
    
    /* Image loading message */
    .image-loading {
        background: #f8f9fa;
        border: 2px dashed #dee2e6;
        border-radius: 15px;
        padding: 40px;
        text-align: center;
        color: #6c757d;
        font-style: italic;
    }
    
    /* Scrollbar styling */
    .chat-messages::-webkit-scrollbar,
    .json-container::-webkit-scrollbar {
        width: 8px;
    }
    
    .chat-messages::-webkit-scrollbar-track,
    .json-container::-webkit-scrollbar-track {
        background: #f1f1f1;
        border-radius: 10px;
    }
    
    .chat-messages::-webkit-scrollbar-thumb,
    .json-container::-webkit-scrollbar-thumb {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        border-radius: 10px;
    }
    
    .chat-messages::-webkit-scrollbar-thumb:hover,
    .json-container::-webkit-scrollbar-thumb:hover {
        background: linear-gradient(135deg, #5a6fd8 0%, #6a4190 100%);
    }
    
    /* Responsive adjustments */
    @media (max-width: 768px) {
        .header-title {
            font-size: 1.8rem;
        }
        
        .header-subtitle {
            font-size: 1rem;
        }
        
        .logo {
            width: 60px;
            height: 60px;
        }
        
        .image-column, .chat-column {
            margin: 0;
            padding: 20px;
        }
        
        .slide-image {
            height: 300px;
        }
    }
</style>
""", unsafe_allow_html=True)

# Initialize session state
if 'messages' not in st.session_state:
    st.session_state.messages = []

if 'generated_json' not in st.session_state:
    st.session_state.generated_json = None

if 'processing' not in st.session_state:
    st.session_state.processing = False

if 'current_slide' not in st.session_state:
    st.session_state.current_slide = 0

if 'last_slide_change' not in st.session_state:
    st.session_state.last_slide_change = time.time()

# Header with logos and title
st.markdown("""
<div class="header-container">
    <div class="header-content">
        <div>
            <img src="https://ih1.redbubble.net/image.2734932759.2157/st,small,507x507-pad,600x600,f8f8f8.jpg" 
                 alt="ISRO Logo" class="logo" onerror="this.src='data:image/svg+xml;base64,PHN2ZyB3aWR0aD0iODAiIGhlaWdodD0iODAiIHZpZXdCb3g9IjAgMCA4MCA4MCIgZmlsbD0ibm9uZSIgeG1sbnM9Imh0dHA6Ly93d3cudzMub3JnLzIwMDAvc3ZnIj4KPGNpcmNsZSBjeD0iNDAiIGN5PSI0MCIgcj0iNDAiIGZpbGw9IiMwMDY2Q0MiLz4KPHRleHQgeD0iNDAiIHk9IjQ1IiB0ZXh0LWFuY2hvcj0ibWlkZGxlIiBmaWxsPSJ3aGl0ZSIgZm9udC1zaXplPSIxNCIgZm9udC13ZWlnaHQ9ImJvbGQiPklTUk88L3RleHQ+Cjwvc3ZnPgo='">
        </div>
        <div>
            <h1 class="header-title">Bharatvarsha Hackathon 2025-26</h1>
            <p class="header-subtitle">Geospatial AI Task Planner</p>
        </div>
        <div>
            <img src="https://via.placeholder.com/80x80/FF6B6B/FFFFFF?text=TECH" 
                 alt="Tech Logo" class="logo" onerror="this.src='data:image/svg+xml;base64,PHN2ZyB3aWR0aD0iODAiIGhlaWdodD0iODAiIHZpZXdCb3g9IjAgMCA4MCA4MCIgZmlsbD0ibm9uZSIgeG1sbnM9Imh0dHA6Ly93d3cudzMub3JnLzIwMDAvc3ZnIj4KPGNpcmNsZSBjeD0iNDAiIGN5PSI0MCIgcj0iNDAiIGZpbGw9IiNGRjZCNkIiLz4KPHRleHQgeD0iNDAiIHk9IjQ1IiB0ZXh0LWFuY2hvcj0ibWlkZGxlIiBmaWxsPSJ3aGl0ZSIgZm9udC1zaXplPSIxMiIgZm9udC13ZWlnaHQ9ImJvbGQiPlRFQ0g8L2h0ZXh0Pgo8L3N2Zz4K'>
        </div>
    </div>
</div>
""", unsafe_allow_html=True)

# Auto-advance slideshow
current_time = time.time()
if current_time - st.session_state.last_slide_change > 5:  # 5 seconds
    st.session_state.current_slide = (st.session_state.current_slide + 1) % len(SAMPLE_ANALYSES)
    st.session_state.last_slide_change = current_time

# Create two columns using st.columns
col1, col2 = st.columns([1, 1], gap="medium")

# Left column - Sample Analysis Slideshow
with col1:
    st.markdown("""
            <h2 class="sample-analysis-title">Sample Analysis</h2>
    """, unsafe_allow_html=True)
    
    # Display current query
    current_analysis = SAMPLE_ANALYSES[st.session_state.current_slide]
    st.markdown(f"""
        <div class="query-display">
            <div class="query-label">Query:</div>
            <p class="query-text">{current_analysis['query']}</p>
        </div>
    """, unsafe_allow_html=True)
    
    # Display the current slide image using Streamlit's native image display
    current_image_path = current_analysis['image_path']
    
    try:
        if os.path.exists(current_image_path):
            # Use Streamlit's native image display for better performance
            st.image(
                current_image_path,
                caption=current_analysis['caption'],
                
            )
        else:
            # Show a placeholder message when image is not found
            st.markdown(f"""
                <div class="image-loading">
                    <h3>📁 Image not found</h3>
                    <p>Please place your image at: <code>{current_image_path}</code></p>
                    <p>Or update the image path in the SAMPLE_ANALYSES list</p>
                </div>
            """, unsafe_allow_html=True)
    except Exception as e:
        st.markdown(f"""
            <div class="image-loading">
                <h3>⚠️ Error loading image</h3>
                <p>Path: <code>{current_image_path}</code></p>
                <p>Error: {str(e)}</p>
            </div>
        """, unsafe_allow_html=True)
    
    # Navigation buttons
    _,col_prev, col_next = st.columns([0.5,0.8, 1])
    with col_prev:
        if st.button("◀", key="prev_slide"):
            st.session_state.current_slide = (st.session_state.current_slide - 1) % len(SAMPLE_ANALYSES)
            st.session_state.last_slide_change = time.time()
            st.rerun()
    
    with col_next:
        if st.button("▶", key="next_slide"):
            st.session_state.current_slide = (st.session_state.current_slide + 1) % len(SAMPLE_ANALYSES)
            st.session_state.last_slide_change = time.time()
            st.rerun()
    

# Right column - Chat section
with col2:
    # Chat header
    st.markdown("""
    <div class="chat-header">
        💬 AI Task Planner Chat
    </div>
    """, unsafe_allow_html=True)
    


    
    # Display chat messages
    for message in st.session_state.messages:
        if message["role"] == "user":
            st.markdown(f"""
            <div class="user-message">
                <strong>You:</strong> {message["content"]}
            </div>
            """, unsafe_allow_html=True)
        else:  # AI message
            if message["content"] == "task_plan":
                # Render task plan using native Streamlit components
                st.markdown("""
                <div class="ai-message">
                    <strong>🤖 AI Task Planner:</strong>
                </div>
                """, unsafe_allow_html=True)

                
                for task in message["tasks"]:
                    st.markdown(f"""
                    <div class="task-item">
                        <strong>Task {task['id']}:</strong> {task['task']}
                        <br>
                        <span class="tool-badge">{task['tool']}</span>
                    </div>
                    """, unsafe_allow_html=True)
                
                st.markdown("""
                <div class="download-section">
                    <strong>📄 Task Plan JSON Generated Successfully!</strong>
                    <br><br>
                    View the complete JSON structure below and download it for your records.
                </div>
                </div>
                """, unsafe_allow_html=True)
            else:
                # This is plain text content
                st.markdown(f"""
                <div class="ai-message">
                    <strong>🤖 AI Task Planner:</strong>
                    <br><br>
                    {message["content"]}
                </div>
                """, unsafe_allow_html=True)
    
    # If no messages, show welcome message
    if not st.session_state.messages:
        st.markdown("""
        <div class="ai-message">
            <strong>🤖 Welcome to Geospatial AI Task Planner!</strong>
            <br><br>
            I'm here to help you break down complex geospatial analysis tasks into manageable steps.
            <br><br>
            Simply describe your geospatial analysis needs in the text area below, and I'll create a detailed task plan for you.
        </div>
        """, unsafe_allow_html=True)
    

    

    
    # Text input for query
    user_input = st.text_area(
        "Enter your geospatial analysis query:",
        height=120,
        placeholder="e.g., Create a flood risk assessment map for Mumbai during monsoon season...",
        key="user_input",
        disabled=st.session_state.processing
    )
    
    # Submit button
    if st.button("🚀 Generate Task Plan", type="primary", use_container_width=True, disabled=st.session_state.processing):
        if user_input.strip():
            # Set processing state
            st.session_state.processing = True
            
            # Add user message
            st.session_state.messages.append({"role": "user", "content": user_input})
            
            # Show processing message with spinner
            with st.spinner("🔄 Generating your geospatial task plan..."):
                try:
                    # Call your agent
                    result = call_agent(user_input)
                    
                    # Create JSON structure
                    task_plan_json = {
                        "query": user_input,
                        "timestamp": datetime.now().isoformat(),
                        "task_id": str(uuid.uuid4()),
                        "tasks": result.response,
                        "output_files": result.output_files_path,
                        "status": "completed"
                    }
                    
                    # Store the tasks in session state for rendering
                    st.session_state.generated_json = task_plan_json
                    
                    # Instead of storing HTML, store the task data
                    st.session_state.messages.append({
                        "role": "assistant", 
                        "content": "task_plan",
                        "tasks": result.response
                    })
                    
                    # Show success message
                    st.success("✅ Task plan generated successfully!")
                    
                except Exception as e:
                    error_msg = f"❌ Error processing request: {str(e)}"
                    st.session_state.messages.append({"role": "assistant", "content": error_msg})
                    st.error(error_msg)
                    st.session_state.generated_json = None
                
                # Reset processing state
                st.session_state.processing = False
                
                # Rerun to update the interface
                st.rerun()
        else:
            st.warning("Please enter a query before submitting.")
    
    # Display JSON and download button if JSON is generated
    if st.session_state.generated_json:
        st.markdown("### 📋 Generated Task Plan JSON")
        
        # Display JSON in a formatted container
        json_str = json.dumps(st.session_state.generated_json, indent=2)
        st.markdown(f'<div class="json-container"><pre>{json_str}</pre></div>', unsafe_allow_html=True)
        
        # Download button
        filename = f"geospatial_task_plan_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        st.download_button(
            label="📥 Download Task Plan JSON",
            data=json_str,
            file_name=filename,
            mime="application/json",
            use_container_width=True
        )
    


# Clear chat button (at the bottom)
if st.session_state.messages:
    col_clear1, col_clear2 = st.columns([1, 1])
    with col_clear1:
        if st.button("🗑️ Clear Chat", key="clear_chat", use_container_width=True):
            st.session_state.messages = []
            st.session_state.generated_json = None
            st.session_state.processing = False
            st.rerun()
    
    with col_clear2:
        if st.session_state.generated_json and st.button("🔄 Generate New Plan", key="new_plan", use_container_width=True):
            st.session_state.generated_json = None
            st.rerun()