import streamlit as st
from datetime import datetime, timedelta
import base64
import pandas as pd
import pydeck as pdk
from geopy.geocoders import Nominatim
import qrcode
from io import BytesIO

from travel import (
    destination_research_task, accommodation_task, transportation_task,
    activities_task, dining_task, itinerary_task, chatbot_task,
    run_task, update_llm
)

st.set_page_config(
    page_title="Your Agentic AI Travelling Partner ❤️ ",
    page_icon="✈️ Namude Yatra",
    layout="wide",
    initial_sidebar_state="expanded"
)

with st.sidebar:
    st.image("image.png", width=250)
    st.markdown("""
<h3>Namude Yatra: The beginning of your dream journey</h3>
<p style="font-size: 14px; color: var(--text-muted);">AI-Powered Agentic Travel Planning</p>
""", unsafe_allow_html=True)
    
    api_key = st.text_input("Enter your GEMINI API Key", type="password")
    if api_key:
        from travel import init_travel_system  
        init_travel_system(api_key)
        st.success("Travel system initialized with the provided API key!")

    st.markdown("""
<style>
    :root {
        --primary: #3a86ff;
        --primary-light: #4895ef;
        --primary-dark: #2667ff;
        --secondary: #4cc9f0;
        --accent: #4361ee;
        --background: #f8f9fa;
        --card-bg: #ffffff;
        --text: #212529;
        --text-muted: #495057;
        --border: #e9ecef;
    }

    .agent-item {
        margin-bottom: 10px;
    }
    .agent-item strong {
        font-size: 18px;
        color: var(--primary-dark);
    }
    .agent-item small {
        font-size: 15px;
    }
</style>
""", unsafe_allow_html=True)
    
    st.markdown("""
    <div class="sidebar-section">
        <h2 style="color: var(--primary);">About</h2>
        <p>This tool creates a personalized travel itinerary based on your preferences. Fill in the form on the main page and let our expert travel agents plan your perfect trip!</p>
    </div>
    <div class="sidebar-section">
        <h2 style="color: var(--primary);">How It Works</h2>
        <ul>
            <li><b>Enter:</b> Your travel details</li>
            <li><b>Analyze:</b> Preferences by AI</li>
            <li><b>Generate:</b> Your comprehensive itinerary</li>
            <li><b>Download:</b> Save and share your itinerary</li>
        </ul>
    </div>
    <div class="sidebar-section">
        <h2>Travel Agents</h2>
        <div class="agent-item"><strong>🔭 Research Specialist</strong><br><small>Finds the best destinations based on your preferences.</small></div>
        <div class="agent-item"><strong>🏨 Accommodation Expert</strong><br><small>Suggests suitable hotels and stays.</small></div>
        <div class="agent-item"><strong>🚆 Transportation Planner</strong><br><small>Plans efficient travel routes.</small></div>
        <div class="agent-item"><strong>🎯 Activities Curator</strong><br><small>Recommends activities tailored to your interests.</small></div>
        <div class="agent-item"><strong>🍽️ Dining Connoisseur</strong><br><small>Finds the best dining experiences.</small></div>
        <div class="agent-item"><strong>📅 Itinerary Creator</strong><br><small>Puts everything together in a daily plan.</small></div>
    </div>
    """, unsafe_allow_html=True)

def get_download_link(text_content, filename):
    b64 = base64.b64encode(text_content.encode()).decode()
    return f'<a class="download-link" href="data:text/plain;base64,{b64}" download="{filename}"><i>📥</i> Save Your Itinerary</a>'

