Spaces:

Alakbarli
/

Assistant

Sleeping

Update app.py
4c154fd verified 10 months ago
51.2 kB
import streamlit as st
import os.path
import os
import time
import base64
import openai
import re
import datetime
import json
import pytz
import extra_streamlit_components as st_tags
from datetime import date, timedelta

from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials
from google_auth_oauthlib.flow import InstalledAppFlow
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError


# --- CONFIGURATION ---
os.environ['OAUTHLIB_INSECURE_TRANSPORT'] = '1'

OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")
GCP_CREDENTIALS_STR = os.environ.get("GCP_CREDENTIALS")
CHATGPT_MODEL = "openai/gpt-oss-20b:free"
OPENROUTER_BASE_URL = 'https://openrouter.ai/api/v1'
TOKEN_FILE = "token.json"
TIMEZONES = pytz.common_timezones

GMAIL_SCOPES = ['https://www.googleapis.com/auth/gmail.modify']
CALENDAR_SCOPES = ['https://www.googleapis.com/auth/calendar.events']
ALL_SCOPES = GMAIL_SCOPES + CALENDAR_SCOPES

if 'cookie_manager' not in st.session_state:
    st.session_state.cookie_manager = st_tags.CookieManager(key="maint_cm")

# --- LLM Client Initialization ---
client = None
if OPENROUTER_API_KEY:
    try:
        client = openai.OpenAI(
            base_url=OPENROUTER_BASE_URL,
            api_key=OPENROUTER_API_KEY,
            default_headers={'HTTP-Referer': 'https://huggingface.co', 'X-Title': 'AI Org Hub'}
        )
    except Exception as e:
        st.error(f"❌ Error during LLM client initialization: {e}")


import extra_streamlit_components as st_tags
from google.oauth2.credentials import Credentials

# Initialize Cookie Manager at the top of main()
cookie_manager = st_tags.CookieManager()


def authenticate_service():
    cm = st.session_state.cookie_manager
    
    # 1. Check current session memory (Fastest)
    if 'google_creds' in st.session_state and st.session_state.google_creds:
        return st.session_state.google_creds

    # 2. Try to recover from Browser Cookie
    cookie_token = cm.get(cookie="google_creds_token")
    
    # If cookie not immediately found, wait for JS component to sync
    if cookie_token is None and "code" not in st.query_params:
        with st.spinner("Restoring session..."):
            time.sleep(1.2) # Essential delay for CookieManager
            cookie_token = cm.get(cookie="google_creds_token")
    
    if cookie_token:
        try:
            creds_info = json.loads(cookie_token)
            creds = Credentials.from_authorized_user_info(creds_info, ALL_SCOPES)
            
            if creds and creds.expired and creds.refresh_token:
                creds.refresh(Request())
                cm.set("google_creds_token", creds.to_json(), key="sync_refresh")
            
            st.session_state.google_creds = creds
            return creds
        except Exception:
            cm.delete("google_creds_token")

    # 3. Handle OAuth Redirect (Code Exchange)
    if "code" in st.query_params:
        try:
            client_config = json.loads(GCP_CREDENTIALS_STR)
            flow = InstalledAppFlow.from_client_config(
                client_config, 
                scopes=ALL_SCOPES,
                redirect_uri="https://alakbarli-assistant.hf.space/"
            )
            flow.fetch_token(code=st.query_params["code"])
            creds = flow.credentials
            
            st.session_state.google_creds = creds
            cm.set("google_creds_token", creds.to_json(), key="save_new_login")
            
            st.query_params.clear()
            st.rerun()
        except Exception:
            st.query_params.clear()
            st.rerun()
            
    # 4. Show Login Button if all above fails
    client_config = json.loads(GCP_CREDENTIALS_STR)
    flow = InstalledAppFlow.from_client_config(
        client_config, 
        scopes=ALL_SCOPES,
        redirect_uri="https://alakbarli-assistant.hf.space/"
    )
    auth_url, _ = flow.authorization_url(prompt='consent', access_type='offline')
    st.info("👋 Welcome! Please sign in to continue.")
    st.link_button("🔑 Login with Google", auth_url)
    return None


# --- CALENDAR OPERATIONS ---
def normalize_event_json(event_data):
    """Ensures JSON from AI or Google uses the same keys."""
    return {
        'id': event_data.get('id'),
        'title': event_data.get('title') or event_data.get('summary', 'Untitled'),
        'start_time': event_data.get('start_time') or event_data.get('start'),
        'end_time': event_data.get('end_time') or event_data.get('end'),
        'location': event_data.get('location', 'N/A'),
        'attendees': event_data.get('attendees', []),
        'description': event_data.get('description', '')
    }

def create_calendar_event(creds, event_json, user_timezone_str, create_meet=False):
    try:
        service = build('calendar', 'v3', credentials=creds)
        attendees_list = [{'email': e.strip()} for e in event_json.get('attendees', []) if '@' in e]

        event = {
            'summary': event_json.get('title', 'Untitled Event'),
            'location': event_json.get('location', ''),
            'description': event_json.get('description', ''),
            'start': {'dateTime': event_json['start_time'], 'timeZone': user_timezone_str},
            'end': {'dateTime': event_json['end_time'], 'timeZone': user_timezone_str},
            'attendees': attendees_list,
        }

        if create_meet:
            event['conferenceData'] = {
                'createRequest': {
                    'requestId': f"meet_{int(time.time())}",
                    'conferenceSolutionKey': {'type': 'hangoutsMeet'}
                }
            }

        res = service.events().insert(
            calendarId='primary',
            body=event,
            conferenceDataVersion=1
        ).execute()

        meet_link = res.get('hangoutLink', '')
        success_msg = f"✅ Event **'{res.get('summary')}'** created!"
        if meet_link:
            success_msg += f"\n\n🔗 **Google Meet Link:** {meet_link}"
        return success_msg
    except Exception as e:
        return f"❌ Error: {str(e)}"

