import streamlit as st
import pandas as pd
from datetime import datetime
import os

class TaskManager:
    TASKS_FILE = "tasks.csv"
    CATEGORIES = ['Learning', 'Gym', 'Personal', 'Family', 'Work', 'Prayer']

    def __init__(self):
        self.local_timezone = 'Asia/Dhaka'  # GMT+6
        if 'tasks' not in st.session_state:
            st.session_state.tasks = self.load_tasks()

    def load_tasks(self):
        if os.path.exists(self.TASKS_FILE):
            if os.path.getsize(self.TASKS_FILE) > 0:
                tasks = pd.read_csv(self.TASKS_FILE, parse_dates=['Task Time']).to_dict(orient='records')
                # Ensure timezone-awareness when loading tasks
                for task in tasks:
                    task['Task Time'] = pd.Timestamp(task['Task Time']).tz_localize('UTC').tz_convert(self.local_timezone)
                return tasks
            else:
                return []
        else:
            df = pd.DataFrame(columns=["Task Name", "Task Time", "Task Duration (hours)", "Task Duration (minutes)", "Category"])
            df.to_csv(self.TASKS_FILE, index=False)
            return []

    def save_tasks(self, tasks):
        df = pd.DataFrame(tasks)
        df['Task Time'] = df['Task Time'].apply(lambda x: x.tz_convert('UTC'))
        df.to_csv(self.TASKS_FILE, index=False)

    def add_task(self, task_name, task_time, task_duration_hours, task_duration_minutes, task_category):
        task_time_full = datetime.combine(datetime.today(), task_time).astimezone(self.local_timezone)
        task_entry = {
            "Task Name": task_name,
            "Task Time": task_time_full,
            "Task Duration (hours)": int(task_duration_hours),
            "Task Duration (minutes)": int(task_duration_minutes),
            "Category": task_category
        }
        st.session_state.tasks.append(task_entry)
        self.save_tasks(st.session_state.tasks)

    def delete_task_by_name(self, task_name):
        for index, task in enumerate(st.session_state.tasks):
            if task['Task Name'] == task_name:
                st.session_state.tasks.pop(index)
                self.save_tasks(st.session_state.tasks)
                return True
        return False

    def generate_report(self, timeframe):
        df = pd.DataFrame(st.session_state.tasks)
        if df.empty:
            return pd.DataFrame()  # Return empty DataFrame if no tasks are present

        # Ensure 'Task Time' is already timezone-aware and convert it to local time for the report
        df['Task Time'] = pd.to_datetime(df['Task Time']).apply(lambda x: x.tz_convert(self.local_timezone))

        today = pd.Timestamp.now(tz=self.local_timezone).date()

        if timeframe == 'daily':
            report = df[df['Task Time'].dt.date == today]
        elif timeframe == 'weekly':
            week_start = pd.Timestamp.now(tz=self.local_timezone) - pd.DateOffset(days=pd.Timestamp.now(tz=self.local_timezone).dayofweek)
            report = df[(df['Task Time'] >= week_start) & (df['Task Time'] < pd.Timestamp.now(tz=self.local_timezone) + pd.DateOffset(days=1))]
        elif timeframe == 'monthly':
            report = df[df['Task Time'].dt.month == pd.Timestamp.now(tz=self.local_timezone).month]
        elif timeframe == 'yearly':
            report = df[df['Task Time'].dt.year == pd.Timestamp.now(tz=self.local_timezone).year]
        else:
            report = pd.DataFrame()  # Empty DataFrame for unsupported timeframes

        report['Total Duration'] = report['Task Duration (hours)'] + report['Task Duration (minutes)'] / 60.0
        return report

    def download_report(self):
        df = pd.DataFrame(st.session_state.tasks)
        csv = df.to_csv(index=False)
        st.download_button("Download CSV", data=csv, file_name="task_report.csv", mime='text/csv')