app.py

import streamlit as st
import pandas as pd
from datetime import datetime
import database as db 
import sqlite3
import llm_handler

# --- Page Configuration ---
st.set_page_config(page_title="AI Task Manager", layout="centered", initial_sidebar_state="collapsed")

# --- Helper Functions ---
def initialize_session_state():
    """Initializes session state variables."""
    if "user_name" not in st.session_state: st.session_state.user_name = ""
    if "user_email" not in st.session_state: st.session_state.user_email = ""
    if "logged_in" not in st.session_state: st.session_state.logged_in = False
    if "llm" not in st.session_state:
        try:
            st.session_state.llm = llm_handler.get_llm()
        except ValueError as e:
            st.error(f"LLM Initialization Error: {e}. Please ensure GOOGLE_API_KEY is set in .env")
            st.stop()
    
    if "last_interacted_task_details" not in st.session_state:
        st.session_state.last_interacted_task_details = None
    if "chat_history_for_context" not in st.session_state: 
        st.session_state.chat_history_for_context = []

# --- Initialize ---
initialize_session_state()
db.create_db_and_table() 
if "table_info" not in st.session_state or st.session_state.get("reload_table_info", False):
    st.session_state.table_info = db.get_db_info()
    st.session_state.reload_table_info = False


# --- UI Sections ---

if not st.session_state.logged_in:
    st.title("AI Task Manager Login")
    with st.form("login_form"):
        name = st.text_input("Your Name", value=st.session_state.get("user_name_input", ""))
        email = st.text_input("Your Email", value=st.session_state.get("user_email_input", ""))
        login_button = st.form_submit_button("Login")

        if login_button:
            error_messages = []
            if not name:
                error_messages.append("Please enter your name.")
            if not email:
                error_messages.append("Please enter your email.")
            elif not email.strip().lower().endswith(".com"):
                error_messages.append("Invalid email format.")

            if not error_messages:
                st.session_state.user_name = name
                st.session_state.user_email = email
                st.session_state.logged_in = True
                st.session_state.user_name_input = name 
                st.session_state.user_email_input = email
                st.session_state.last_interacted_task_details = None 
                st.session_state.reload_table_info = True
                st.rerun()
            else:
                for msg in error_messages:
                    st.error(msg)