def update_calendar_event(creds, event_id, updated_json, user_timezone_str, create_meet=False):
    try:
        service = build('calendar', 'v3', credentials=creds)
        attendees_list = [{'email': e.strip()} for e in updated_json.get('attendees', []) if '@' in e]

        body = {
            'summary': updated_json.get('title'),
            'location': updated_json.get('location'),
            'description': updated_json.get('description'),
            'start': {'dateTime': updated_json['start_time'], 'timeZone': user_timezone_str},
            'end': {'dateTime': updated_json['end_time'], 'timeZone': user_timezone_str},
            'attendees': attendees_list,
        }

        if create_meet:
            body['conferenceData'] = {
                'createRequest': {
                    'requestId': f"mod_meet_{int(time.time())}",
                    'conferenceSolutionKey': {'type': 'hangoutsMeet'}
                }
            }
        else:
            body['conferenceData'] = None

        res = service.events().patch(
            calendarId='primary',
            eventId=event_id,
            body=body,
            conferenceDataVersion=1
        ).execute()

        meet_link = res.get('hangoutLink', '')
        success_msg = "✅ Event updated successfully!"
        if meet_link and create_meet:
            success_msg += f"\n\n🔗 **New Google Meet Link:** {meet_link}"
        return success_msg
    except Exception as e:
        return f"❌ Update failed: {str(e)}"

def delete_calendar_event(creds, event_id):
    try:
        service = build('calendar', 'v3', credentials=creds)
        service.events().delete(calendarId='primary', eventId=event_id).execute()
        return "🗑️ Event deleted successfully!"
    except Exception as e:
        return f"❌ Deletion failed: {str(e)}"

def get_conflicting_events(creds, start_iso, end_iso, exclude_id=None):
    try:
        service = build('calendar', 'v3', credentials=creds)
        results = service.events().list(
            calendarId='primary',
            timeMin=start_iso,
            timeMax=end_iso,
            singleEvents=True,
            orderBy='startTime'
        ).execute()
        events = results.get('items', [])
        if exclude_id:
            events = [e for e in events if e.get('id') != exclude_id]
        return events
    except Exception as e:
        st.error(f"Error checking conflicts: {e}")
        return []

def fetch_upcoming_events(creds, start_date=None, end_date=None):
    try:
        service = build('calendar', 'v3', credentials=creds)
        time_min = start_date.isoformat() + 'Z' if start_date else datetime.datetime.utcnow().isoformat() + 'Z'
        time_max = end_date.isoformat() + 'Z' if end_date else None

        results = service.events().list(
            calendarId='primary',
            timeMin=time_min,
            timeMax=time_max,
            singleEvents=True,
            orderBy='startTime'
        ).execute()

        events = results.get('items', [])
        return [{
            'id': e['id'],
            'summary': e.get('summary', 'No Title'),
            'start': e['start'].get('dateTime', e['start'].get('date')),
            'end': e['end'].get('dateTime', e['end'].get('date')),
            'attendees': [a.get('email') for a in e.get('attendees', [])],
            'location': e.get('location', 'N/A')
        } for e in events]
    except Exception as e:
        st.error(f"Error fetching: {e}")
        return []

# --- GMAIL OPERATIONS ---
# def fetch_unread_emails(service):
#     structured_emails = []
#     try:
#         results = service.users().messages().list(userId="me", maxResults=5, q="is:unread").execute()
#         for m_info in results.get("messages", []):
#             msg = service.users().messages().get(userId='me', id=m_info['id']).execute()
#             headers = {h['name']: h['value'] for h in msg['payload']['headers']}

#             body = ""
#             payload = msg.get('payload', {})

#             def get_body(parts):
#                 for part in parts:
#                     mimeType = part.get('mimeType')
#                     data = part.get('body', {}).get('data')
#                     if mimeType == 'text/plain' and data:
#                         return base64.urlsafe_b64decode(data).decode('utf-8')
#                     elif 'parts' in part:
#                         result = get_body(part['parts'])
#                         if result: return result
#                 return ""

#             if 'parts' in payload:
#                 body = get_body(payload['parts'])
#             else:
#                 data = payload.get('body', {}).get('data')
#                 if data:
#                     body = base64.urlsafe_b64decode(data).decode('utf-8')

#             structured_emails.append({
#                 'id': msg['id'],
#                 'time': time.strftime('%Y-%m-%d %I:%M %p', time.localtime(int(msg['internalDate'])/1000)),
#                 'from': headers.get('From', 'N/A'),
#                 'subject': headers.get('Subject', 'N/A'),
#                 'body': body if body else "No readable text found",
#                 'summary': ''
#             })
#         return structured_emails
#     except Exception as e:
#         st.error(f"Error fetching emails: {e}")
#         return []



