import streamlit as st
import pandas as pd
from typing import Union, List, Dict, Optional
from groq import Groq
import os
from duckduckgo_search import DDGS
import json
from datetime import datetime, timedelta
import time

# Set page configuration with fullscreen layout and custom theme
st.set_page_config(
    page_title="SmolAgent Travel Planner",
    layout="wide",
    initial_sidebar_state="collapsed",
    page_icon="✈️"
)

# Enhanced CSS for a beautiful travel-themed styling
st.markdown("""
<style>
    /* Base styles */
    html, body, .stApp, .main {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
        color: #ffffff !important;
    }
    
    /* Typography */
    h1, h2, h3, h4, h5, h6 {
        color: #ffffff !important;
        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
        text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
    }
    
    /* Cards and containers */
    .card {
        background: rgba(255, 255, 255, 0.1) !important;
        backdrop-filter: blur(10px) !important;
        border: 1px solid rgba(255, 255, 255, 0.2) !important;
        border-radius: 15px !important;
        box-shadow: 0 8px 32px rgba(31, 38, 135, 0.37) !important;
        padding: 2rem !important;
        margin-bottom: 2rem !important;
    }
    
    /* Agent status cards */
    .agent-card {
        background: rgba(255, 255, 255, 0.15) !important;
        border: 1px solid rgba(255, 255, 255, 0.3) !important;
        border-radius: 12px !important;
        padding: 1.5rem !important;
        margin: 1rem 0 !important;
        backdrop-filter: blur(5px) !important;
    }
    
    /* Input fields */
    .stTextInput input, .stTextArea textarea, .stSelectbox div[data-baseweb="select"] {
        background: rgba(255, 255, 255, 0.2) !important;
        color: #ffffff !important;
        border: 1px solid rgba(255, 255, 255, 0.3) !important;
        border-radius: 8px !important;
    }
    
    .stTextInput input::placeholder, .stTextArea textarea::placeholder {
        color: rgba(255, 255, 255, 0.7) !important;
    }
    
    /* Buttons */
    .stButton > button {
        background: linear-gradient(45deg, #ff6b6b, #ff8e8e) !important;
        color: #ffffff !important;
        border: none !important;
        border-radius: 25px !important;
        font-weight: bold !important;
        padding: 0.75rem 2rem !important;
        transition: all 0.3s ease !important;
        box-shadow: 0 4px 15px rgba(255, 107, 107, 0.4) !important;
    }
    
    .stButton > button:hover {
        transform: translateY(-2px) !important;
        box-shadow: 0 8px 25px rgba(255, 107, 107, 0.6) !important;
    }
    
    /* Progress indicators */
    .agent-status {
        display: flex;
        align-items: center;
        gap: 10px;
        padding: 10px;
        border-radius: 8px;
        margin: 5px 0;
        background: rgba(255, 255, 255, 0.1);
        backdrop-filter: blur(5px);
    }
    
    .status-working {
        border-left: 4px solid #4CAF50;
    }
    
    .status-complete {
        border-left: 4px solid #2196F3;
    }
    
    .status-error {
        border-left: 4px solid #F44336;
    }
    
    /* Spinning animation */
    .spinner {
        border: 2px solid rgba(255, 255, 255, 0.3);
        border-radius: 50%;
        border-top: 2px solid #ffffff;
        width: 20px;
        height: 20px;
        animation: spin 1s linear infinite;
    }
    
    @keyframes spin {
        0% { transform: rotate(0deg); }
        100% { transform: rotate(360deg); }
    }
    
    /* Results styling */
    .travel-result {
        background: rgba(255, 255, 255, 0.05) !important;
        border-left: 4px solid #4CAF50 !important;
        border-radius: 8px !important;
        padding: 1.5rem !important;
        margin: 1rem 0 !important;
    }
    
    /* Tabs */
    .stTabs [aria-selected="true"] {
        background: rgba(255, 255, 255, 0.2) !important;
        color: #ffffff !important;
        border-bottom: 2px solid #ff6b6b !important;
    }
    
    /* Logo styling */
    .logo-text {
        font-size: 3rem !important;
        background: linear-gradient(45deg, #FFD700, #FFA500) !important;
        -webkit-background-clip: text !important;
        -webkit-text-fill-color: transparent !important;
        text-shadow: 2px 2px 4px rgba(0,0,0,0.3) !important;
    }
    
    /* Success indicators */
    .success-badge {
        background: linear-gradient(45deg, #4CAF50, #66BB6A) !important;
        color: white !important;
        padding: 0.5rem 1rem !important;
        border-radius: 20px !important;
        font-size: 0.9rem !important;
    }
</style>
""", unsafe_allow_html=True)