def display_modern_progress(current_step, total_steps=6):
    if 'progress_steps' not in st.session_state:
        st.session_state.progress_steps = {
            0: {'status': 'pending', 'name': "Trip Details"},
            1: {'status': 'pending', 'name': "About"},
            2: {'status': 'pending', 'name': "Travel Style"},
            3: {'status': 'pending', 'name': "Live Agent Outputs"},
            4: {'status': 'pending', 'name': "Download & Share"},
            5: {'status': 'pending', 'name': "Full Itinerary"}
        }
    
    for i in range(total_steps):
        if i < current_step:
            st.session_state.progress_steps[i]['status'] = 'complete'
        elif i == current_step:
            st.session_state.progress_steps[i]['status'] = 'active'
        else:
            st.session_state.progress_steps[i]['status'] = 'pending'
    
    progress_percentage = (current_step / total_steps) * 100
    st.progress(progress_percentage / 100)
    
    st.markdown("""
    <style>
    .compact-progress {
        background: var(--background);
        border-radius: 10px;
        padding: 15px;
        box-shadow: 0 4px 6px rgba(0,0,0,0.1);
        margin-bottom: 20px;
    }
    .step-grid {
        display: grid;
        grid-template-columns: repeat(3, 1fr);
        gap: 10px;
    }
    .step-item {
        display: flex;
        align-items: center;
        padding: 8px 10px;
        border-radius: 6px;
        background: var(--card-bg);
        box-shadow: 0 1px 3px rgba(0,0,0,0.05);
    }
    .step-item.complete {
        border-left: 3px solid #4CAF50;
        background: #f1f8e9;
    }
    .step-item.active {
        border-left: 3px solid #2196F3;
        background: #e3f2fd;
        font-weight: bold;
    }
    .step-item.pending {
        border-left: 3px solid #9e9e9e;
        opacity: 0.7;
    }
    .step-icon {
        margin-right: 8px;
        font-size: 14px;
    }
    .step-text {
        font-size: 13px;
        color: var(--text);
    }
    </style>
    <div class="compact-progress">
      <div class="step-grid">
    """, unsafe_allow_html=True)
    
    for i, step_info in st.session_state.progress_steps.items():
        status = step_info['status']
        name = step_info['name']
        if status == 'complete':
            icon = "✅"
            status_class = "complete"
        elif status == 'active':
            icon = "🔄"
            status_class = "active"
        else:
            icon = "⭕"
            status_class = "pending"
        st.markdown(f"""
        <div class="step-item {status_class}">
            <span class="step-icon">{icon}</span>
            <span class="step-text">{name}</span>
        </div>
        """, unsafe_allow_html=True)
    
    st.markdown('</div></div>', unsafe_allow_html=True)
    return progress_percentage

def update_step_status(step_index, status):
    if 'progress_steps' in st.session_state and step_index in st.session_state.progress_steps:
        st.session_state.progress_steps[step_index]['status'] = status

def run_task_with_logs(task, input_text, log_container, output_container, results_key=None):
    if 'log_messages' not in st.session_state:
        st.session_state.log_messages = []
    log_message = f"🤖 Starting {task.agent.role}..."
    st.session_state.log_messages.append(log_message)
    
    with log_container:
        st.markdown("### Agent Activity")
        for msg in st.session_state.log_messages:
            st.markdown(msg)
    
    result = run_task(task, input_text)
    
    if results_key:
        st.session_state.results[results_key] = result
    log_message = f"✅ {task.agent.role} completed!"
    st.session_state.log_messages.append(log_message)
    
    with log_container:
        st.markdown("### Agent Activity")
        for msg in st.session_state.log_messages:
            st.markdown(msg)
    
    with output_container:
        st.markdown(f"### {task.agent.role} Output")
        st.markdown("<div class='agent-output'>" + result + "</div>", unsafe_allow_html=True)
    return result

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
if 'form_submitted' not in st.session_state:
    st.session_state.form_submitted = False

st.markdown(f"""
<div class="animate-in" style="text-align: center;">
    <div style="margin-bottom: 20px;">
    </div>
    <h1 class="main-header">Your Agentic AI Travelling Partner ❤️</h1>
    <p style="font-size: 1.2rem; color: var(--text-muted); margin-bottom: 25px;">
        ✨ Create your personalized AI-powered travel itinerary in minutes! ✨
    </p>
</div>
""", unsafe_allow_html=True)