def fetch_emails_by_date(service, target_date, max_results=50):
    structured_emails = []
    # Format dates for Gmail query (YYYY/MM/DD)
    after_date = target_date.strftime('%Y/%m/%d')
    before_date = (target_date + timedelta(days=1)).strftime('%Y/%m/%d')
    
    # Query: Emails received on the specific day
    query = f"after:{after_date} before:{before_date}"
    
    try:
        results = service.users().messages().list(userId="me", maxResults=max_results, q=query).execute()
        messages = results.get("messages", [])
        
        if not messages:
            return []

        for m_info in messages:
            msg = service.users().messages().get(userId='me', id=m_info['id']).execute()
            headers = {h['name']: h['value'] for h in msg['payload']['headers']}
            
            body = ""
            payload = msg.get('payload', {})

            def get_body(parts):
                for part in parts:
                    mimeType = part.get('mimeType')
                    data = part.get('body', {}).get('data')
                    if mimeType == 'text/plain' and data:
                        return base64.urlsafe_b64decode(data).decode('utf-8')
                    elif 'parts' in part:
                        result = get_body(part['parts'])
                        if result: return result
                return ""

            if 'parts' in payload:
                body = get_body(payload['parts'])
            else:
                data = payload.get('body', {}).get('data')
                if data:
                    body = base64.urlsafe_b64decode(data).decode('utf-8')

            structured_emails.append({
                'id': msg['id'],
                'time': datetime.datetime.fromtimestamp(int(msg['internalDate'])/1000).strftime('%I:%M %p'),
                'from': headers.get('From', 'N/A'),
                'subject': headers.get('Subject', 'N/A'),
                'body': body if body else "No readable text found",
                'summary': ''
            })
        return structured_emails
    except Exception as e:
        st.error(f"Error fetching emails for {after_date}: {e}")
        return []

        

# --- AI AGENT PARSERS ---
def parse_schedule_prompt(prompt, current_date, user_timezone_str):
    tz = pytz.timezone(user_timezone_str)
    now_aware = datetime.datetime.now(tz)
    system_prompt = f"You are a scheduling assistant. Return ONLY JSON. Date: {current_date}. TZ: {user_timezone_str}. Keys: title, start_time, end_time, location, attendees, description."
    try:
        resp = client.chat.completions.create(
            model=CHATGPT_MODEL,
            messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}],
            temperature=0.1
        )
        return resp.choices[0].message.content.strip().strip("```json").strip("```")
    except Exception as e: 
        return json.dumps({"error": str(e)})

def parse_modification_prompt(original_event, mod_request, current_date, user_timezone_str):
    system_prompt = f"Update the following JSON based on user request. Maintain ISO 8601. Date: {current_date}. Return ONLY JSON."
    user_input = f"EXISTING: {json.dumps(original_event)}\nREQUEST: {mod_request}"
    try:
        resp = client.chat.completions.create(
            model=CHATGPT_MODEL,
            messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": user_input}],
            temperature=0.1
        )
        return resp.choices[0].message.content.strip().strip("```json").strip("```")
    except Exception as e: 
        return json.dumps({"error": str(e)})

# def summarize_email(email_text, email_from, email_subject):
#     if not client:
#         return "• ⚠️ Summary failed: LLM client is not configured."
    
#     match = re.search(r'^(.*?) <', email_from)
#     sender_name = match.group(1).strip() if match else email_from.split('<')[0].strip()
#     system_prompt = "You are an AI assistant. Summarize into one concise bullet point (<15 words) starting with the sender's name."
#     user_input = f"Sender: {sender_name}\nSubject: {email_subject}\nBody: {email_text}"

#     try:
#         response = client.chat.completions.create(
#             model=CHATGPT_MODEL, 
#             messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": user_input}], 
#             temperature=0.3
#         )
#         summary = response.choices[0].message.content.strip()
#         return summary if summary.startswith('•') else '• ' + summary
#     except Exception as e:
#         return f"• Error during summarization: {str(e)[:50]}..."


def summarize_email(email_text, email_from, email_subject):
    if not client:
        return "• ⚠️ Summary failed: LLM client is not configured."
    
    # Ensure this logic runs (not skipped by a return above)
    match = re.search(r'^(.*?) <', email_from)
    sender_name = match.group(1).strip() if match else email_from.split('<')[0].strip()
    
    system_prompt = (
        "You are an executive assistant. Summarize the email into ONE line following this exact pattern: "
        "• [Sender Name] [Action Taken/Topic] — [Required Action]. "
        "Keep it under 15 words. If no action is needed, use 'for information only'."
    )
    
    user_input = f"Sender: {sender_name}\nSubject: {email_subject}\nBody: {email_text}"
    
    try:
        response = client.chat.completions.create(
            model=CHATGPT_MODEL, 
            messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": user_input}], 
            temperature=0.1
        )
        summary = response.choices[0].message.content.strip()
        if not summary.startswith('•'):
            summary = '• ' + summary
        return summary
    except Exception as e:
        return f"• Error during summarization: {str(e)[:50]}..."


def draft_specific_notification(event_json, change_summary, is_new_invite=False):
    if not client: return "LLM client not configured."
    system_prompt = "You are a professional corporate coordinator. Draft a change-first email."
    
    if is_new_invite:
        user_content = f"Write a warm invitation for a NEW attendee to: {event_json.get('title')}."
    else:
        user_content = f"Event '{event_json.get('title')}' updated. Changes: {change_summary}."

    try:
        res = client.chat.completions.create(
            model=CHATGPT_MODEL,
            messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": user_content}]
        )
        return res.choices[0].message.content.strip()
    except Exception as e:
        return f"Error drafting: {e}"