class SmolAgent:
    """
    Base agent class implementing the smol agent architecture.
    Each agent has a specific role and can collaborate with others.
    """
    def __init__(self, name: str, role: str, llm_client, search_tool):
        self.name = name
        self.role = role
        self.llm = llm_client
        self.search = search_tool
        self.status = "idle"
        self.last_result = None
        
    def set_status(self, status: str, message: str = ""):
        self.status = status
        self.status_message = message
        
    def execute_task(self, task_prompt: str, search_query: str = None) -> str:
        """Execute a task with optional web search"""
        try:
            self.set_status("working", f"Processing {self.role.lower()} task...")
            
            # Perform search if needed
            search_results = ""
            if search_query:
                self.set_status("working", "Searching for relevant information...")
                search_results = self.search(search_query)
                
            # Create enhanced prompt with search results
            enhanced_prompt = f"""
            Role: {self.role}
            Task: {task_prompt}
            
            {f"Search Results: {search_results}" if search_results else ""}
            
            Please provide a detailed response based on your role as {self.role}.
            """
            
            # Get response from LLM
            self.set_status("working", "Generating response...")
            result = self.llm(enhanced_prompt)
            
            self.last_result = result
            self.set_status("complete", "Task completed successfully")
            return result
            
        except Exception as e:
            self.set_status("error", f"Error: {str(e)}")
            return f"Error in {self.name}: {str(e)}"


class TravelSearchTool:
    """Enhanced search tool specifically designed for travel planning"""
    def __init__(self):
        self.ddgs = DDGS()
    
    def __call__(self, query: str, search_type: str = "general", max_results: int = 5) -> str:
        try:
            # Customize search based on type
            if search_type == "flights":
                query = f"best flight deals {query} 2024 booking sites"
            elif search_type == "hotels":
                query = f"best hotels accommodation {query} reviews booking"
            elif search_type == "activities":
                query = f"top attractions activities things to do {query} tourist guide"
            elif search_type == "restaurants":
                query = f"best restaurants food dining {query} local cuisine"
            elif search_type == "weather":
                query = f"weather forecast climate {query} best time to visit"
            
            # Perform search
            search_results = list(self.ddgs.text(
                query,
                max_results=max_results,
                region='wt-wt',
                safesearch='on'
            ))
            
            if not search_results:
                return "No relevant information found. Please try a different search."
            
            # Format results
            formatted_results = []
            for idx, result in enumerate(search_results, 1):
                title = result.get('title', 'No title')
                snippet = result.get('body', 'No description')
                url = result.get('href', 'No URL')
                
                formatted_results.append(
                    f"{idx}. {title}\n"
                    f"   Description: {snippet}\n"
                    f"   Source: {url}\n"
                )
            
            return "\n".join(formatted_results)
            
        except Exception as e:
            return f"Search error: {str(e)}"


class GroqLLM:
    """Enhanced LLM client with travel-specific optimizations"""
    def __init__(self, model_name="llama-3.1-70b-versatile"):
        self.client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
        self.model_name = model_name
    
    def __call__(self, prompt: Union[str, dict, List[Dict]]) -> str:
        try:
            completion = self.client.chat.completions.create(
                model=self.model_name,
                messages=[{
                    "role": "user",
                    "content": str(prompt)
                }],
                temperature=0.7,
                max_tokens=2048,
                stream=False
            )
            
            return completion.choices[0].message.content if completion.choices else "No response generated"
            
        except Exception as e:
            return f"LLM Error: {str(e)}"


