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()