database.py

import sqlite3
import json
from datetime import datetime
from typing import List, Dict, Optional
from config import DATABASE_PATH

class NewsDatabase:
    def __init__(self, db_path: str = DATABASE_PATH):
        self.db_path = db_path
        self.init_database()
    
    def init_database(self):
        """Initialize the database with required tables."""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            
            # Create articles table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS articles (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    title TEXT NOT NULL,
                    link TEXT UNIQUE NOT NULL,
                    description TEXT,
                    content TEXT,
                    summary TEXT,
                    source TEXT NOT NULL,
                    published_date TEXT,
                    relevance_score REAL DEFAULT 0.0,
                    keywords_found TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # Create feeds table to track feed status
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS feeds (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT UNIQUE NOT NULL,
                    url TEXT NOT NULL,
                    last_fetch TIMESTAMP,
                    status TEXT DEFAULT 'active',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # Create keywords table for relevance tracking
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS keywords (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    keyword TEXT UNIQUE NOT NULL,
                    category TEXT,
                    weight REAL DEFAULT 1.0,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            conn.commit()
    
    def add_article(self, article_data: Dict) -> bool:
        """Add a new article to the database."""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                cursor.execute('''
                    INSERT OR REPLACE INTO articles 
                    (title, link, description, content, summary, source, published_date, 
                     relevance_score, keywords_found, updated_at)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                ''', (
                    article_data.get('title', ''),
                    article_data.get('link', ''),
                    article_data.get('description', ''),
                    article_data.get('content', ''),
                    article_data.get('summary', ''),
                    article_data.get('source', ''),
                    article_data.get('published_date', ''),
                    article_data.get('relevance_score', 0.0),
                    json.dumps(article_data.get('keywords_found', [])),
                    datetime.now().isoformat()
                ))
                
                conn.commit()
                return True
        except Exception as e:
            print(f"Error adding article: {e}")
            return False
    
    def get_articles(self, limit: int = 50, source: Optional[str] = None, 
                    min_relevance: float = 0.0) -> List[Dict]:
        """Retrieve articles from the database."""
        try:
            with sqlite3.connect(self.db_path) as conn:
                conn.row_factory = sqlite3.Row
                cursor = conn.cursor()
                
                query = '''
                    SELECT * FROM articles 
                    WHERE relevance_score >= ?
                '''
                params = [min_relevance]
                
                if source:
                    query += ' AND source = ?'
                    params.append(source)
                
                query += ' ORDER BY published_date DESC, relevance_score DESC LIMIT ?'
                params.append(limit)
                
                cursor.execute(query, params)
                rows = cursor.fetchall()
                
                articles = []
                for row in rows:
                    article = dict(row)
                    article['keywords_found'] = json.loads(article['keywords_found'] or '[]')
                    articles.append(article)
                
                return articles
        except Exception as e:
            print(f"Error retrieving articles: {e}")
            return []
    
    def get_articles_by_keyword(self, keyword: str, limit: int = 20) -> List[Dict]:
        """Retrieve articles that contain a specific keyword."""
        try:
            with sqlite3.connect(self.db_path) as conn:
                conn.row_factory = sqlite3.Row
                cursor = conn.cursor()
                
                cursor.execute('''
                    SELECT * FROM articles 
                    WHERE keywords_found LIKE ? OR title LIKE ? OR description LIKE ?
                    ORDER BY published_date DESC, relevance_score DESC 
                    LIMIT ?
                ''', (f'%{keyword}%', f'%{keyword}%', f'%{keyword}%', limit))
                
                rows = cursor.fetchall()
                articles = []
                for row in rows:
                    article = dict(row)
                    article['keywords_found'] = json.loads(article['keywords_found'] or '[]')
                    articles.append(article)
                
                return articles
        except Exception as e:
            print(f"Error retrieving articles by keyword: {e}")
            return []
    
    def update_feed_status(self, feed_name: str, last_fetch: str, status: str = 'active'):
        """Update the status of a feed."""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                cursor.execute('''
                    INSERT OR REPLACE INTO feeds (name, url, last_fetch, status)
                    VALUES (?, ?, ?, ?)
                ''', (feed_name, '', last_fetch, status))
                
                conn.commit()
        except Exception as e:
            print(f"Error updating feed status: {e}")
    
    def get_feed_status(self) -> List[Dict]:
        """Get the status of all feeds."""
        try:
            with sqlite3.connect(self.db_path) as conn:
                conn.row_factory = sqlite3.Row
                cursor = conn.cursor()
                
                cursor.execute('SELECT * FROM feeds ORDER BY last_fetch DESC')
                rows = cursor.fetchall()
                
                return [dict(row) for row in rows]
        except Exception as e:
            print(f"Error retrieving feed status: {e}")
            return []
    
    def get_article_count(self) -> int:
        """Get the total number of articles in the database."""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute('SELECT COUNT(*) FROM articles')
                return cursor.fetchone()[0]
        except Exception as e:
            print(f"Error getting article count: {e}")
            return 0
    
    def cleanup_old_articles(self, days: int = 30):
        """Remove articles older than specified days."""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                cursor.execute('''
                    DELETE FROM articles 
                    WHERE created_at < datetime('now', '-{} days')
                '''.format(days))
                
                deleted_count = cursor.rowcount
                conn.commit()
                print(f"Cleaned up {deleted_count} old articles")
        except Exception as e:
            print(f"Error cleaning up old articles: {e}")