database_manager.py

import sqlite3
import pandas as pd
from datetime import datetime
from pathlib import Path
import json
from typing import List, Dict, Any, Optional

class DatabaseManager:
    def __init__(self, db_path: str = "business_database.db"):
        self.db_path = db_path
        self.init_database()
    
    def init_database(self):
        """Initialize database with required tables"""
        with sqlite3.connect(self.db_path) as conn:
            # Companies table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS companies (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT NOT NULL,
                    address TEXT,
                    phone TEXT,
                    website TEXT,
                    email TEXT,
                    rating TEXT,
                    reviews TEXT,
                    category TEXT,
                    hours TEXT,
                    price TEXT,
                    description TEXT,
                    dataset_name TEXT NOT NULL,
                    business_type TEXT,
                    location TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE(name, address, dataset_name)
                )
            """)
            
            # Business categories table (for multi-category support)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS company_categories (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    company_id INTEGER,
                    category TEXT,
                    FOREIGN KEY (company_id) REFERENCES companies (id),
                    UNIQUE(company_id, category)
                )
            """)
            
            # Contact collection status
            conn.execute("""
                CREATE TABLE IF NOT EXISTS contact_collection_status (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    company_id INTEGER,
                    website_scraped BOOLEAN DEFAULT FALSE,
                    email_found BOOLEAN DEFAULT FALSE,
                    phone_found BOOLEAN DEFAULT FALSE,
                    last_scraped TIMESTAMP,
                    FOREIGN KEY (company_id) REFERENCES companies (id)
                )
            """)
            
            # Datasets table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS datasets (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT UNIQUE NOT NULL,
                    description TEXT,
                    business_type TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    total_companies INTEGER DEFAULT 0
                )
            """)
            
            conn.commit()
    
    def create_dataset(self, name: str, description: str = "", business_type: str = "") -> bool:
        """Create a new dataset"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                conn.execute("""
                    INSERT INTO datasets (name, description, business_type)
                    VALUES (?, ?, ?)
                """, (name, description, business_type))
                conn.commit()
                return True
        except sqlite3.IntegrityError:
            return False  # Dataset already exists
    
    def get_datasets(self) -> List[Dict]:
        """Get all datasets"""
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.execute("""
                SELECT d.*, COUNT(c.id) as company_count
                FROM datasets d
                LEFT JOIN companies c ON d.name = c.dataset_name
                GROUP BY d.id
                ORDER BY d.created_at DESC
            """)
            return [dict(row) for row in cursor.fetchall()]
    
    def add_companies(self, companies: List[Dict], dataset_name: str, business_type: str = "") -> int:
        """Add companies to database, return number of new companies added"""
        added_count = 0
        
        with sqlite3.connect(self.db_path) as conn:
            for company in companies:
                try:
                    # Insert company
                    cursor = conn.execute("""
                        INSERT INTO companies (
                            name, address, phone, website, email, rating, reviews,
                            category, hours, price, description, dataset_name, business_type, location
                        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """, (
                        company.get('name', ''),
                        company.get('address', ''),
                        company.get('phone', ''),
                        company.get('website', ''),
                        company.get('email', ''),
                        company.get('rating', ''),
                        company.get('reviews', ''),
                        company.get('category', ''),
                        company.get('hours', ''),
                        company.get('price', ''),
                        company.get('description', ''),
                        dataset_name,
                        business_type,
                        company.get('location', '')
                    ))
                    
                    company_id = cursor.lastrowid
                    added_count += 1
                    
                    # Add categories if multiple
                    categories = company.get('categories', [])
                    if not categories and company.get('category'):
                        categories = [company.get('category')]
                    
                    for category in categories:
                        if category:
                            conn.execute("""
                                INSERT OR IGNORE INTO company_categories (company_id, category)
                                VALUES (?, ?)
                            """, (company_id, category))
                    
                    # Initialize contact collection status
                    conn.execute("""
                        INSERT OR IGNORE INTO contact_collection_status (company_id)
                        VALUES (?)
                    """, (company_id,))
                    
                except sqlite3.IntegrityError:
                    # Company already exists, skip
                    continue
            
            # Update dataset company count
            conn.execute("""
                UPDATE datasets 
                SET total_companies = (
                    SELECT COUNT(*) FROM companies WHERE dataset_name = ?
                )
                WHERE name = ?
            """, (dataset_name, dataset_name))
            
            conn.commit()
        
        return added_count
    
    def get_companies(self, dataset_name: str = None, business_type: str = None) -> pd.DataFrame:
        """Get companies from database"""
        with sqlite3.connect(self.db_path) as conn:
            query = """
                SELECT c.*, 
                       GROUP_CONCAT(cc.category, '; ') as all_categories,
                       ccs.website_scraped, ccs.email_found, ccs.phone_found, ccs.last_scraped
                FROM companies c
                LEFT JOIN company_categories cc ON c.id = cc.company_id
                LEFT JOIN contact_collection_status ccs ON c.id = ccs.company_id
            """
            params = []
            conditions = []
            
            if dataset_name:
                conditions.append("c.dataset_name = ?")
                params.append(dataset_name)
            
            if business_type:
                conditions.append("c.business_type = ?")
                params.append(business_type)
            
            if conditions:
                query += " WHERE " + " AND ".join(conditions)
            
            query += " GROUP BY c.id ORDER BY c.created_at DESC"
            
            return pd.read_sql_query(query, conn, params=params)
    
    def update_company(self, company_id: int, updates: Dict) -> bool:
        """Update company information"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                if not updates:
                    return True
                
                # Add updated_at to updates
                updates['updated_at'] = datetime.now().isoformat()
                
                # Build the SET clause and values
                set_clause = ", ".join([f"{key} = ?" for key in updates.keys()])
                values = list(updates.values()) + [company_id]
                
                query = f"""
                    UPDATE companies 
                    SET {set_clause}
                    WHERE id = ?
                """
                
                conn.execute(query, values)
                conn.commit()
                return True
        except Exception as e:
            print(f"Error updating company {company_id}: {e}")
            print(f"Updates: {updates}")
            print(f"Values: {values}")
            return False
    
    def delete_company(self, company_id: int) -> bool:
        """Delete a company from database"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                # Delete from contact_collection_status first (foreign key constraint)
                conn.execute("DELETE FROM contact_collection_status WHERE company_id = ?", (company_id,))
                
                # Delete from company_categories
                conn.execute("DELETE FROM company_categories WHERE company_id = ?", (company_id,))
                
                # Delete the company
                conn.execute("DELETE FROM companies WHERE id = ?", (company_id,))
                conn.commit()
                return True
        except Exception as e:
            print(f"Error deleting company {company_id}: {e}")
            return False
    
    def mark_contact_collected(self, company_id: int, email_found: bool = False, phone_found: bool = False):
        """Mark contact collection status"""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                UPDATE contact_collection_status 
                SET website_scraped = TRUE, 
                    email_found = ?, 
                    phone_found = ?, 
                    last_scraped = CURRENT_TIMESTAMP
                WHERE company_id = ?
            """, (email_found, phone_found, company_id))
            conn.commit()
    
    def get_companies_for_contact_collection(self, dataset_name: str = None, search_type: str = "both") -> List[Dict]:
        """Get companies that need contact information collection - only those missing emails or phones"""
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            
            # Build conditions based on search type
            conditions = ["c.website IS NOT NULL AND c.website != ''"]
            
            if search_type == "emails_only":
                conditions.append("(c.email IS NULL OR c.email = '')")
            elif search_type == "phones_only":
                conditions.append("(c.phone IS NULL OR c.phone = '')")
            else:  # both or all
                conditions.append("""
                    (
                        (c.email IS NULL OR c.email = '') OR 
                        (c.phone IS NULL OR c.phone = '') OR
                        (ccs.website_scraped IS NULL OR ccs.website_scraped = 0)
                    )
                """)
            
            query = f"""
                SELECT c.*, ccs.website_scraped, ccs.email_found, ccs.phone_found
                FROM companies c
                LEFT JOIN contact_collection_status ccs ON c.id = ccs.company_id
                WHERE {' AND '.join(conditions)}
            """
            params = []
            
            if dataset_name:
                query += " AND c.dataset_name = ?"
                params.append(dataset_name)
            
            query += " ORDER BY ccs.last_scraped ASC NULLS FIRST"
            
            cursor = conn.execute(query, params)
            return [dict(row) for row in cursor.fetchall()]
    
    def delete_dataset(self, dataset_name: str) -> bool:
        """Delete a dataset and all its companies"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                # Get company IDs for this dataset
                cursor = conn.execute("SELECT id FROM companies WHERE dataset_name = ?", (dataset_name,))
                company_ids = [row[0] for row in cursor.fetchall()]
                
                if company_ids:
                    # Delete related records
                    conn.execute("DELETE FROM company_categories WHERE company_id IN ({})".format(
                        ','.join('?' * len(company_ids))), company_ids)
                    conn.execute("DELETE FROM contact_collection_status WHERE company_id IN ({})".format(
                        ','.join('?' * len(company_ids))), company_ids)
                    conn.execute("DELETE FROM companies WHERE dataset_name = ?", (dataset_name,))
                
                # Delete dataset
                conn.execute("DELETE FROM datasets WHERE name = ?", (dataset_name,))
                conn.commit()
                return True
        except Exception as e:
            print(f"Error deleting dataset: {e}")
            return False
    
    def get_dataset_stats(self, dataset_name: str) -> Dict:
        """Get statistics for a dataset"""
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            
            # Basic stats
            cursor = conn.execute("""
                SELECT 
                    COUNT(*) as total_companies,
                    COUNT(CASE WHEN email IS NOT NULL AND email != '' THEN 1 END) as with_email,
                    COUNT(CASE WHEN phone IS NOT NULL AND phone != '' THEN 1 END) as with_phone,
                    COUNT(CASE WHEN website IS NOT NULL AND website != '' THEN 1 END) as with_website
                FROM companies 
                WHERE dataset_name = ?
            """, (dataset_name,))
            basic_stats = dict(cursor.fetchone())
            
            # Contact collection stats
            cursor = conn.execute("""
                SELECT 
                    COUNT(*) as scraped_count,
                    COUNT(CASE WHEN email_found = 1 THEN 1 END) as email_found_count,
                    COUNT(CASE WHEN phone_found = 1 THEN 1 END) as phone_found_count
                FROM companies c
                JOIN contact_collection_status ccs ON c.id = ccs.company_id
                WHERE c.dataset_name = ? AND ccs.website_scraped = 1
            """, (dataset_name,))
            contact_stats = dict(cursor.fetchone())
            
            # Companies needing contact collection
            cursor = conn.execute("""
                SELECT COUNT(*) as need_contact_collection
                FROM companies c
                LEFT JOIN contact_collection_status ccs ON c.id = ccs.company_id
                WHERE c.dataset_name = ? 
                AND c.website IS NOT NULL AND c.website != ''
                AND (
                    (c.email IS NULL OR c.email = '') OR 
                    (c.phone IS NULL OR c.phone = '') OR
                    (ccs.website_scraped IS NULL OR ccs.website_scraped = 0)
                )
            """, (dataset_name,))
            need_collection = dict(cursor.fetchone())
            
            return {**basic_stats, **contact_stats, **need_collection}