tracker.py

"""
Productivity tracker data model and database operations.

This module handles all task-related data operations including
CRUD operations, database management, and data validation.
"""

import sqlite3
import os
from datetime import datetime, date, timedelta
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass
from utils import validate_duration, validate_focus_score, format_duration, parse_duration


@dataclass
class Task:
    """Data class representing a productivity task."""
    id: Optional[int] = None
    date: str = ""
    task_name: str = ""
    duration_minutes: int = 0
    category: str = ""
    focus_score: int = 5
    notes: str = ""

    def __post_init__(self):
        """Validate data after initialization."""
        if self.duration_minutes < 0:
            raise ValueError("Duration cannot be negative")
        if not validate_focus_score(self.focus_score):
            raise ValueError("Focus score must be between 1 and 10")
        if not self.task_name.strip():
            raise ValueError("Task name cannot be empty")


class ProductivityTracker:
    """Manages task data and database operations."""
    
    def __init__(self, db_path: str = "data/time_log.db"):
        """Initialize the productivity tracker with database path."""
        self.db_path = db_path
        self._ensure_data_directory()
        self._init_database()
    
    def _ensure_data_directory(self) -> None:
        """Ensure the data directory exists."""
        os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
    
    def _init_database(self) -> None:
        """Initialize the SQLite database with tasks table."""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS tasks (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    date TEXT NOT NULL,
                    task_name TEXT NOT NULL,
                    duration_minutes INTEGER NOT NULL,
                    category TEXT NOT NULL,
                    focus_score INTEGER NOT NULL,
                    notes TEXT DEFAULT '',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            conn.commit()
    
    def add_task(self, task: Task) -> int:
        """
        Add a new task to the database.
        
        Args:
            task: Task object to add
            
        Returns:
            int: ID of the created task
            
        Raises:
            ValueError: If task data is invalid
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute("""
                    INSERT INTO tasks (date, task_name, duration_minutes, category, focus_score, notes)
                    VALUES (?, ?, ?, ?, ?, ?)
                """, (
                    task.date,
                    task.task_name,
                    task.duration_minutes,
                    task.category,
                    task.focus_score,
                    task.notes
                ))
                conn.commit()
                return cursor.lastrowid
        except sqlite3.Error as e:
            raise ValueError(f"Database error: {e}")
    
    def get_all_tasks(self) -> List[Task]:
        """
        Retrieve all tasks from the database.
        
        Returns:
            List[Task]: List of all tasks
        """
        tasks = []
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute("SELECT * FROM tasks ORDER BY date DESC, created_at DESC")
                rows = cursor.fetchall()
                
                for row in rows:
                    task = Task(
                        id=row[0],
                        date=row[1],
                        task_name=row[2],
                        duration_minutes=row[3],
                        category=row[4],
                        focus_score=row[5],
                        notes=row[6] or ""
                    )
                    tasks.append(task)
        except sqlite3.Error as e:
            print(f"Database error: {e}")
        
        return tasks
    
    def get_task_by_id(self, task_id: int) -> Optional[Task]:
        """
        Retrieve a task by its ID.
        
        Args:
            task_id: ID of the task to retrieve
            
        Returns:
            Optional[Task]: Task object or None if not found
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute("SELECT * FROM tasks WHERE id = ?", (task_id,))
                row = cursor.fetchone()
                
                if row:
                    return Task(
                        id=row[0],
                        date=row[1],
                        task_name=row[2],
                        duration_minutes=row[3],
                        category=row[4],
                        focus_score=row[5],
                        notes=row[6] or ""
                    )
        except sqlite3.Error as e:
            print(f"Database error: {e}")
        
        return None
    
    def update_task(self, task: Task) -> bool:
        """
        Update an existing task.
        
        Args:
            task: Task object with updated data
            
        Returns:
            bool: True if update was successful, False otherwise
        """
        if not task.id:
            raise ValueError("Task ID is required for update")
        
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute("""
                    UPDATE tasks 
                    SET date = ?, task_name = ?, duration_minutes = ?, 
                        category = ?, focus_score = ?, notes = ?
                    WHERE id = ?
                """, (
                    task.date,
                    task.task_name,
                    task.duration_minutes,
                    task.category,
                    task.focus_score,
                    task.notes,
                    task.id
                ))
                conn.commit()
                return cursor.rowcount > 0
        except sqlite3.Error as e:
            print(f"Database error: {e}")
            return False
    
    def delete_task(self, task_id: int) -> bool:
        """
        Delete a task by ID.
        
        Args:
            task_id: ID of the task to delete
            
        Returns:
            bool: True if deletion was successful, False otherwise
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute("DELETE FROM tasks WHERE id = ?", (task_id,))
                conn.commit()
                return cursor.rowcount > 0
        except sqlite3.Error as e:
            print(f"Database error: {e}")
            return False
    
    def get_tasks_by_date(self, target_date: str) -> List[Task]:
        """
        Get all tasks for a specific date.
        
        Args:
            target_date: Date in YYYY-MM-DD format
            
        Returns:
            List[Task]: List of tasks for the specified date
        """
        tasks = []
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute("SELECT * FROM tasks WHERE date = ? ORDER BY created_at", (target_date,))
                rows = cursor.fetchall()
                
                for row in rows:
                    task = Task(
                        id=row[0],
                        date=row[1],
                        task_name=row[2],
                        duration_minutes=row[3],
                        category=row[4],
                        focus_score=row[5],
                        notes=row[6] or ""
                    )
                    tasks.append(task)
        except sqlite3.Error as e:
            print(f"Database error: {e}")
        
        return tasks
    
    def get_tasks_by_category(self, category: str) -> List[Task]:
        """
        Get all tasks in a specific category.
        
        Args:
            category: Category to filter by
            
        Returns:
            List[Task]: List of tasks in the category
        """
        tasks = []
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute("SELECT * FROM tasks WHERE category = ? ORDER BY date DESC", (category,))
                rows = cursor.fetchall()
                
                for row in rows:
                    task = Task(
                        id=row[0],
                        date=row[1],
                        task_name=row[2],
                        duration_minutes=row[3],
                        category=row[4],
                        focus_score=row[5],
                        notes=row[6] or ""
                    )
                    tasks.append(task)
        except sqlite3.Error as e:
            print(f"Database error: {e}")
        
        return tasks
    
    def get_tasks_by_date_range(self, start_date: str, end_date: str) -> List[Task]:
        """
        Get all tasks within a date range.
        
        Args:
            start_date: Start date in YYYY-MM-DD format
            end_date: End date in YYYY-MM-DD format
            
        Returns:
            List[Task]: List of tasks within the date range
        """
        tasks = []
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute("""
                    SELECT * FROM tasks 
                    WHERE date >= ? AND date <= ? 
                    ORDER BY date DESC, created_at DESC
                """, (start_date, end_date))
                rows = cursor.fetchall()
                
                for row in rows:
                    task = Task(
                        id=row[0],
                        date=row[1],
                        task_name=row[2],
                        duration_minutes=row[3],
                        category=row[4],
                        focus_score=row[5],
                        notes=row[6] or ""
                    )
                    tasks.append(task)
        except sqlite3.Error as e:
            print(f"Database error: {e}")
        
        return tasks
    
    def get_total_time_by_date(self, target_date: str) -> int:
        """
        Calculate total time spent on a specific date.
        
        Args:
            target_date: Date in YYYY-MM-DD format
            
        Returns:
            int: Total minutes spent on the date
        """
        tasks = self.get_tasks_by_date(target_date)
        return sum(task.duration_minutes for task in tasks)
    
    def get_average_focus_by_date(self, target_date: str) -> float:
        """
        Calculate average focus score for a specific date.
        
        Args:
            target_date: Date in YYYY-MM-DD format
            
        Returns:
            float: Average focus score for the date
        """
        tasks = self.get_tasks_by_date(target_date)
        if not tasks:
            return 0.0
        
        total_focus = sum(task.focus_score for task in tasks)
        return total_focus / len(tasks)
    
    def get_weekly_summary(self, start_date: str) -> Dict[str, Dict]:
        """
        Get weekly summary starting from a specific date.
        
        Args:
            start_date: Start date in YYYY-MM-DD format
            
        Returns:
            Dict[str, Dict]: Dictionary with daily summaries
        """
        start = datetime.strptime(start_date, '%Y-%m-%d').date()
        weekly_data = {}
        
        for i in range(7):
            current_date = start + timedelta(days=i)
            date_str = current_date.strftime('%Y-%m-%d')
            
            tasks = self.get_tasks_by_date(date_str)
            total_time = sum(task.duration_minutes for task in tasks)
            avg_focus = sum(task.focus_score for task in tasks) / len(tasks) if tasks else 0
            
            weekly_data[date_str] = {
                'total_time_minutes': total_time,
                'total_time_hours': total_time / 60,
                'average_focus': avg_focus,
                'task_count': len(tasks)
            }
        
        return weekly_data
    
    def get_category_summary(self) -> Dict[str, Dict]:
        """
        Get summary statistics by category.
        
        Returns:
            Dict[str, Dict]: Dictionary with category summaries
        """
        tasks = self.get_all_tasks()
        category_data = {}
        
        for task in tasks:
            if task.category not in category_data:
                category_data[task.category] = {
                    'total_time_minutes': 0,
                    'total_time_hours': 0,
                    'task_count': 0,
                    'average_focus': 0,
                    'focus_scores': []
                }
            
            category_data[task.category]['total_time_minutes'] += task.duration_minutes
            category_data[task.category]['total_time_hours'] += task.duration_minutes / 60
            category_data[task.category]['task_count'] += 1
            category_data[task.category]['focus_scores'].append(task.focus_score)
        
        # Calculate average focus scores
        for category in category_data:
            scores = category_data[category]['focus_scores']
            category_data[category]['average_focus'] = sum(scores) / len(scores) if scores else 0
            del category_data[category]['focus_scores']  # Remove raw scores
        
        return category_data
    
    def search_tasks(self, query: str) -> List[Task]:
        """
        Search tasks by name or category.
        
        Args:
            query: Search query string
            
        Returns:
            List[Task]: List of matching tasks
        """
        results = []
        query_lower = query.lower()
        
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute("""
                    SELECT * FROM tasks 
                    WHERE LOWER(task_name) LIKE ? OR LOWER(category) LIKE ? OR LOWER(notes) LIKE ?
                    ORDER BY date DESC, created_at DESC
                """, (f"%{query_lower}%", f"%{query_lower}%", f"%{query_lower}%"))
                rows = cursor.fetchall()
                
                for row in rows:
                    task = Task(
                        id=row[0],
                        date=row[1],
                        task_name=row[2],
                        duration_minutes=row[3],
                        category=row[4],
                        focus_score=row[5],
                        notes=row[6] or ""
                    )
                    results.append(task)
        except sqlite3.Error as e:
            print(f"Database error: {e}")
        
        return results
    
    def get_productivity_stats(self) -> Dict[str, float]:
        """
        Get overall productivity statistics.
        
        Returns:
            Dict[str, float]: Dictionary with productivity metrics
        """
        tasks = self.get_all_tasks()
        
        if not tasks:
            return {
                'total_tasks': 0,
                'total_time_hours': 0,
                'average_focus': 0,
                'average_task_duration': 0,
                'most_productive_category': '',
                'most_productive_day': ''
            }
        
        # Basic stats
        total_tasks = len(tasks)
        total_time_minutes = sum(task.duration_minutes for task in tasks)
        total_time_hours = total_time_minutes / 60
        average_focus = sum(task.focus_score for task in tasks) / total_tasks
        average_task_duration = total_time_minutes / total_tasks
        
        # Most productive category
        category_times = {}
        for task in tasks:
            if task.category not in category_times:
                category_times[task.category] = 0
            category_times[task.category] += task.duration_minutes
        
        most_productive_category = max(category_times.items(), key=lambda x: x[1])[0] if category_times else ''
        
        # Most productive day
        daily_times = {}
        for task in tasks:
            if task.date not in daily_times:
                daily_times[task.date] = 0
            daily_times[task.date] += task.duration_minutes
        
        most_productive_day = max(daily_times.items(), key=lambda x: x[1])[0] if daily_times else ''
        
        return {
            'total_tasks': total_tasks,
            'total_time_hours': total_time_hours,
            'average_focus': average_focus,
            'average_task_duration': average_task_duration,
            'most_productive_category': most_productive_category,
            'most_productive_day': most_productive_day
        }


# TODO: Add Pomodoro Timer Integration (countdowns, breaks, notifications)
# TODO: Add Productivity Scoring Algorithm (based on focus + time spent)
# TODO: Fix occasional bug: chart not updating if data added after chart generation
# TODO: Add Export Reports to Excel or PDF feature
# TODO: Integrate Rich library for colored CLI output
# TODO: Add Daily Goal Tracker (target hours per day)
# TODO: Implement Unit Tests for utils.py functions