class TravelPlannerSystem:
    """Main orchestrator for the travel planning system"""
    def __init__(self):
        self.llm = GroqLLM()
        self.search_tool = TravelSearchTool()
        self.agents = self._initialize_agents()
        
    def _initialize_agents(self) -> Dict[str, SmolAgent]:
        """Initialize specialized travel agents"""
        return {
            "destination_expert": SmolAgent(
                "Destination Expert",
                "Travel destination researcher and recommender",
                self.llm, self.search_tool
            ),
            "itinerary_planner": SmolAgent(
                "Itinerary Planner", 
                "Trip itinerary creator and scheduler",
                self.llm, self.search_tool
            ),
            "accommodation_specialist": SmolAgent(
                "Accommodation Specialist",
                "Hotel and lodging finder and recommender",
                self.llm, self.search_tool
            ),
            "activity_curator": SmolAgent(
                "Activity Curator",
                "Tourist attraction and activity recommender",
                self.llm, self.search_tool
            ),
            "budget_advisor": SmolAgent(
                "Budget Advisor",
                "Travel cost estimator and budget optimizer",
                self.llm, self.search_tool
            )
        }
    
    def create_travel_plan(self, destination: str, duration: str, budget: str, 
                          travel_style: str, interests: str) -> Dict[str, str]:
        """Create a comprehensive travel plan using multiple agents"""
        
        # Prepare base context for all agents
        base_context = f"""
        Destination: {destination}
        Duration: {duration}
        Budget: {budget}
        Travel Style: {travel_style}
        Interests: {interests}
        """
        
        results = {}
        
        # Agent 1: Destination Research
        dest_task = f"""
        Research {destination} as a travel destination. Include:
        - Best time to visit and current weather
        - Cultural highlights and local customs
        - Transportation options
        - Safety and travel requirements
        - Local currency and tipping customs
        
        Context: {base_context}
        """
        results["destination_info"] = self.agents["destination_expert"].execute_task(
            dest_task, f"{destination} travel guide weather best time visit"
        )
        
        # Agent 2: Accommodation Research
        accommodation_task = f"""
        Find accommodation options for {destination}. Consider:
        - Hotels, hostels, and alternative lodging based on budget: {budget}
        - Location preferences for {travel_style} travelers
        - Amenities and reviews
        - Booking tips and best deals
        
        Context: {base_context}
        """
        results["accommodation"] = self.agents["accommodation_specialist"].execute_task(
            accommodation_task, f"{destination} hotels accommodation {budget} {travel_style}"
        )
        
        # Agent 3: Activities and Attractions
        activity_task = f"""
        Curate activities and attractions for {destination} based on interests: {interests}
        Include:
        - Must-see attractions and hidden gems
        - Activities matching {travel_style} style
        - Local experiences and cultural activities
        - Restaurant and dining recommendations
        
        Context: {base_context}
        """
        results["activities"] = self.agents["activity_curator"].execute_task(
            activity_task, f"{destination} attractions activities {interests} restaurants"
        )
        
        # Agent 4: Itinerary Planning
        itinerary_task = f"""
        Create a detailed {duration} itinerary for {destination}. Include:
        - Day-by-day schedule optimized for {travel_style} travelers
        - Time management and logistics
        - Mix of planned activities and free time
        - Transportation between locations
        
        Use information about accommodations and activities to create a cohesive plan.
        Context: {base_context}
        """
        results["itinerary"] = self.agents["itinerary_planner"].execute_task(
            itinerary_task, f"{destination} {duration} itinerary travel schedule"
        )
        
        # Agent 5: Budget Planning
        budget_task = f"""
        Create a detailed budget breakdown for the {destination} trip:
        - Accommodation costs based on research
        - Transportation (flights, local transport)
        - Food and dining expenses
        - Activities and attraction fees
        - Miscellaneous expenses and emergency fund
        - Money-saving tips and deals
        
        Context: {base_context}
        """
        results["budget_breakdown"] = self.agents["budget_advisor"].execute_task(
            budget_task, f"{destination} travel costs budget {duration} {budget}"
        )
        
        return results


def display_agent_status(agents: Dict[str, SmolAgent]):
    """Display the status of all agents"""
    st.markdown("### 🤖 Agent Status Dashboard")
    
    cols = st.columns(len(agents))
    for idx, (name, agent) in enumerate(agents.items()):
        with cols[idx]:
            status_class = f"status-{agent.status}" if hasattr(agent, 'status') else "status-idle"
            status_message = getattr(agent, 'status_message', 'Ready')
            
            if agent.status == "working":
                icon = '<div class="spinner"></div>'
            elif agent.status == "complete":
                icon = '✅'
            elif agent.status == "error":
                icon = '❌'
            else:
                icon = '⏸️'
            
            st.markdown(f"""
            <div class="agent-status {status_class}">
                {icon} <strong>{agent.name}</strong><br>
                <small>{status_message}</small>
            </div>
            """, unsafe_allow_html=True)


