app.py

import streamlit as st
import json
from datetime import datetime
import uuid
import os
import time

# Import your agent function (assuming it's in the same directory or properly imported)
# from your_agent_file import call_agent

# Mock function for demonstration - replace with your actual call_agent function
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"}
    ]
    
    # Mock state object
    class MockState:
        def __init__(self):
            self.query = query
            self.response = mock_response
            self.output_files_path = [f"outputs/geospatial_plan_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:6]}.json"]
    
    return MockState()

# Sample analysis data - Replace with your actual image URLs and captions
SAMPLE_ANALYSES = [
    {
        "image_url": "output.png",
        "caption": "Flood vulnerability analysis of Chennai during monsoon season",
        "query": "Create a flood risk assessment map for Chennai during monsoon season"
    },
    {
        "image_url": "https://via.placeholder.com/800x600/2196F3/FFFFFF?text=Urban+Heat+Island+Mumbai",
        "caption": "Urban heat island analysis of Mumbai metropolitan area",
        "query": "Analyze urban heat island effect in Mumbai using satellite imagery"
    },
    {
        "image_url": "https://via.placeholder.com/800x600/FF9800/FFFFFF?text=Deforestation+Amazon",
        "caption": "Deforestation monitoring in Amazon rainforest",
        "query": "Monitor deforestation patterns in Amazon rainforest using multi-temporal satellite data"
    },
    {
        "image_url": "https://via.placeholder.com/800x600/9C27B0/FFFFFF?text=Crop+Yield+Punjab",
        "caption": "Crop yield prediction analysis for Punjab agricultural regions",
        "query": "Predict crop yield for wheat cultivation in Punjab using satellite data"
    },
    {
        "image_url": "https://via.placeholder.com/800x600/F44336/FFFFFF?text=Wildfire+California",
        "caption": "Wildfire risk assessment and spread modeling",
        "query": "Create wildfire risk maps and predict fire spread patterns for California"
    }
]

# 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
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; }
    }
    
    /* 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://www.isro.gov.in/sites/default/files/flipping_book/ISRO%20Annual%20Report%202022-23/assets/basic-html/page-1_1.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)
    
    # Slideshow container
    st.markdown("""
        <div class="slideshow-container">
    """, unsafe_allow_html=True)
    
    # Generate slides
    for i, analysis in enumerate(SAMPLE_ANALYSES):
        active_class = "active" if i == st.session_state.current_slide else ""
        st.markdown(f"""
            <div class="slide {active_class}">
                <img src="{analysis['image_url']}" alt="Sample Analysis" class="slide-image">
                <div class="slide-caption">{analysis['caption']}</div>
            </div>
        """, unsafe_allow_html=True)
    
    # Navigation buttons
    col_prev, col_next = st.columns([1, 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()
    
    # Slide indicators
    st.markdown("""
        <div class="slide-indicators">
    """, unsafe_allow_html=True)
    
    indicator_cols = st.columns(len(SAMPLE_ANALYSES))
    for i, col in enumerate(indicator_cols):
        with col:
            active_class = "active" if i == st.session_state.current_slide else ""
            if st.button("●", key=f"indicator_{i}"):
                st.session_state.current_slide = i
                st.session_state.last_slide_change = time.time()
                st.rerun()
    
    st.markdown("""
        </div>
        </div>
    </div>
    """, unsafe_allow_html=True)
# Right column - Chat section
with col2:
    # Chat header
    st.markdown("""
    <div class="chat-header">
        💬 AI Task Planner Chat
    </div>
    """, unsafe_allow_html=True)
    
    # Chat messages container
    st.markdown('<div class="chat-messages">', 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)
                
                # Use Streamlit's native components for better rendering
                st.markdown("""
                <div class="task-plan-container">
                """, 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)
    
    st.markdown('</div>', unsafe_allow_html=True)
    
    # Input section
    st.markdown('<div class="input-section">', 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
        )
    
    st.markdown('</div>', unsafe_allow_html=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()