# --- NAVIGATION HELPER ---
def render_window_controls():
    if 'calendar_window_start' not in st.session_state:
        st.session_state.calendar_window_start = datetime.datetime.now(datetime.timezone.utc)
        
    start_dt = st.session_state.calendar_window_start
    end_dt = start_dt + timedelta(days=90)
    st.write(f"📅 Window: **{start_dt.strftime('%b %Y')}** — **{end_dt.strftime('%b %Y')}**")

    col_prev, col_next = st.columns(2)
    if col_next.button("➡️ Next 3 Months", use_container_width=True):
        st.session_state.calendar_window_start = end_dt
        st.rerun()
    if col_prev.button("⬅️ Back to Today", use_container_width=True):
        st.session_state.calendar_window_start = datetime.datetime.utcnow()
        st.rerun()
    return start_dt, end_dt


# --- HELPERS FOR MODIFICATION ---
def get_change_summary(old, new):
    """Compares two event objects and returns a string describing changes."""
    changes = []
    if old.get('title') != new.get('title'): 
        changes.append(f"Title changed to '{new['title']}'")
    if old.get('start_time') != new.get('start_time'): 
        changes.append("The event time was updated")
    if old.get('location') != new.get('location'): 
        changes.append(f"Location moved to {new['location']}")
    return ", ".join(changes) if changes else "General details were updated"

def find_email_by_sender(email_list, name_snippet):
    """Finds an email object by searching the sender's name."""
    for e in email_list:
        if name_snippet.lower() in e['from'].lower():
            return e
    return None

def draft_reply_email(subject, body, summary):
    """Drafts a professional reply based on the email content."""
    system_prompt = "You are a helpful assistant. Draft a professional, brief reply to this email."
    user_input = f"Subject: {subject}\nBody: {body}\nSummary: {summary}"
    try:
        res = client.chat.completions.create(
            model=CHATGPT_MODEL,
            messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": user_input}]
        )
        return res.choices[0].message.content.strip()
    except:
        return "Could not generate draft."




# --- MAIN APP ---