st.markdown('<hr style="height:3px;border:none;background-color:#f0f0f0;margin-bottom:25px;">', unsafe_allow_html=True)

if not st.session_state.generation_complete:
    st.markdown('<div class="modern-card animate-in">', unsafe_allow_html=True)
    st.markdown("<h3 style='font-weight: 600; color: var(--primary-dark); text-align: center;'>✈️ Create Your Itinerary</h3>", unsafe_allow_html=True)
    st.markdown("""
    <p style="text-align: center; color: var(--text-muted); font-size: 14px;">Complete the form below for a personalized travel plan.</p>
    """, unsafe_allow_html=True)
    with st.form("travel_form"):
        col1, col2 = st.columns(2)
        with col1:
            st.markdown('<p style="font-weight: 500; color: var(--primary); font-size: 14px;">Trip Details</p>', unsafe_allow_html=True)
            origin = st.text_input("Origin", placeholder="e.g., New York, USA")
            destination = st.text_input("Destination", placeholder="e.g., Paris, France")
            st.markdown('<p style="font-size: 14px;">Travel Dates</p>', unsafe_allow_html=True)
            start_date = st.date_input("Start Date", min_value=datetime.now(), label_visibility="collapsed")
            duration = st.slider("Duration (days)", min_value=1, max_value=30, value=7)
            end_date = start_date + timedelta(days=duration-1)
            st.markdown(f'<p style="font-size: 13px; color: var(--text-muted); text-align: center;">{start_date.strftime("%b %d")} - {end_date.strftime("%b %d, %Y")}</p>', unsafe_allow_html=True)
        with col2:
            st.markdown('<p style="font-weight: 500; color: var(--primary); font-size: 14px;">Preferences</p>', unsafe_allow_html=True)
            travelers = st.number_input("Travelers", min_value=1, max_value=15, value=2)
            budget_options = ["Budget", "Moderate", "Luxury"]
            budget = st.selectbox("Budget", budget_options, help="Budget: Economy options | Moderate: Mid-range | Luxury: High-end experiences")
            travel_style = st.multiselect("Travel Style", options=["Culture", "Adventure", "Relaxation", "Food & Dining", "Nature", "Shopping", "Nightlife", "Family-friendly"], default=["Culture", "Food & Dining"])
        with st.expander("Additional Preferences"):
            preferences = st.text_area("Interests", placeholder="History museums, local cuisine, hiking, art...")
            special_requirements = st.text_area("Special Requirements", placeholder="Dietary restrictions, accessibility needs...")
        submit_button = st.form_submit_button("🚀 Create My Personal Travel Itinerary")
    st.markdown('</div>', unsafe_allow_html=True)

    if submit_button:
        if not origin or not destination:
            st.error("Please enter both origin and destination.")
        else:
            st.session_state.form_submitted = True
            st.session_state.destination = destination  
            user_input = {
                "origin": origin,
                "destination": destination,
                "duration": str(duration),
                "travel_dates": f"{start_date.strftime('%Y-%m-%d')} to {end_date.strftime('%Y-%m-%d')}",
                "travelers": str(travelers),
                "budget": budget.lower(),
                "travel_style": ", ".join(travel_style),
                "preferences": preferences,
                "special_requirements": special_requirements
            }
            st.session_state.user_input = user_input
            input_context = f"""Travel Request Details:
Origin: {user_input['origin']}
Destination: {user_input['destination']}
Duration: {user_input['duration']} days
Travel Dates: {user_input['travel_dates']}
Travelers: {user_input['travelers']}
Budget Level: {user_input['budget']}
Travel Style: {user_input['travel_style']}
Preferences/Interests: {user_input['preferences']}
Special Requirements: {user_input['special_requirements']}
"""
            modified_input_context = "Please output the response in English.\n" + input_context

            # Processing Animation
            st.markdown("""
            <div style="text-align: center; padding:20px 0;">
              <div style="position:relative; width:50px; height:50px;">
                <div style="position:absolute; width:100%; height:100%; border:2px solid #4361ee; border-radius:50%; animation:pulse 1.5s ease-out infinite;"></div>
                <div style="position:absolute; left:50%; top:50%; transform:translate(-50%, -50%); width:12px; height:12px; background-color:#4361ee; border-radius:50%; box-shadow:0 0 8px rgba(67,97,238,0.6);"></div>
              </div>
            </div>
            <style>
            @keyframes pulse {
              0% { transform: scale(0.1); opacity: 0; }
              50% { opacity: 0.5; }
              100% { transform: scale(1); opacity: 0; }
            }
            </style>
            """, unsafe_allow_html=True)

            st.markdown('<div class="modern-card">', unsafe_allow_html=True)
            progress_tab, logs_tab, details_tab = st.tabs(["📊 Progress", "🔄 Live Activity", "📋 Your Travel Request"])
            with details_tab:
                st.markdown("#### Your Travel Request")
                st.markdown("**Destination:** " + user_input['destination'])
                st.markdown("**From:** " + user_input['origin'])
                st.markdown("**When:** " + user_input['travel_dates'] + " (" + user_input['duration'] + " days)")
                st.markdown("**Budget:** " + user_input['budget'].title())
                st.markdown("**Travel Style:** " + user_input['travel_style'])
                if user_input['preferences']:
                    st.markdown("**Interests:** " + user_input['preferences'])
                if user_input['special_requirements']:
                    st.markdown("**Special Requirements:** " + user_input['special_requirements'])
            with progress_tab:
                if 'progress_placeholder' not in st.session_state:
                    st.session_state.progress_placeholder = st.empty()
                with st.session_state.progress_placeholder.container():
                    display_modern_progress(st.session_state.current_step)
            with logs_tab:
                log_container = st.container()
                st.session_state.log_messages = []
            st.markdown('</div>', unsafe_allow_html=True)
            
            output_container = st.container()
            with output_container:
                st.markdown('<div class="modern-card">', unsafe_allow_html=True)
                st.markdown("### Live Agent Outputs")
                st.info("Our AI agents will show their work here as they create your itinerary")
                st.markdown('</div>', unsafe_allow_html=True)
            
            st.session_state.current_step = 0
            
            # Step 1: Destination Research
            update_step_status(0, 'active')
            with st.session_state.progress_placeholder.container():
                display_modern_progress(st.session_state.current_step)
            destination_info = run_task_with_logs(
                destination_research_task,
                modified_input_context.format(destination=user_input['destination'], preferences=user_input['preferences']),
                log_container,
                output_container,
                "destination_info"
            )
            update_step_status(0, 'complete')
            st.session_state.current_step = 1
            
            # Step 2: Accommodation
            update_step_status(1, 'active')
            with st.session_state.progress_placeholder.container():
                display_modern_progress(st.session_state.current_step)
            accommodation_info = run_task_with_logs(
                accommodation_task,
                modified_input_context.format(destination=user_input['destination'], budget=user_input['budget'], preferences=user_input['preferences']),
                log_container,
                output_container,
                "accommodation_info"
            )
            update_step_status(1, 'complete')
            st.session_state.current_step = 2
            
            # Step 3: Transportation
            update_step_status(2, 'active')
            with st.session_state.progress_placeholder.container():
                display_modern_progress(st.session_state.current_step)
            transportation_info = run_task_with_logs(
                transportation_task,
                modified_input_context.format(origin=user_input['origin'], destination=user_input['destination']),
                log_container,
                output_container,
                "transportation_info"
            )
            update_step_status(2, 'complete')
            st.session_state.current_step = 3
            
            # Step 4: Activities
            update_step_status(3, 'active')
            with st.session_state.progress_placeholder.container():
                display_modern_progress(st.session_state.current_step)
            activities_info = run_task_with_logs(
                activities_task,
                modified_input_context.format(destination=user_input['destination'], preferences=user_input['preferences']),
                log_container,
                output_container,
                "activities_info"
            )
            update_step_status(3, 'complete')
            st.session_state.current_step = 4
            
            # Step 5: Dining
            update_step_status(4, 'active')
            with st.session_state.progress_placeholder.container():
                display_modern_progress(st.session_state.current_step)
            dining_info = run_task_with_logs(
                dining_task,
                modified_input_context.format(destination=user_input['destination'], preferences=user_input['preferences']),
                log_container,
                output_container,
                "dining_info"
            )
            update_step_status(4, 'complete')
            st.session_state.current_step = 5
            
            # Step 6: Final Itinerary
            update_step_status(5, 'active')
            with st.session_state.progress_placeholder.container():
                display_modern_progress(st.session_state.current_step)
            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"
            )
            update_step_status(5, 'complete')
            st.session_state.current_step = 6
            with st.session_state.progress_placeholder.container():
                display_modern_progress(st.session_state.current_step)
            
            st.session_state.generated_itinerary = itinerary
            st.session_state.generation_complete = True
            date_str = datetime.now().strftime("%Y-%m-%d")
            st.session_state.filename = f"{user_input['destination'].replace(' ', '_')}_{date_str}_itinerary.txt"

