database.py

import sqlite3
import logging
from typing import Dict, Any, Optional

logger = logging.getLogger(__name__)

class DatabaseManager:
    def __init__(self, db_path: str = "leads.db"):
        self.db_path = db_path
        self._init_db()

    def _init_db(self):
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS leads (
                        id INTEGER PRIMARY KEY AUTOINCREMENT,
                        company_name TEXT NOT NULL,
                        industry_segment TEXT,
                        city TEXT NOT NULL,
                        market_priority TEXT,
                        contact_email TEXT NOT NULL,
                        website_url TEXT,
                        source_platform TEXT,
                        generated_pitch TEXT,
                        status TEXT DEFAULT 'PENDING_REVIEW'
                    )
                ''')
                conn.commit()
                logger.info(f"Database initialized at {self.db_path}")
        except sqlite3.Error as e:
            logger.error(f"Failed to initialize database: {e}")
            raise

    def insert_lead(self, lead_data: Dict[str, Any]) -> Optional[int]:
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute('''
                    INSERT INTO leads (
                        company_name, industry_segment, city, market_priority,
                        contact_email, website_url, source_platform, generated_pitch, status
                    ) VALUES (
                        :company_name, :industry_segment, :city, :market_priority,
                        :contact_email, :website_url, :source_platform, :generated_pitch, :status
                    )
                ''', {
                    'company_name': lead_data.get('company_name'),
                    'industry_segment': lead_data.get('industry_segment'),
                    'city': lead_data.get('city'),
                    'market_priority': lead_data.get('market_priority', 'National'),
                    'contact_email': lead_data.get('contact_email'),
                    'website_url': lead_data.get('website_url'),
                    'source_platform': lead_data.get('source_platform'),
                    'generated_pitch': lead_data.get('generated_pitch'),
                    'status': lead_data.get('status', 'PENDING_REVIEW')
                })
                conn.commit()
                return cursor.lastrowid
        except sqlite3.Error as e:
            logger.error(f"Failed to insert lead {lead_data.get('company_name')}: {e}")
            return None