db.py

import sqlite3
import logging
from datetime import datetime
from typing import Optional, Dict, List
import os
from dotenv import load_dotenv

load_dotenv()

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

DATABASE_PATH = os.getenv('DATABASE_PATH', 'slack_agent.db')

def init_database():
    """Initialize the SQLite database with the required schema."""
    try:
        conn = sqlite3.connect(DATABASE_PATH)
        cursor = conn.cursor()
        
        # schema required
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS tasks (
                thread_ts TEXT PRIMARY KEY,
                linear_task_id TEXT,
                title TEXT,
                brain_dump TEXT,
                draft TEXT,
                tone TEXT,
                status TEXT DEFAULT 'waiting_dump',
                created_at TEXT,
                updated_at TEXT
            )
        ''')
        
        # Only keep the primary key index (thread_ts)
        
        conn.commit()
        conn.close()
        
        logger.info("Database initialized successfully")
        
    except Exception as e:
        logger.error(f"Error initializing database: {e}")
        raise

def create_task(thread_ts: str, linear_task_id: str, title: str) -> bool:
    """Create a new task in the database."""
    try:
        conn = sqlite3.connect(DATABASE_PATH)
        cursor = conn.cursor()
        
        now = datetime.now().isoformat()
        
        cursor.execute('''
            INSERT INTO tasks (thread_ts, linear_task_id, title, status, created_at, updated_at)
            VALUES (?, ?, ?, 'waiting_dump', ?, ?)
        ''', (thread_ts, linear_task_id, title, now, now))
        
        conn.commit()
        conn.close()
        
        logger.info(f"Created task: {title} (thread_ts: {thread_ts}, linear_id: {linear_task_id})")
        return True
        
    except Exception as e:
        logger.error(f"Error creating task: {e}")
        return False

def get_task_by_thread_ts(thread_ts: str) -> Optional[Dict]:
    """Get a task by its thread timestamp."""
    try:
        conn = sqlite3.connect(DATABASE_PATH)
        cursor = conn.cursor()
        
        cursor.execute('SELECT * FROM tasks WHERE thread_ts = ?', (thread_ts,))
        row = cursor.fetchone()
        
        conn.close()
        
        if row:
            return {
                'thread_ts': row[0],
                'linear_task_id': row[1],
                'title': row[2],
                'brain_dump': row[3],
                'draft': row[4],
                'tone': row[5],
                'status': row[6],
                'created_at': row[7],
                'updated_at': row[8]
            }
        return None
        
    except Exception as e:
        logger.error(f"Error getting task by thread_ts {thread_ts}: {e}")
        return None

def update_brain_dump(thread_ts: str, brain_dump: str) -> bool:
    """Update the brain dump for a task."""
    try:
        conn = sqlite3.connect(DATABASE_PATH)
        cursor = conn.cursor()
        
        now = datetime.now().isoformat()
        
        cursor.execute('''
            UPDATE tasks 
            SET brain_dump = ?, updated_at = ?
            WHERE thread_ts = ?
        ''', (brain_dump, now, thread_ts))
        
        conn.commit()
        conn.close()
        
        logger.info(f"Updated brain dump for thread_ts: {thread_ts}")
        return True
        
    except Exception as e:
        logger.error(f"Error updating brain dump: {e}")
        return False

def update_task_status(thread_ts: str, status: str) -> bool:
    """Update the status of a task."""
    try:
        conn = sqlite3.connect(DATABASE_PATH)
        cursor = conn.cursor()
        
        now = datetime.now().isoformat()
        
        cursor.execute('''
            UPDATE tasks 
            SET status = ?, updated_at = ?
            WHERE thread_ts = ?
        ''', (status, now, thread_ts))
        
        conn.commit()
        conn.close()
        
        logger.info(f"Updated task status to '{status}' for thread_ts: {thread_ts}")
        return True
        
    except Exception as e:
        logger.error(f"Error updating task status: {e}")
        return False

def update_draft(thread_ts: str, draft: str, tone: str = None) -> bool:
    """Update the draft and optionally the tone for a task."""
    try:
        conn = sqlite3.connect(DATABASE_PATH)
        cursor = conn.cursor()
        
        now = datetime.now().isoformat()
        
        if tone:
            cursor.execute('''
                UPDATE tasks 
                SET draft = ?, tone = ?, updated_at = ?
                WHERE thread_ts = ?
            ''', (draft, tone, now, thread_ts))
        else:
            cursor.execute('''
                UPDATE tasks 
                SET draft = ?, updated_at = ?
                WHERE thread_ts = ?
            ''', (draft, now, thread_ts))
        
        conn.commit()
        conn.close()
        
        logger.info(f"Updated draft for thread_ts: {thread_ts}")
        return True
        
    except Exception as e:
        logger.error(f"Error updating draft: {e}")
        return False

def append_brain_dump(thread_ts: str, new_content: str) -> bool:
    """Append new content to existing brain dump."""
    try:
        task = get_task_by_thread_ts(thread_ts)
        if not task:
            logger.error(f"Task not found for thread_ts: {thread_ts}")
            return False
        
        existing_dump = task.get('brain_dump', '') or ''
        updated_dump = existing_dump + '\n' + new_content if existing_dump else new_content
        
        return update_brain_dump(thread_ts, updated_dump)
        
    except Exception as e:
        logger.error(f"Error appending brain dump: {e}")
        return False

def delete_task(thread_ts: str) -> bool:
    """Delete a task by thread_ts. Used primarily for testing cleanup."""
    try:
        conn = sqlite3.connect(DATABASE_PATH)
        cursor = conn.cursor()
        
        cursor.execute('DELETE FROM tasks WHERE thread_ts = ?', (thread_ts,))
        
        conn.commit()
        conn.close()
        
        logger.info(f"Deleted task with thread_ts: {thread_ts}")
        return True
        
    except Exception as e:
        logger.error(f"Error deleting task: {e}")
        return False

# Initialize database when module is imported
if __name__ == "__main__":
    init_database()
    logger.info("Database module initialized")