if st.session_state.generation_complete:
    st.markdown("""
    <div class="modern-card animate-in">
      <div style="display: flex; justify-content: center; margin-bottom: 20px;">
        <div class="success-animation">
          <svg class="checkmark" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 52 52">
            <circle class="checkmark__circle" cx="26" cy="26" r="25" fill="none" />
            <path class="checkmark__check" fill="none" d="M14.1 27.2l7.1 7.2 16.7-16.8" />
          </svg>
        </div>
      </div>
      <h2 style="text-align: center; color: var(--primary-dark);">Your Travel Itinerary is Ready! 🎉</h2>
      <p style="text-align: center; color: var(--text-muted); margin-bottom: 20px;">We've created a personalized travel experience just for you. Explore your itinerary below.</p>
    </div>
    <style>
    .success-animation {
      width: 100px;
      height: 100px;
      position: relative;
    }
    .checkmark {
      width: 100px;
      height: 100px;
      border-radius: 50%;
      display: block;
      stroke-width: 2;
      stroke: var(--primary-dark);
      stroke-miterlimit: 10;
      box-shadow: 0 0 20px rgba(38,103,255,0.3);
      animation: fill 0.4s ease-in-out 0.4s forwards, scale 0.3s ease-in-out 0.9s both;
    }
    .checkmark__circle {
      stroke-dasharray: 166;
      stroke-dashoffset: 166;
      stroke-width: 2;
      stroke-miterlimit: 10;
      stroke: var(--primary-dark);
      fill: none;
      animation: stroke 0.6s cubic-bezier(0.65, 0, 0.45, 1) forwards;
    }
    .checkmark__check {
      transform-origin: 50% 50%;
      stroke-dasharray: 48;
      stroke-dashoffset: 48;
      animation: stroke 0.3s cubic-bezier(0.65, 0, 0.45, 1) 0.8s forwards;
    }
    @keyframes stroke {
      100% { stroke-dashoffset: 0; }
    }
    @keyframes scale {
      0%, 100% { transform: none; }
      50% { transform: scale3d(1.1, 1.1, 1); }
    }
    @keyframes fill {
      100% { box-shadow: 0 0 20px rgba(38,103,255,0.3); }
    }
    </style>
    """, unsafe_allow_html=True)
    
    itinerary_tab, details_tab, download_tab, map_tab, chatbot_tab = st.tabs([
        "🗒️ Full Itinerary", 
        "💼 Details", 
        "💾 Download & Share",
        "🗺️ Map & Visualization",
        "🤖 Chatbot Interface"
    ])
    
    with itinerary_tab:
        st.text_area("Your Itinerary", st.session_state.generated_itinerary, height=600)
    
    with details_tab:
        agent_tabs = st.tabs(["🌎 Destination", "🏨 Accommodation", "🚗 Transportation", "🎭 Activities", "🍽️ Dining"])
        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 download_tab:
        col1, col2 = st.columns([2, 1])
        with col1:
            st.markdown("### Save Your Itinerary")
            st.markdown("Download your personalized travel plan to access it offline or share with your travel companions.")
            st.markdown("""
            <div style="background-color: var(--background); padding: 15px; border-radius: 10px; margin-top: 20px;">
                <h4>Your Itinerary File</h4>
                <p style="font-size: 0.9rem; color: var(--text-muted);">Text format - Can be opened in any text editor</p>
            </div>
            """, unsafe_allow_html=True)
            st.markdown(
                "<div style='margin: 10px 0;'>"
                + get_download_link(st.session_state.generated_itinerary, st.session_state.filename)
                + "</div>",
                unsafe_allow_html=True
            )
            st.markdown("### Share Your Itinerary")
            st.markdown("*Coming soon: Additional sharing features (email, phone, etc.)*")
        with col2:
            st.markdown("### Save for Mobile")
            st.markdown("*Coming soon: Additional mobile features such as a dedicated app or push notifications*")
    
    with map_tab:
        st.markdown("### Destination Map")
        dest = st.session_state.get("destination", "Paris")
        try:
            geolocator = Nominatim(user_agent="travel_app")
            location = geolocator.geocode(dest)
            if location:
                lat, lon = location.latitude, location.longitude
            else:
                st.error("The destination could not be found. Using default coordinates.")
                lat, lon = 48.8566, 2.3522 
        except Exception as e:
            st.error("Geocoding error: " + str(e))
            lat, lon = 48.8566, 2.3522
        
        map_data = pd.DataFrame({
            "lat": [lat],
            "lon": [lon],
            "name": [dest]
        })
        st.map(map_data)
        st.markdown("#### Interactive Map with Pydeck")
        layer = pdk.Layer(
            "ScatterplotLayer",
            data=map_data,
            get_position='[lon, lat]',
            get_color='[200, 30, 0, 160]',
            get_radius=200,
        )
        view_state = pdk.ViewState(
            latitude=lat,
            longitude=lon,
            zoom=12,
            pitch=50,
        )
        deck_chart = pdk.Deck(layers=[layer], initial_view_state=view_state)
        st.pydeck_chart(deck_chart)
    
    with chatbot_tab:
        st.markdown("### AI Chatbot Interface")
        if "chat_history" not in st.session_state:
            st.session_state.chat_history = []
        user_message = st.text_input("Enter your message:", key="chat_input")
        if st.button("Send", key="send_button"):
            if user_message:
                response = run_task(chatbot_task, user_message)
                st.session_state.chat_history.append({
                    "speaker": "User",
                    "message": user_message,
                    "time": datetime.now()
                })
                st.session_state.chat_history.append({
                    "speaker": "AI",
                    "message": response,
                    "time": datetime.now()
                })
        st.markdown("<div style='max-height:400px; overflow-y:auto; padding:10px; border:1px solid var(--border); border-radius:6px;'>", unsafe_allow_html=True)
        for chat in st.session_state.chat_history:
            time_str = chat["time"].strftime("%H:%M:%S")
            st.markdown(f"**{chat['speaker']}** ({time_str}): {chat['message']}")
        st.markdown("</div>", unsafe_allow_html=True)

st.markdown("""
<div style="margin-top: 50px; text-align: center; padding: 20px; color: var(--text-muted); font-size: 0.8rem;">
    <p>Built with ❤️ for you</p>
</div>
""", unsafe_allow_html=True)