# Main Application Interface
st.markdown("""
<div class="card" style="text-align: center; padding: 3rem;">
    <h1 class="logo-text">✈️ SmolAgent Travel Planner</h1>
    <p style="font-size: 1.2rem; opacity: 0.9; margin-top: 1rem;">
        AI-powered travel planning with specialized agents working together to create your perfect trip
    </p>
</div>
""", unsafe_allow_html=True)

# Initialize the system
try:
    if 'travel_system' not in st.session_state:
        st.session_state.travel_system = TravelPlannerSystem()
    
    travel_system = st.session_state.travel_system
    
    # Input Section
    st.markdown('<div class="card">', unsafe_allow_html=True)
    st.markdown("## 🌍 Plan Your Perfect Trip")
    
    col1, col2 = st.columns(2)
    
    with col1:
        destination = st.text_input(
            "🎯 Where do you want to go?",
            placeholder="e.g., Tokyo, Japan or Paris, France",
            help="Enter your dream destination"
        )
        
        duration = st.selectbox(
            "📅 Trip Duration",
            ["3 days", "1 week", "2 weeks", "1 month", "Custom"],
            help="How long will you be traveling?"
        )
        
        if duration == "Custom":
            duration = st.text_input("Specify custom duration:", placeholder="e.g., 10 days")
    
    with col2:
        budget = st.selectbox(
            "💰 Budget Range",
            ["Budget ($50-100/day)", "Mid-range ($100-250/day)", "Luxury ($250+/day)", "No specific budget"],
            help="What's your daily budget range?"
        )
        
        travel_style = st.selectbox(
            "🎒 Travel Style",
            ["Adventure & Outdoor", "Cultural & Historical", "Relaxation & Spa", "Nightlife & Entertainment", 
             "Food & Culinary", "Family-Friendly", "Business Travel", "Solo Traveler", "Romantic Getaway"],
            help="What kind of traveler are you?"
        )
    
    interests = st.text_area(
        "🎨 Interests & Preferences",
        placeholder="Tell us about your interests: art, food, history, adventure, shopping, nightlife, nature, etc.",
        help="The more specific you are, the better recommendations we can provide!"
    )
    
    st.markdown('</div>', unsafe_allow_html=True)
    
    # Planning Button
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        plan_trip = st.button("🚀 CREATE MY TRAVEL PLAN", type="primary")
    
    # Results Section
    if plan_trip and destination:
        with st.spinner(""):
            # Display agent status dashboard
            display_agent_status(travel_system.agents)
            
            # Create the travel plan
            st.markdown("### 🔄 Agents Working Together...")
            
            # Progress tracking
            progress_bar = st.progress(0)
            status_placeholder = st.empty()
            
            results = {}
            total_agents = len(travel_system.agents)
            
            # Update progress as each agent completes
            for i, (agent_name, agent) in enumerate(travel_system.agents.items()):
                progress = (i + 1) / total_agents
                progress_bar.progress(progress)
                status_placeholder.text(f"Agent {i+1}/{total_agents}: {agent.name} working...")
                time.sleep(0.5)  # Brief delay for visual effect
            
            # Execute the travel planning
            results = travel_system.create_travel_plan(
                destination, duration, budget, travel_style, interests
            )
            
            progress_bar.progress(1.0)
            status_placeholder.text("✅ All agents completed successfully!")
            
            # Display Results in Tabs
            st.markdown(f'<div class="card"><h2>🎯 Your {destination} Travel Plan</h2></div>', unsafe_allow_html=True)
            
            tab1, tab2, tab3, tab4, tab5 = st.tabs([
                "🌍 Destination Info", "🏨 Accommodation", "🎭 Activities", 
                "📋 Itinerary", "💰 Budget"
            ])
            
            with tab1:
                st.markdown('<div class="travel-result">', unsafe_allow_html=True)
                st.markdown(results.get("destination_info", "Information not available"))
                st.markdown('</div>', unsafe_allow_html=True)
            
            with tab2:
                st.markdown('<div class="travel-result">', unsafe_allow_html=True)
                st.markdown(results.get("accommodation", "Information not available"))
                st.markdown('</div>', unsafe_allow_html=True)
            
            with tab3:
                st.markdown('<div class="travel-result">', unsafe_allow_html=True)
                st.markdown(results.get("activities", "Information not available"))
                st.markdown('</div>', unsafe_allow_html=True)
            
            with tab4:
                st.markdown('<div class="travel-result">', unsafe_allow_html=True)
                st.markdown(results.get("itinerary", "Information not available"))
                st.markdown('</div>', unsafe_allow_html=True)
            
            with tab5:
                st.markdown('<div class="travel-result">', unsafe_allow_html=True)
                st.markdown(results.get("budget_breakdown", "Information not available"))
                st.markdown('</div>', unsafe_allow_html=True)
            
            # Final agent status
            st.markdown("### ✅ Mission Complete!")
            display_agent_status(travel_system.agents)
            
    elif plan_trip:
        st.warning("🚨 Please enter a destination to start planning your trip!")
    
    # Tips and Examples Section
    if not plan_trip or not destination:
        st.markdown('<div class="card">', unsafe_allow_html=True)
        st.markdown("## 💡 How It Works")
        
        col1, col2, col3 = st.columns(3)
        
        with col1:
            st.markdown("""
            ### 🔍 Research Phase
            Our **Destination Expert** searches for:
            - Current weather & best times to visit
            - Local customs & cultural insights  
            - Safety requirements & travel tips
            - Transportation options
            """)
        
        with col2:
            st.markdown("""
            ### 🏨 Planning Phase
            Our specialists find:
            - **Accommodation** matching your budget
            - **Activities** based on your interests
            - **Restaurants** & local dining spots
            - Hidden gems & local experiences
            """)
        
        with col3:
            st.markdown("""
            ### 📋 Organization Phase
            Our planners create:
            - **Day-by-day itinerary** optimized for you
            - **Budget breakdown** with cost estimates
            - **Logistics** and travel coordination
            - **Money-saving tips** and deals
            """)
        
        st.markdown('</div>', unsafe_allow_html=True)
        
        # Example destinations
        st.markdown('<div class="card">', unsafe_allow_html=True)
        st.markdown("## 🌟 Popular Destinations")
        
        example_cols = st.columns(4)
        
        examples = [
            ("🗾 Tokyo, Japan", "tokyo japan"),
            ("🥖 Paris, France", "paris france"), 
            ("🏛️ Rome, Italy", "rome italy"),
            ("🏝️ Bali, Indonesia", "bali indonesia")
        ]
        
        for i, (display_name, destination_value) in enumerate(examples):
            with example_cols[i]:
                if st.button(display_name, key=f"example_{i}"):
                    st.session_state.destination_input = destination_value
                    st.rerun()
        
        st.markdown('</div>', unsafe_allow_html=True)