def main():
    st.set_page_config(page_title="AI Gmail & Calendar Assistant", layout="wide")
    st.title("🤖 AI Gmail & Calendar Assistant")

    # 1. Initialize Sidebar UI
    if 'user_timezone' not in st.session_state:
        st.session_state.user_timezone = "UTC"
    
    st.sidebar.selectbox("🌎 Timezone", options=TIMEZONES, key='user_timezone')

    # 2. Corrected Logout/Switch Account Logic
    # This must be outside the 'creds' check to allow users to escape a session
    if st.sidebar.button("🔄 Switch Account / Logout"):
        for key in list(st.session_state.keys()):
            del st.session_state[key]
        cookie_manager.delete("google_creds_token")
        st.rerun()

    
    
    # --- Authentication ---
    # --- UPDATED AUTHENTICATION LOGIC ---
    creds = authenticate_service()
    
    if not creds:
        # We stop the app here. The authenticate_service function 
        # is already displaying the "Login with Google" button.
        st.stop() 

    # --- Everything below this line runs ONLY after successful login ---
    st.sidebar.success("✅ Connected to Google")

    if 'email_selected_date' not in st.session_state or st.session_state.email_selected_date is None:
        st.session_state.email_selected_date = date.today()

    if 'calendar_window_start' not in st.session_state or st.session_state.calendar_window_start is None:
        st.session_state.calendar_window_start = datetime.datetime.utcnow()

    # --- Session State Initialization ---
    state_keys = [
        'email_data', 'calendar_data', 'summarized', 'calendar_action', 
        'parsed_schedule_json', 'manual_schedule_json', 'calendar_window_start', 'email_selected_date'
    ]
    for key in state_keys:
        if key not in st.session_state:
            st.session_state[key] = None

    if st.session_state.email_selected_date is None:
        st.session_state.email_selected_date = date.today()

    # --- Main Tabs ---
    tab1, tab2 = st.tabs(["📧 Email Inbox Assistant", "🗓️ Calendar Agenda"])

    # ---------------------------------------------------------
    # TAB 1: EMAIL INBOX ASSISTANT
    # ---------------------------------------------------------
    # with tab1:
    #     st.header("1. 📧 Agent 1: Email Fetching")
        
    #     # --- Step 1: Fetch Raw Data ---
    #     if st.button("▶️ Step 1: Fetch Raw Email Data (Gmail)"):
    #         gmail_service = build("gmail", "v1", credentials=creds)
    #         with st.spinner("Connecting to Gmail and fetching structured email data..."):
    #             st.session_state.email_data = fetch_unread_emails(gmail_service) 
    #             st.session_state.summarized = False 
        
    #     # --- Display Agent 1 Results ---
    #     if st.session_state.email_data:
    #         st.subheader("Raw Data Extracted from Gmail")
    #         raw_table_data = [
    #             {
    #                 'Date & Time': e['time'], 
    #                 'From': e['from'], 
    #                 'Subject': e['subject'], 
    #                 'Body': e['body'],
    #                 'Attachments': 'Yes' if e.get('attachments') else 'No'
    #             }
    #             for e in st.session_state.email_data
    #         ]
            
    #         st.dataframe(
    #             raw_table_data, 
    #             use_container_width=True,
    #             column_order=('Date & Time', 'From', 'Subject', 'Body', 'Attachments'), 
    #             column_config={
    #                 "Subject": st.column_config.Column(width="large"),
    #                 "Body": st.column_config.TextColumn(
    #                     "Email Body (Scroll/Click)", 
    #                     help="The full text of the email. Click to read all content.",
    #                     width="medium",
    #                     max_chars=200, 
    #                 )
    #             }
    #         )
            
    #         st.divider()

    with tab1:
        st.header("1. 📧 Agent 1: Email Fetching")
        
        # --- NEW: Date Navigation UI (Similar to your Calendar navigation) ---
        col_prev, col_date, col_next = st.columns([1, 2, 1])
        
        if col_prev.button("⬅️ Previous Day", use_container_width=True):
            st.session_state.email_selected_date -= timedelta(days=1)
            st.rerun()
            
        # Allow manual date selection too
        st.session_state.email_selected_date = col_date.date_input(
            "Target Date", 
            value=st.session_state.email_selected_date,
            label_visibility="collapsed"
        )
            
        if col_next.button("Next Day ➡️", use_container_width=True):
            st.session_state.email_selected_date += timedelta(days=1)
            st.rerun()

        st.info(f"📅 Currently targeting emails for: **{st.session_state.email_selected_date.strftime('%Y-%m-%d')}**")
        
        # --- Modified Step 1: Use the new fetch_emails_by_date ---
        if st.button(f"▶️ Step 1: Fetch Emails for {st.session_state.email_selected_date}", use_container_width=True):
            gmail_service = build("gmail", "v1", credentials=creds)
            with st.spinner(f"Searching Gmail for {st.session_state.email_selected_date}..."):
                # Calling your new function instead of fetch_unread_emails
                st.session_state.email_data = fetch_emails_by_date(
                    gmail_service, 
                    st.session_state.email_selected_date,
                    max_results=50 # Increased for high-volume handling
                ) 
                st.session_state.summarized = False 
        
        # --- Display Agent 1 Results (No changes to your view logic here) ---
        if st.session_state.email_data:
            st.subheader(f"Raw Data Extracted ({len(st.session_state.email_data)} emails)")
            raw_table_data = [
                {
                    'Date & Time': e['time'], 
                    'From': e['from'], 
                    'Subject': e['subject'], 
                    'Body': e['body'],
                    'Attachments': 'Yes' if e.get('attachments') else 'No'
                }
                for e in st.session_state.email_data
            ]
            
            st.dataframe(
                raw_table_data, 
                use_container_width=True,
                column_order=('Date & Time', 'From', 'Subject', 'Body', 'Attachments'), 
                column_config={
                    "Subject": st.column_config.Column(width="large"),
                    "Body": st.column_config.TextColumn(
                        "Email Body (Scroll/Click)", 
                        width="medium",
                        max_chars=200, 
                    )
                }
            )
            
            st.divider()
            
            # --- Step 2: Agent 2 (Summarization) Trigger ---
            # st.header("2. 🎯 Agent 2: Summarize Actions")
            # if st.button("▶️ Step 2: Generate Actionable Summaries (LLM Call)"):
            #     with st.spinner("Generating actionable summaries using the LLM..."):
            #         for email in st.session_state.email_data:
            #             if not email.get('summary') or 'Error' in email.get('summary', ''): 
            #                 email['summary'] = summarize_email(
            #                     email['body'], 
            #                     email['from'], 
            #                     email['subject']
            #                 )
            #         st.session_state.summarized = True
                
            # # --- Display Agent 2 Results & Step 3 ---
            # if st.session_state.summarized:
            #     st.success("✅ Summaries generated!")
            #     st.subheader("Actionable Summary List")
            #     summary_table_data = [
            #         {'Date & Time': e['time'], 'Summary': e['summary']}
            #         for e in st.session_state.email_data
            #     ]
            #     st.dataframe(
            #         summary_table_data, 
            #         use_container_width=True, 
            #         column_config={"Summary": st.column_config.Column(width="large")}
            #     )
                
            #     st.divider()




            
            # --- Step 2: Agent 2 (Summarization) Trigger ---
            # st.header("2. 🎯 Agent 2: Summarize Actions")
            # if st.button("▶️ Step 2: Generate Actionable Summaries (LLM Call)"):
            #     if not st.session_state.email_data:
            #         st.warning("Please fetch emails first!")
            #     else:
            #         # NEW: Add a progress bar for high-volume emails
            #         progress_bar = st.progress(0)
            #         status_text = st.empty()
                    
            #         num_emails = len(st.session_state.email_data)
                    
            #         for i, email in enumerate(st.session_state.email_data):
            #             status_text.text(f"Processing email {i+1} of {num_emails}...")
                        
            #             # Only summarize if not already done or if it failed previously
            #             if not email.get('summary') or 'Error' in email.get('summary', ''): 
            #                 email['summary'] = summarize_email(
            #                     email['body'], 
            #                     email['from'], 
            #                     email['subject']
            #                 )
                        
            #             # Update progress
            #             progress_bar.progress((i + 1) / num_emails)
                    
            #         st.session_state.summarized = True
            #         status_text.empty()
            #         progress_bar.empty()
            #         st.success(f"✅ Successfully processed {num_emails} actionable summaries!")
    
            # # --- Display Agent 2 Results (Table) ---
            # if st.session_state.summarized:
            #     summary_table_data = [
            #         {'Time': e['time'], 'Actionable Summary': e['summary']}
            #         for e in st.session_state.email_data
            #     ]
            #     st.dataframe(
            #         summary_table_data, 
            #         use_container_width=True, 
            #         column_config={"Actionable Summary": st.column_config.Column(width="large")}
            #     )
            #     st.divider()


            # --- Step 2: Agent 2 (Summarization) Trigger ---
            st.header("2. 🎯 Agent 2: Summarize Actions")

            if st.session_state.email_data:
                # 1. Calculate status
                pending_emails = [e for e in st.session_state.email_data if not e.get('summary') or 'Error' in e.get('summary', '')]
                total_ready = len(st.session_state.email_data) - len(pending_emails)
                
                # 2. Show the Action Button if there's work to do
                if pending_emails:
                    btn_label = f"▶️ Summarize Next 5 Emails ({total_ready}/{len(st.session_state.email_data)} done)"
                    if st.button(btn_label):
                        batch_to_process = pending_emails[:5]
                        progress_bar = st.progress(0)
                        
                        for i, email in enumerate(batch_to_process):
                            email['summary'] = summarize_email(email['body'], email['from'], email['subject'])
                            time.sleep(1.5) 
                            progress_bar.progress((i + 1) / len(batch_to_process))
                        
                        st.session_state.summarized = True
                        st.rerun() 
                else:
                    st.success("✅ All emails for this date have been summarized!")
            
            if st.session_state.get('summarized'):
                # Look for any email that has a valid summary string
                summarized_emails = [e for e in st.session_state.email_data if e.get('summary') and "•" in e['summary']]
                
                if summarized_emails:
                    st.subheader("📋 Actionable Summary List")
                    summary_table_data = [
                        {'Time': e['time'], 'Actionable Summary': e['summary']}
                        for e in summarized_emails
                    ]
                    st.dataframe(
                        summary_table_data, 
                        use_container_width=True, 
                        column_config={"Actionable Summary": st.column_config.Column(width="large")}
                    )
                    st.divider()


            
            else:
                st.info("Please fetch emails in Step 1 first.")
            
            # --- Display Agent 2 Results (Table) ---
            # Filter the display to only show emails that HAVE a summary
            if st.session_state.summarized:
                summarized_only = [e for e in st.session_state.email_data if e.get('summary') and 'Error' not in e['summary']]
                
                if summarized_only:
                    st.subheader("Actionable Summary List")
                    summary_table_data = [
                        {'Time': e['time'], 'Actionable Summary': e['summary']}
                        for e in summarized_only
                    ]
                    st.dataframe(
                        summary_table_data, 
                        use_container_width=True, 
                        column_config={"Actionable Summary": st.column_config.Column(width="large")}
                    )
                    st.divider()
                
                # --- Step 3: Agent 3 (Reply Drafting) ---
                st.header("3. ✍️ Agent 3: Reply Drafting")
                st.markdown("Type a command (e.g., `prepare reply for Alex`) to draft an email response.")

                user_prompt = st.text_input(
                    "Enter your drafting command:", 
                    placeholder="e.g., Prepare a quick reply for Alex"
                )
                
                if user_prompt and st.button("▶️ Step 3: Draft Reply"):
                    if "reply for" in user_prompt.lower():
                        sender_name = user_prompt.lower().split("reply for")[-1].strip()
                        target_email = find_email_by_sender(st.session_state.email_data, sender_name)

                        if target_email:
                            with st.spinner(f"Finding email from {sender_name} and drafting reply..."):
                                reply_draft = draft_reply_email(
                                    target_email['subject'], 
                                    target_email['body'], 
                                    target_email['summary']
                                )
                                
                            st.success(f"✅ Reply drafted for **{target_email['from'].split('<')[0].strip()}**.")
                            st.subheader("Copy & Paste Reply Draft")
                            st.markdown(f"> **Context Used:** {target_email['summary']}") 
                            st.code(reply_draft, language='email')
                        else:
                            st.error(f"❌ Could not find an email from '{sender_name}' in the fetched list.")
                    else:
                        st.warning("Please phrase your command using 'prepare reply for [Sender Name]'.")
        else:
            st.info("Click 'Step 1: Fetch Raw Email Data' to begin processing your unread emails.")

    # ---------------------------------------------------------
    # TAB 2: CALENDAR MANAGEMENT
    # ---------------------------------------------------------
    with tab2:
        st.header("🗓️ Calendar Management")
        c1, c2, c3, c4 = st.columns(4)
        
        if c1.button("List Events", use_container_width=True): 
            st.session_state.calendar_action = "list"
        if c2.button("Schedule New", use_container_width=True): 
            st.session_state.calendar_action = "schedule"
        if c3.button("Modify Existing", use_container_width=True): 
            st.session_state.calendar_action = "modify"
        if c4.button("Delete Event", use_container_width=True): 
            st.session_state.calendar_action = "delete"

        st.divider()

        # --- LIST LOGIC ---
        if st.session_state.calendar_action == "list":
            start_dt, end_dt = render_window_controls()
            st.session_state.calendar_data = fetch_upcoming_events(creds, start_date=start_dt, end_date=end_dt)
            if st.session_state.calendar_data:
                st.table(st.session_state.calendar_data)
            else:
                st.info("No events in this window.")



        # elif st.session_state.calendar_action == "schedule":
        #     prompt = st.text_input("Plan details:", "Meeting tomorrow at 10am")
            
        #     meet_keywords = ["online", "meeting link", "google meet", "video call", "zoom", "virtual"]
        #     wants_meet = any(word in prompt.lower() for word in meet_keywords)

        #     if st.button("Parse Plan"):
        #         # 1. Get the raw text from the AI
        #         raw_ai_response = parse_schedule_prompt(
        #             prompt, 
        #             date.today(), 
        #             st.session_state.user_timezone
        #         )
                
        #         try:
        #             # 2. Convert text to a temporary dictionary
        #             temp_dict = json.loads(raw_ai_response)
                    
        #             # 3. Use your NORMALIZE function to fix mismatched keys (like 'start' vs 'start_time')
        #             clean_data = normalize_event_json(temp_dict)
                    
        #             # 4. Save the CLEANED version back to session state
        #             st.session_state.parsed_schedule_json = json.dumps(clean_data, indent=2)
        #         except Exception as e:
        #             st.error(f"Failed to parse AI response: {e}")
            
        #     # --- Display and Process Cleaned Data ---
        #     if st.session_state.parsed_schedule_json:
        #         # Now we are 100% sure 'start_time' exists because of normalize_event_json
        #         event_data = json.loads(st.session_state.parsed_schedule_json)
                
        #         with st.spinner("Checking for calendar conflicts..."):
        #             # This used to crash; now it's safe!
        #             conflicts = get_conflicting_events(
        #                 creds, 
        #                 event_data['start_time'], 
        #                 event_data['end_time']
        #             )
                
        #         if conflicts:
        #             st.warning("⚠️ **Schedule Conflict Detected!**")
        #             for cf in conflicts:
        #                 st.info(f"**Conflict:** {cf.get('summary', 'No Title')} ({cf['start'].get('dateTime', 'All Day')})")
        #         else:
        #             st.success("✅ This time slot is free!")
                
        #         create_meet_toggle = st.checkbox("🌐 Create Google Meet link", value=wants_meet)
                
        #         # Show the cleaned JSON in the text area
        #         edited = st.text_area("Review JSON Details", value=st.session_state.parsed_schedule_json, height=150)
                
        #         if st.button("Add to Calendar"):
        #             result = create_calendar_event(
        #                 creds, 
        #                 json.loads(edited), 
        #                 st.session_state.user_timezone,
        #                 create_meet=create_meet_toggle 
        #             )
        #             st.success(result)
        #             # Optional: Clear the state so the form resets
        #             st.session_state.parsed_schedule_json = None

        elif st.session_state.calendar_action == "schedule":
            prompt = st.text_input("Plan details:", placeholder="e.g., Meeting tomorrow at 10am")
            
            meet_keywords = ["online", "meeting link", "google meet", "video call", "zoom", "virtual"]
            wants_meet = any(word in prompt.lower() for word in meet_keywords)

            if st.button("1. Parse Plan", use_container_width=True):
                raw_ai_response = parse_schedule_prompt(
                    prompt, 
                    date.today(), 
                    st.session_state.user_timezone
                )
                try:
                    temp_dict = json.loads(raw_ai_response)
                    clean_data = normalize_event_json(temp_dict)
                    st.session_state.parsed_schedule_json = json.dumps(clean_data, indent=2)
                    # Clear previous conflict results when parsing a new plan
                    st.session_state.conflict_results = None 
                except Exception as e:
                    st.error(f"Failed to parse AI response: {e}")
            
            if st.session_state.parsed_schedule_json:
                st.subheader("📝 Review & Verify")
                # Let user see/edit the JSON first
                edited_json_str = st.text_area("Event JSON Details", value=st.session_state.parsed_schedule_json, height=150)
                event_data = json.loads(edited_json_str)

                col1, col2 = st.columns(2)
                
                # --- MANUAL AVAILABILITY CHECK ---
                if col1.button("🔍 Check Availability", use_container_width=True):
                    with st.spinner("Checking exact time slot..."):
                        # Ensure time strings are in full ISO format (YYYY-MM-DDTHH:MM:SSZ)
                        start_iso = event_data['start_time']
                        end_iso = event_data['end_time']
                        
                        conflicts = get_conflicting_events(creds, start_iso, end_iso)
                        st.session_state.conflict_results = conflicts

                # --- ADD TO CALENDAR ---
                create_meet_toggle = st.checkbox("🌐 Create Google Meet link", value=wants_meet)
                if col2.button("🚀 Add to Calendar", use_container_width=True, type="primary"):
                    result = create_calendar_event(
                        creds, 
                        event_data, 
                        st.session_state.user_timezone,
                        create_meet=create_meet_toggle 
                    )
                    st.success(result)
                    st.session_state.parsed_schedule_json = None
                    st.session_state.conflict_results = None

                # --- DISPLAY CONFLICTS ONLY IF BUTTON WAS PRESSED ---
                if st.session_state.get('conflict_results') is not None:
                    if st.session_state.conflict_results:
                        st.warning(f"⚠️ **Conflict Found:** You have {len(st.session_state.conflict_results)} event(s) during this exact time.")
                        for cf in st.session_state.conflict_results:
                            st.info(f"**Blocked by:** {cf.get('summary')} \n\n Time: {cf['start'].get('dateTime')}")
                    else:
                        st.success("✅ This exact time slot is free!")




        elif st.session_state.calendar_action == "modify":
            st.subheader("📝 Modify Events")
            start_dt, end_dt = render_window_controls()
            
            if st.session_state.calendar_data is None:
                st.session_state.calendar_data = fetch_upcoming_events(creds, start_date=start_dt, end_date=end_dt)
            
            if st.session_state.calendar_data:
                df = st.dataframe(
                    st.session_state.calendar_data, 
                    on_select="rerun", 
                    selection_mode="single-row", 
                    use_container_width=True,
                    key="modify_df"
                )

                if df.selection.rows:
                    selected = st.session_state.calendar_data[df.selection.rows[0]]
                    st.markdown(f"### Selected: {selected['summary']}")
                    
                    mod_input = st.text_input("What should change?", placeholder="e.g., change from online to offline in IT Office")
                    
                    # --- AI Intent Detection for Location Type ---
                    is_hybrid = any(word in mod_input.lower() for word in ["hybrid", "both", "keep location and add meet"])
                    is_offline = any(word in mod_input.lower() for word in ["offline", "in person", "remove meeting link", "at the office"])
                    wants_online = any(word in mod_input.lower() for word in ["online", "virtual", "google meet", "add link"]) and not is_offline

                    if st.button("1. Preview AI Changes"):
                        st.session_state.manual_schedule_json = parse_modification_prompt(
                            selected, mod_input, date.today(), st.session_state.user_timezone
                        )

                    if st.session_state.manual_schedule_json:
                        raw_ai_json = json.loads(st.session_state.manual_schedule_json)
                        
                        # Logic to handle location field based on detected intent
                        if wants_online and not is_hybrid:
                            raw_ai_json['location'] = "Google Meet (Online)"
                        elif is_offline:
                            # Ensure the AI didn't leave "Google Meet" in the location
                            if "google meet" in raw_ai_json.get('location', '').lower():
                                raw_ai_json['location'] = "IT Service Office" # Default or AI-parsed location

                        new_event_preview = normalize_event_json(raw_ai_json)
                        
                        with st.spinner("Checking availability..."):
                            conflicts = get_conflicting_events(
                                creds, 
                                new_event_preview['start_time'], 
                                new_event_preview['end_time'],
                                exclude_id=selected['id']
                            )
                        
                        if conflicts:
                            st.error("🚫 **Conflict:** The new time slot is occupied.")
                            disable_update = True
                        else:
                            st.success("✅ The time slot is available.")
                            disable_update = False

                        # Auto-set the Meet checkbox based on AI intent
                        create_meet_toggle = st.checkbox("🌐 Enable Google Meet Link", value=(wants_online or is_hybrid))
                        
                        final_json_str = st.text_area("Review JSON", value=json.dumps(raw_ai_json, indent=2))
                        
                      

                        if st.button("2. Confirm & Update Calendar", disabled=disable_update):
                            old_event = normalize_event_json(selected)
                            new_event = normalize_event_json(json.loads(final_json_str))
                            
                            # A. Identify what changed
                            change_summary = get_change_summary(old_event, new_event)
                            
                            # B. Perform the Calendar Update
                            result = update_calendar_event(
                                creds, 
                                old_event['id'], 
                                new_event, 
                                st.session_state.user_timezone,
                                create_meet=create_meet_toggle
                            )
                            
                            st.success(result)

                            # C. Agent 4: Draft Notification if attendees exist
                            if new_event.get('attendees'):
                                with st.spinner("Agent 4: Drafting notification for attendees..."):
                                    notification_draft = draft_specific_notification(
                                        new_event, 
                                        change_summary, 
                                        is_new_invite=False
                                    )
                                    
                                st.subheader("📧 Notification Draft for Attendees")
                                st.info(f"**Changes detected:** {change_summary}")
                                st.code(notification_draft, language='email')
                                st.caption("Copy this draft to send a manual update to: " + ", ".join(new_event['attendees']))
                            
                            # Reset session state
                            st.session_state.manual_schedule_json = None



        # --- DELETE LOGIC ---
        elif st.session_state.calendar_action == "delete":
            st.subheader("🗑️ Delete Events")
            start_dt, end_dt = render_window_controls()
            st.session_state.calendar_data = fetch_upcoming_events(creds, start_date=start_dt, end_date=end_dt)
            
            if st.session_state.calendar_data:
                df_del = st.dataframe(
                    st.session_state.calendar_data, 
                    on_select="rerun", 
                    selection_mode="single-row", 
                    use_container_width=True
                )
                if df_del.selection.rows:
                    event_to_del = st.session_state.calendar_data[df_del.selection.rows[0]]
                    st.warning(f"Are you sure you want to delete **'{event_to_del['summary']}'**?")
                    col_a, col_b = st.columns(2)
                    if col_a.button("✅ Confirm Delete", use_container_width=True):
                        st.success(delete_calendar_event(creds, event_to_del['id']))
                        st.session_state.calendar_data = None
                        st.rerun()
                    if col_b.button("❌ Cancel", use_container_width=True):
                        st.info("Action cancelled.")

if __name__ == "__main__":
    main()