else:
    st.sidebar.header("User Info")
    st.sidebar.write(f"👤 **Name:** {st.session_state.user_name}")
    st.sidebar.write(f"📧 **Email:** {st.session_state.user_email}")
    if st.sidebar.button("Logout"):
        st.session_state.logged_in = False
        st.session_state.user_name = ""
        st.session_state.user_email = ""
        st.session_state.last_interacted_task_details = None
        st.rerun()

    st.title(f"📝 AI Task Manager for {st.session_state.user_name}")
    st.markdown("Enter your task command (e.g., 'I have a meeting at 2pm on next Friday' , 'Show pending tasks').")

    if st.session_state.last_interacted_task_details:
        task_ctx = st.session_state.last_interacted_task_details
        
        # Prepare task context display message
        task_name = task_ctx.get('task_name', 'N/A')
        task_id = task_ctx.get('id', 'N/A')
        due_date = task_ctx.get('due_date', '')
        due_time = task_ctx.get('due_time', '')
        ctx_display = f"Last task: '{task_name}' (ID: {task_id})"
        if due_date:
            ctx_display += f" - Due: {due_date}"
        if due_time:
            ctx_display += f" at {due_time}"
        st.caption(ctx_display)


    user_input_query = st.text_input("Your command:", key="task_input", placeholder="e.g., Add task 'Submit report' due next Friday 3pm category Work")

    if st.button("Process Command", type="primary"):
        if not user_input_query:
            st.warning("Please enter a command.")
        else:
            llm = st.session_state.llm
            table_info = st.session_state.table_info
            
            context_str = "None"
            if st.session_state.last_interacted_task_details:
                details = st.session_state.last_interacted_task_details
                context_list = []
                if details.get('id'): context_list.append(f"id: {details['id']}")
                if details.get('task_name'): context_list.append(f"name: '{details['task_name']}'")
                if details.get('due_date'): context_list.append(f"due_date: {details['due_date']}")
                if details.get('due_time'): context_list.append(f"due_time: {details['due_time']}")
                context_str = f"Task context - {', '.join(context_list)}"


            with st.spinner("🤖 Thinking and generating SQL..."):
                try:
                    generated_sql = llm_handler.generate_sql_query(
                        llm,
                        user_input_query,
                        table_info,
                        st.session_state.user_name,
                        st.session_state.user_email,
                        context_str
                    )
                except Exception as e:
                    st.error(f"Error generating SQL: {e}")
                    generated_sql = None

            if generated_sql:
                st.write("⚙️ **Generated SQL Query:**")
                st.code(generated_sql, language="sql")

                is_select_query = generated_sql.strip().upper().startswith("SELECT")
                is_insert_query = generated_sql.strip().upper().startswith("INSERT")

                with st.spinner("💾 Executing query..."):
                    try:
                        summary_context_for_llm = ""
                        if is_select_query:
                            data, columns = db.execute_select_query(generated_sql)
                            if data:
                                df = pd.DataFrame(data, columns=columns)
                                major_display_columns = ['id', 'task_name', 'status', 'category', 'due_date', 'due_time', 'created_at']
                                display_cols_in_df = [col for col in major_display_columns if col in df.columns]
                                
                                if not display_cols_in_df and df.columns.any():
                                    display_cols_in_df = df.columns.tolist()

                                if display_cols_in_df:
                                    st.dataframe(df[display_cols_in_df], use_container_width=True)
                                else:
                                    st.info("The query ran but returned no columns to display.")
                                
                                summary_context_for_llm = f"Retrieved {len(data)} task(s)."
                                if len(data) == 1: 
                                    st.session_state.last_interacted_task_details = dict(zip(columns, data[0]))
                                elif len(data) == 0:
                                    summary_context_for_llm = "No tasks found matching your criteria."
                            else:
                                st.info("No tasks found matching your criteria.")
                                summary_context_for_llm = "No tasks found matching your criteria."
                                st.session_state.last_interacted_task_details = None 

                        else: 
                            result = db.execute_dml_query(generated_sql)
                            
                            action = "processed"
                            if is_insert_query: 
                                action = "added"
                                if result: 
                                    inserted_task_details = db.get_task_by_id(result, st.session_state.user_email)
                                    if inserted_task_details:
                                        st.session_state.last_interacted_task_details = inserted_task_details
                                        summary_context_for_llm = f"Task '{inserted_task_details.get('task_name')}' (ID: {result}) was {action}."
                                    else:
                                        summary_context_for_llm = f"Task was {action}, but details couldn't be retrieved post-insertion."
                                else:
                                    summary_context_for_llm = f"Task addition was attempted but may not have completed as expected (no ID returned)."
                            elif "UPDATE" in generated_sql.upper(): 
                                action = "updated"
                                summary_context_for_llm = f"Task(s) {action}. {result} row(s) affected."
                            elif "DELETE" in generated_sql.upper(): 
                                action = "deleted"
                                summary_context_for_llm = f"Task(s) {action}. {result} row(s) affected."
                                st.session_state.last_interacted_task_details = None 
                            
                            st.success(f"Task command '{action}' processed successfully.")


                        with st.spinner("📜 Generating friendly summary..."):
                           final_summary = llm_handler.summarize_query_result(
                               llm,
                               user_input_query,
                               generated_sql,
                               summary_context_for_llm
                           )
                        st.markdown(f"**🤖 Summary:**\n {final_summary}")

                    except ValueError as ve:
                         st.error(f"⚠️ Action failed: {ve}")
                    except sqlite3.Error as e: 
                        st.error(f"❌ Database Error: {e}")
                        st.error(f"   Failed Query: {generated_sql}")
                    except Exception as e:
                        st.error(f"❌ An unexpected error occurred: {e}")
                        st.error(f"   Query attempted: {generated_sql}")

            elif not user_input_query.strip() == "":
                st.error("Could not generate an SQL query for your request. Please try rephrasing.")

    st.markdown("---")
    st.markdown("Example commands:")
    st.caption("""
    - Add: "Schedule a 'Team sync meeting' for next Tuesday at 10:00 category Work"
    - Add: "I have a 'Doctor's appointment' tomorrow at 3:30 PM"
    - View: "Show my tasks for this week", "What's pending?"
    - Update: "Mark 'Team sync meeting' as completed" (if it was the last task discussed)
    - Update: "I've finished the 'Doctor's appointment'"
    - Delete: "Cancel the 'Team sync meeting'"
    - Delete: "Remove task 'Old project idea'"
    """)