insucompass/services/database.py

import sqlite3
import logging
import json
from contextlib import contextmanager
from datetime import datetime
from pathlib import Path
from urllib.parse import urlparse
from typing import Optional, Dict, Any

from ..config import settings

# Configure logging
logging.basicConfig(level=settings.LOG_LEVEL)
logger = logging.getLogger(__name__)

@contextmanager
def get_db_connection():
    """Context manager for handling database connections."""
    conn = None
    try:
        conn = sqlite3.connect(settings.DATABASE_URL)
        conn.row_factory = sqlite3.Row
        logger.debug("Database connection established.")
        yield conn
    except sqlite3.Error as e:
        logger.error(f"Database connection error: {e}")
        raise
    finally:
        if conn:
            conn.close()
            logger.debug("Database connection closed.")

def setup_database():
    """Creates/updates all necessary tables in the database."""
    logger.info("Setting up database schema...")
    ddl_statements = [
        """
        CREATE TABLE IF NOT EXISTS data_sources (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name TEXT NOT NULL,
            url TEXT NOT NULL UNIQUE,
            data_type TEXT NOT NULL,
            category TEXT,
            local_path TEXT,
            content_hash TEXT,
            status TEXT DEFAULT 'pending',
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP
        );
        """,
        """
        CREATE TABLE IF NOT EXISTS knowledge_chunks (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            source_id INTEGER NOT NULL,
            chunk_text TEXT NOT NULL,
            metadata_json TEXT,
            vector_id TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (source_id) REFERENCES data_sources (id)
        );
        """,
        """
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT NOT NULL UNIQUE,
            hashed_password TEXT NOT NULL,
            role TEXT NOT NULL DEFAULT 'user',
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        );
        """,
        """
        CREATE TABLE IF NOT EXISTS user_profiles (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL UNIQUE,
            profile_data_json TEXT,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (user_id) REFERENCES users (id)
        );
        """
    ]
    with get_db_connection() as conn:
        for statement in ddl_statements:
            conn.cursor().execute(statement)
        conn.commit()
    logger.info("Database schema setup complete.")

# --- Crawler-related Helpers ---

def initialize_crawl_jobs():
    """Populates the data_sources table with the initial crawl job start URLs."""
    logger.info("Initializing crawl jobs from config...")
    query = "INSERT OR IGNORE INTO data_sources (name, url, data_type, category, status, created_at) VALUES (?, ?, ?, ?, ?, ?)"
    with get_db_connection() as conn:
        for job in settings.CRAWLING_JOBS:
            conn.cursor().execute(query, (
                job['name'],
                job['start_url'],
                'start_url',
                job.get('domain_lock', 'General'),
                'pending',
                datetime.now()
            ))
        conn.commit()
    logger.info("Crawl jobs initialized successfully.")


def find_or_create_web_source(url: str, name: str) -> int:
    """
    Finds an existing data source by URL or creates a new one if it doesn't exist.
    Used for dynamically ingested web search results.

    Returns:
        The ID of the data source.
    """
    with get_db_connection() as conn:
        cursor = conn.cursor()
        
        # 1. Check if it exists
        cursor.execute("SELECT id FROM data_sources WHERE url = ?", (url,))
        row = cursor.fetchone()
        if row:
            return row['id']
        
        # 2. If not, create it
        logger.info(f"Registering new dynamic web source: {url}")
        insert_query = """
        INSERT INTO data_sources (name, url, data_type, category, status, created_at, updated_at)
        VALUES (?, ?, ?, ?, ?, ?, ?)
        """
        cursor.execute(insert_query, (
            name, url, 'web_search_result', 'Dynamic', 'ingested', datetime.now(), datetime.now()
        ))
        conn.commit()
        return cursor.lastrowid

    
def add_discovered_source(url: str, category: str, data_type: str) -> int:
    """Adds a newly discovered URL to the database if it doesn't exist."""
    with get_db_connection() as conn:
        cursor = conn.cursor()
        cursor.execute("SELECT id FROM data_sources WHERE url = ?", (url,))
        row = cursor.fetchone()
        if row:
            return row['id']
        
        path_name = Path(urlparse(url).path).name
        if not path_name:
            path_name = urlparse(url).path.strip('/').replace('/', '_') or urlparse(url).netloc
        
        insert_query = "INSERT INTO data_sources (name, url, data_type, category, status, created_at) VALUES (?, ?, ?, ?, ?, ?)"
        cursor.execute(insert_query, (path_name, url, data_type, category, 'pending', datetime.now()))
        conn.commit()
        logger.debug(f"Discovered and added new source: {url}")
        return cursor.lastrowid

# --- User Management Helpers ---

def create_user(username: str, hashed_password: str, role: str = 'user') -> Optional[int]:
    """Creates a new user in the database."""
    query = "INSERT INTO users (username, hashed_password, role) VALUES (?, ?, ?)"
    try:
        with get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(query, (username, hashed_password, role))
            conn.commit()
            logger.info(f"User '{username}' created successfully.")
            return cursor.lastrowid
    except sqlite3.IntegrityError:
        logger.warning(f"Attempted to create a user that already exists: {username}")
        return None
    except sqlite3.Error as e:
        logger.error(f"Database error while creating user {username}: {e}")
        return None

def get_user_by_username(username: str) -> Optional[Dict[str, Any]]:
    """Retrieves a user by their username."""
    query = "SELECT * FROM users WHERE username = ?"
    with get_db_connection() as conn:
        cursor = conn.cursor()
        cursor.execute(query, (username,))
        row = cursor.fetchone()
        return dict(row) if row else None

# --- User Profile Helpers ---

def create_or_update_user_profile(user_id: int, profile_data: Dict[str, Any]) -> bool:
    """Creates or updates a user's 360-degree profile."""
    profile_json = json.dumps(profile_data)
    current_time = datetime.now()
    
    # Using INSERT OR REPLACE for simplicity (UPSERT)
    query = """
    INSERT INTO user_profiles (user_id, profile_data_json, updated_at)
    VALUES (?, ?, ?)
    ON CONFLICT(user_id) DO UPDATE SET
        profile_data_json = excluded.profile_data_json,
        updated_at = excluded.updated_at;
    """
    try:
        with get_db_connection() as conn:
            conn.cursor().execute(query, (user_id, profile_json, current_time))
            conn.commit()
        logger.info(f"Profile for user_id {user_id} created/updated.")
        return True
    except sqlite3.Error as e:
        logger.error(f"Database error while updating profile for user_id {user_id}: {e}")
        return False

def get_user_profile(user_id: int) -> Optional[Dict[str, Any]]:
    """Retrieves a user's profile."""
    query = "SELECT profile_data_json FROM user_profiles WHERE user_id = ?"
    with get_db_connection() as conn:
        cursor = conn.cursor()
        cursor.execute(query, (user_id,))
        row = cursor.fetchone()
        if row and row['profile_data_json']:
            return json.loads(row['profile_data_json'])
    return None