except Exception as e:
    st.error(f"""
    <div style="display: flex; align-items: center; gap: 12px;">
        <span>❌ Application Error: {str(e)}</span>
    </div>
    """, unsafe_allow_html=True)
    
    st.markdown("""
    <div class="card">
        <h3>🔧 Troubleshooting</h3>
        <ul>
            <li>Ensure your GROQ_API_KEY is set in environment variables</li>
            <li>Check your internet connection for web searches</li>
            <li>Try refreshing the page or restarting the application</li>
        </ul>
    </div>
    """, unsafe_allow_html=True)

# Sidebar controls
with st.sidebar:
    st.markdown("### 🛠️ System Controls")
    
    if st.button("🔄 Reset Application"):
        for key in list(st.session_state.keys()):
            if key.startswith('travel_'):
                del st.session_state[key]
        st.rerun()
    
    st.markdown("### 📊 Agent Performance")
    if 'travel_system' in st.session_state:
        for name, agent in st.session_state.travel_system.agents.items():
            status = getattr(agent, 'status', 'idle')
            st.metric(agent.name, status.title())
    
    st.markdown("### 💬 Feedback")
    feedback = st.text_area("Share your experience:", placeholder="How was your travel plan?")
    if st.button("Submit Feedback"):
        st.success("Thank you for your feedback!")

# Footer
st.markdown("""
<div style="text-align: center; padding: 2rem; opacity: 0.8;">
    <p>SmolAgent Travel Planner v2.0 | Powered by AI Collaboration</p>
</div>
""", unsafe_allow_html=True)