import streamlit as st
import os
import json
from datetime import datetime, timedelta
import base64
from travel import (
    destination_research_task, accommodation_task, transportation_task,
    activities_task, dining_task, itinerary_task, final_report_task,
    run_task
)

# Set page configuration
st.set_page_config(
    page_title="BlockX Travel Itinerary Generator",
    page_icon="✈️",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for better styling
st.markdown("""
<style>
    .main-header {
        font-size: 2.5rem;
        color: #1E88E5;
        text-align: center;
        margin-bottom: 1rem;
    }
    .sub-header {
        font-size: 1.5rem;
        color: #0D47A1;
        margin-top: 2rem;
        margin-bottom: 1rem;
    }
    .info-box {
        background-color: #E3F2FD;
        padding: 1rem;
        border-radius: 0.5rem;
        margin-bottom: 1rem;
    }
    .success-box {
        background-color: #E8F5E9;
        padding: 1rem;
        border-radius: 0.5rem;
        margin-bottom: 1rem;
        border-left: 5px solid #4CAF50;
    }
    .progress-container {
        margin: 1rem 0;
    }
    .step-complete {
        color: #4CAF50;
        font-weight: bold;
    }
    .step-pending {
        color: #9E9E9E;
    }
    .step-active {
        color: #1E88E5;
        font-weight: bold;
    }
    .agent-log {
        background-color: #F5F5F5;
        border-left: 3px solid #1E88E5;
        padding: 0.5rem;
        margin-bottom: 0.5rem;
        font-family: monospace;
    }
    .agent-output {
        background-color: #E8F5E9;
        border-left: 5px solid #4CAF50;
        padding: 1rem;
        margin: 1rem 0;
        border-radius: 0.5rem;
        max-height: 400px;
        overflow-y: auto;
    }
    .footer {
        text-align: center;
        margin-top: 3rem;
        color: #757575;
        font-size: 0.8rem;
    }
    .stButton button {
        background-color: #1E88E5;
        color: white;
        font-weight: bold;
    }
</style>
""", unsafe_allow_html=True)

# Helper function to download HTML file
def get_download_link(html_content, filename):
    b64 = base64.b64encode(html_content.encode()).decode()
    href = f'<a href="data:text/html;base64,{b64}" download="{filename}">Download Itinerary as HTML</a>'
    return href

# Helper function to display progress
def display_progress(current_step, total_steps=7):
    progress_html = '<div class="progress-container">'
    steps = [
        "Destination Research",
        "Accommodation",
        "Transportation",
        "Activities",
        "Dining",
        "Itinerary Creation",
        "Final Report"
    ]
    
    for i, step in enumerate(steps):
        if i < current_step:
            progress_html += f'<p class="step-complete">✓ {step} - Complete</p>'
        elif i == current_step:
            progress_html += f'<p class="step-active">⟳ {step} - In Progress...</p>'
        else:
            progress_html += f'<p class="step-pending">○ {step} - Pending</p>'
    
    progress_html += '</div>'
    return progress_html

# Custom run_task function that updates the UI with logs and shows live agent outputs
def run_task_with_logs(task, input_text, log_container, output_container, results_key=None):
    # Add log message
    log_message = f"🤖 Starting {task.agent.role}..."
    st.session_state.log_messages.append(log_message)
    
    # Update the log container
    with log_container:
        st.markdown("### Agent Activity")
        for msg in st.session_state.log_messages:
            st.markdown(msg)
    
    # Run the actual task
    result = run_task(task, input_text)
    
    # Store result if needed
    if results_key:
        st.session_state.results[results_key] = result
    
    # Add completion log message
    log_message = f"✅ {task.agent.role} completed!"
    st.session_state.log_messages.append(log_message)
    
    # Update the log container again
    with log_container:
        st.markdown("### Agent Activity")
        for msg in st.session_state.log_messages:
            st.markdown(msg)
    
    # Display the agent's output in the output container
    with output_container:
        st.markdown(f"### {task.agent.role} Output")
        st.markdown("""<div class='agent-output'>""" + result + """</div>""", unsafe_allow_html=True)
    
    return result

# Initialize session state
if 'generated_itinerary' not in st.session_state:
    st.session_state.generated_itinerary = None
if 'generation_complete' not in st.session_state:
    st.session_state.generation_complete = False
if 'current_step' not in st.session_state:
    st.session_state.current_step = 0
if 'results' not in st.session_state:
    st.session_state.results = {
        "destination_info": "",
        "accommodation_info": "",
        "transportation_info": "",
        "activities_info": "",
        "dining_info": "",
        "itinerary": "",
        "final_itinerary": ""
    }
if 'log_messages' not in st.session_state:
    st.session_state.log_messages = []
if 'current_output' not in st.session_state:
    st.session_state.current_output = None

# Header
st.markdown('<h1 class="main-header">BlockX Travel Itinerary Generator</h1>', unsafe_allow_html=True)
st.markdown('<p style="text-align: center;">Create your personalized AI-powered travel itinerary in minutes!</p>', unsafe_allow_html=True)

# Sidebar with information
with st.sidebar:
    st.image("https://img.icons8.com/fluency/96/travel-card.png", width=80)
    st.markdown("### About")
    st.info(
        "This AI-powered tool creates a personalized travel itinerary based on your preferences. "
        "Fill in the form and let our specialized travel agents plan your perfect trip!"
    )
    
    st.markdown("### How it works")
    st.markdown(
        "1. Enter your travel details\n"
        "2. Our AI agents analyze your preferences\n"
        "3. Receive a comprehensive itinerary\n"
        "4. Download and enjoy your trip!"
    )
    
    st.markdown("### Travel Agents")
    st.markdown(
        "- 🔍 Destination Research Agent\n"
        "- 🏨 Accommodation Agent\n"
        "- 🚗 Transportation Agent\n"
        "- 🎭 Activities & Attractions Agent\n"
        "- 🍽️ Dining & Culinary Agent\n"
        "- 📅 Itinerary Integration Agent\n"
        "- 📄 Final Report Generation Agent"
    )

# Input Form
if not st.session_state.generation_complete:
    st.markdown('<h2 class="sub-header">Enter Your Travel Details</h2>', unsafe_allow_html=True)
    
    col1, col2 = st.columns(2)
    
    with col1:
        origin = st.text_input("Origin City/Country", placeholder="e.g., New York, USA")
        destination = st.text_input("Destination City/Country", placeholder="e.g., Paris, France")
        duration = st.number_input("Trip Duration (days)", min_value=1, max_value=30, value=7)
    
    with col2:
        budget_options = ["Budget", "Moderate", "Luxury"]
        budget = st.selectbox("Budget Level", budget_options)
        preferences = st.text_area("Travel Preferences & Interests", 
                                 placeholder="e.g., museums, hiking, food, shopping, beaches, history, nightlife, family-friendly")
        special_requirements = st.text_area("Special Requirements", 
                                          placeholder="e.g., dietary restrictions, accessibility needs, etc.")
    
    # Generate button
    if st.button("Generate My Travel Itinerary"):
        if not origin or not destination:
            st.error("Please enter both origin and destination.")
        else:
            user_input = {
                "origin": origin,
                "destination": destination,
                "duration": str(duration),
                "budget": budget.lower(),
                "preferences": preferences,
                "special_requirements": special_requirements
            }
            
            # Format the user input for tasks
            input_context = f"""Travel Request Details:
Origin: {user_input['origin']}
Destination: {user_input['destination']}
Duration: {user_input['duration']} days
Budget Level: {user_input['budget']}
Preferences/Interests: {user_input['preferences']}
Special Requirements: {user_input['special_requirements']}
"""
            
            # Create containers for progress, logs, and live output
            col1, col2 = st.columns([1, 2])
            
            with col1:
                progress_placeholder = st.empty()
                progress_placeholder.markdown(display_progress(0), unsafe_allow_html=True)
                log_container = st.container()
                st.session_state.log_messages = []
            
            with col2:
                output_container = st.container()
                with output_container:
                    st.markdown("### Live Agent Outputs")
                    st.info("Agent outputs will appear here as they are generated")
            
            # Step 1: Destination Research
            st.session_state.current_step = 0
            progress_placeholder.markdown(display_progress(st.session_state.current_step), unsafe_allow_html=True)
            
            destination_info = run_task_with_logs(
                destination_research_task, 
                input_context.format(
                    destination=user_input['destination'],
                    preferences=user_input['preferences']
                ),
                log_container,
                output_container,
                "destination_info"
            )
            
            st.session_state.current_step = 1
            progress_placeholder.markdown(display_progress(st.session_state.current_step), unsafe_allow_html=True)
            
            # Step 2: Accommodation Recommendations
            accommodation_info = run_task_with_logs(
                accommodation_task, 
                input_context.format(
                    destination=user_input['destination'],
                    budget=user_input['budget'],
                    preferences=user_input['preferences']
                ),
                log_container,
                output_container,
                "accommodation_info"
            )
            
            st.session_state.current_step = 2
            progress_placeholder.markdown(display_progress(st.session_state.current_step), unsafe_allow_html=True)
            
            # Step 3: Transportation Planning
            transportation_info = run_task_with_logs(
                transportation_task, 
                input_context.format(
                    origin=user_input['origin'],
                    destination=user_input['destination']
                ),
                log_container,
                output_container,
                "transportation_info"
            )
            
            st.session_state.current_step = 3
            progress_placeholder.markdown(display_progress(st.session_state.current_step), unsafe_allow_html=True)
            
            # Step 4: Activities & Attractions
            activities_info = run_task_with_logs(
                activities_task, 
                input_context.format(
                    destination=user_input['destination'],
                    preferences=user_input['preferences']
                ),
                log_container,
                output_container,
                "activities_info"
            )
            
            st.session_state.current_step = 4
            progress_placeholder.markdown(display_progress(st.session_state.current_step), unsafe_allow_html=True)
            
            # Step 5: Dining Recommendations
            dining_info = run_task_with_logs(
                dining_task, 
                input_context.format(
                    destination=user_input['destination'],
                    preferences=user_input['preferences']
                ),
                log_container,
                output_container,
                "dining_info"
            )
            
            st.session_state.current_step = 5
            progress_placeholder.markdown(display_progress(st.session_state.current_step), unsafe_allow_html=True)
            
            # Step 6: Create Day-by-Day Itinerary
            combined_info = f"""{input_context}

Destination Information:
{destination_info}

Accommodation Options:
{accommodation_info}

Transportation Plan:
{transportation_info}

Recommended Activities:
{activities_info}

Dining Recommendations:
{dining_info}
"""
            
            itinerary = run_task_with_logs(
                itinerary_task, 
                combined_info.format(
                    duration=user_input['duration'],
                    origin=user_input['origin'],
                    destination=user_input['destination']
                ),
                log_container,
                output_container,
                "itinerary"
            )
            
            st.session_state.current_step = 6
            progress_placeholder.markdown(display_progress(st.session_state.current_step), unsafe_allow_html=True)
            
            # Step 7: Generate Final Report
            final_context = f"""{combined_info}

Day-by-Day Itinerary:
{itinerary}
"""
            
            final_itinerary = run_task_with_logs(
                final_report_task, 
                final_context.format(
                    duration=user_input['duration'],
                    origin=user_input['origin'],
                    destination=user_input['destination']
                ),
                log_container,
                output_container,
                "final_itinerary"
            )
            
            # Save the generated itinerary to session state
            st.session_state.generated_itinerary = final_itinerary
            st.session_state.generation_complete = True
            
            # Create a filename based on the destination and date
            date_str = datetime.now().strftime("%Y-%m-%d")
            st.session_state.filename = f"{user_input['destination'].replace(' ', '_')}_{date_str}_itinerary.html"
            
            # No need to rerun, we'll update the UI directly

# Display results if generation is complete
if st.session_state.generation_complete:
    st.markdown('<h2 class="sub-header">Your Travel Itinerary is Ready! 🎉</h2>', unsafe_allow_html=True)
    
    # Display download link
    st.markdown(
        get_download_link(st.session_state.generated_itinerary, st.session_state.filename),
        unsafe_allow_html=True
    )
    
    # Option to view the itinerary directly
    with st.expander("Preview Your Itinerary", expanded=True):
        st.components.v1.html(st.session_state.generated_itinerary, height=600, scrolling=True)
    
    # Option to view individual agent outputs
    with st.expander("View Detailed Agent Outputs"):
        agent_tabs = st.tabs([
            "Destination", "Accommodation", "Transportation", 
            "Activities", "Dining", "Day-by-Day"
        ])
        
        with agent_tabs[0]:
            st.markdown("### Destination Research")
            st.markdown(st.session_state.results["destination_info"])
            
        with agent_tabs[1]:
            st.markdown("### Accommodation Options")
            st.markdown(st.session_state.results["accommodation_info"])
            
        with agent_tabs[2]:
            st.markdown("### Transportation Plan")
            st.markdown(st.session_state.results["transportation_info"])
            
        with agent_tabs[3]:
            st.markdown("### Recommended Activities")
            st.markdown(st.session_state.results["activities_info"])
            
        with agent_tabs[4]:
            st.markdown("### Dining Recommendations")
            st.markdown(st.session_state.results["dining_info"])
            
        with agent_tabs[5]:
            st.markdown("### Day-by-Day Itinerary")
            st.markdown(st.session_state.results["itinerary"])
    
    # Option to start over
    if st.button("Create Another Itinerary"):
        # Reset session state
        st.session_state.generated_itinerary = None
        st.session_state.generation_complete = False
        st.session_state.current_step = 0
        st.session_state.results = {
            "destination_info": "",
            "accommodation_info": "",
            "transportation_info": "",
            "activities_info": "",
            "dining_info": "",
            "itinerary": "",
            "final_itinerary": ""
        }
        st.rerun()

# Footer
st.markdown('<div class="footer">BlockX Travel Itinerary Generator © 2025</div>', unsafe_allow_html=True)