database.py

import sqlite3
import os
import json
import time
import functools
from contextlib import contextmanager
from datetime import datetime, timedelta
from typing import List, Dict, Any, Optional
from logger import get_logger

logger = get_logger(__name__)

# ── SQLite retry decorator for SQLITE_BUSY ───────────────────────────────────

def sqlite_retry(max_retries: int = 5, base_delay: float = 0.1):
    """Decorator that retries DB operations on sqlite3.OperationalError (BUSY/LOCKED)."""
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            last_err = None
            for attempt in range(max_retries):
                try:
                    return func(*args, **kwargs)
                except sqlite3.OperationalError as e:
                    if "locked" in str(e).lower() or "busy" in str(e).lower():
                        last_err = e
                        delay = base_delay * (2 ** attempt)
                        logger.warning(f"DB busy/locked (attempt {attempt+1}/{max_retries}), retrying in {delay:.1f}s...")
                        time.sleep(delay)
                    else:
                        raise
            raise last_err  # type: ignore
        return wrapper
    return decorator


class OutreachDatabase:
    def __init__(self, db_path: str = "outreach.db"):
        self.db_path = db_path
        self.init_database()

    def _conn(self):
        conn = sqlite3.connect(self.db_path, timeout=10)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA foreign_keys=ON")
        conn.execute("PRAGMA busy_timeout=5000")
        return conn

    @contextmanager
    def _get_conn(self):
        """Context manager that guarantees connection cleanup even on exceptions."""
        conn = self._conn()
        try:
            yield conn
            conn.commit()
        except Exception:
            conn.rollback()
            raise
        finally:
            conn.close()

    # ── Schema ───────────────────────────────────────────────────────────────

    def init_database(self):
        """Initialize database with all required tables."""
        conn = self._conn()
        c = conn.cursor()

        # Original tables (kept for backward compat)
        c.execute('''CREATE TABLE IF NOT EXISTS outreach_campaigns (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            campaign_id TEXT UNIQUE NOT NULL,
            user_prompt TEXT NOT NULL,
            target_roles TEXT NOT NULL,
            company_type TEXT NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )''')

        c.execute('''CREATE TABLE IF NOT EXISTS outreach_records (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            campaign_id TEXT NOT NULL,
            company_name TEXT NOT NULL,
            company_domain TEXT NOT NULL,
            contact_name TEXT NOT NULL,
            contact_email TEXT NOT NULL,
            targeted_role TEXT NOT NULL,
            email_draft TEXT NOT NULL,
            status TEXT DEFAULT 'draft',
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (campaign_id) REFERENCES outreach_campaigns(campaign_id)
        )''')

        # ── Master company pool ──────────────────────────────────────────────
        c.execute('''CREATE TABLE IF NOT EXISTS companies (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name TEXT UNIQUE NOT NULL,
            domain TEXT,
            email TEXT NOT NULL,
            contact_name TEXT DEFAULT 'Hiring Team',
            focus TEXT,
            tech TEXT,
            location TEXT,
            role TEXT,
            source TEXT,
            status TEXT DEFAULT 'pending',
            batch TEXT,
            added_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )''')

        # ── Emails sent ─────────────────────────────────────────────────────
        c.execute('''CREATE TABLE IF NOT EXISTS emails_sent (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            company_id INTEGER,
            company_name TEXT NOT NULL,
            to_email TEXT NOT NULL,
            subject TEXT NOT NULL,
            body TEXT NOT NULL,
            email_type TEXT DEFAULT 'initial',
            ab_variant TEXT,
            gmail_message_id TEXT,
            gmail_thread_id TEXT,
            sent_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (company_id) REFERENCES companies(id)
        )''')

        # ── Bounces ──────────────────────────────────────────────────────────
        c.execute('''CREATE TABLE IF NOT EXISTS bounces (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            email_address TEXT UNIQUE NOT NULL,
            company_name TEXT,
            bounce_reason TEXT,
            detected_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )''')

        # ── Replies ──────────────────────────────────────────────────────────
        c.execute('''CREATE TABLE IF NOT EXISTS replies (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            from_email TEXT NOT NULL,
            company_name TEXT,
            gmail_message_id TEXT,
            gmail_thread_id TEXT,
            subject TEXT,
            snippet TEXT,
            category TEXT,
            ai_summary TEXT,
            detected_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )''')
        
        # Dynamically add referred_email column to replies if it doesn't exist
        columns = [col[1] for col in c.execute("PRAGMA table_info(replies)").fetchall()]
        if "referred_email" not in columns:
            c.execute("ALTER TABLE replies ADD COLUMN referred_email TEXT")

        # ── Follow-ups ──────────────────────────────────────────────────────
        c.execute('''CREATE TABLE IF NOT EXISTS followups (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            email_sent_id INTEGER,
            company_name TEXT NOT NULL,
            to_email TEXT NOT NULL,
            scheduled_date TEXT NOT NULL,
            status TEXT DEFAULT 'pending',
            sent_at TIMESTAMP,
            FOREIGN KEY (email_sent_id) REFERENCES emails_sent(id)
        )''')

        # ── A/B tests ───────────────────────────────────────────────────────
        c.execute('''CREATE TABLE IF NOT EXISTS ab_tests (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            company_name TEXT NOT NULL,
            variant_a TEXT NOT NULL,
            variant_b TEXT NOT NULL,
            chosen_variant TEXT NOT NULL,
            got_reply INTEGER DEFAULT 0,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )''')

        # ── Daily schedule log ───────────────────────────────────────────────
        c.execute('''CREATE TABLE IF NOT EXISTS daily_runs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            run_date TEXT NOT NULL,
            companies_sent INTEGER DEFAULT 0,
            followups_sent INTEGER DEFAULT 0,
            bounces_found INTEGER DEFAULT 0,
            replies_found INTEGER DEFAULT 0,
            status TEXT DEFAULT 'completed',
            log TEXT,
            started_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            finished_at TIMESTAMP
        )''')

        # ── Unsubscribes ────────────────────────────────────────────────────
        c.execute('''CREATE TABLE IF NOT EXISTS unsubscribes (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            email_address TEXT UNIQUE NOT NULL,
            company_name TEXT,
            unsubscribed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )''')

        # ── Failed sends (retry queue) ──────────────────────────────────────
        c.execute('''CREATE TABLE IF NOT EXISTS failed_sends (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            company_id INTEGER,
            company_name TEXT NOT NULL,
            to_email TEXT NOT NULL,
            subject TEXT,
            body TEXT,
            error_message TEXT,
            retry_count INTEGER DEFAULT 0,
            max_retries INTEGER DEFAULT 3,
            status TEXT DEFAULT 'pending',
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            last_retry_at TIMESTAMP,
            FOREIGN KEY (company_id) REFERENCES companies(id)
        )''')

        # ── Daily send count tracking ────────────────────────────────────────
        c.execute('''CREATE TABLE IF NOT EXISTS daily_send_counts (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            send_date TEXT NOT NULL,
            send_count INTEGER DEFAULT 0,
            UNIQUE(send_date)
        )''')

        conn.commit()
        conn.close()

    # ── Legacy methods (backward compat) ─────────────────────────────────────

    @sqlite_retry()
    def save_campaign(self, campaign_id: str, user_prompt: str, target_roles: List[str], company_type: str):
        with self._get_conn() as conn:
            roles_str = ",".join(target_roles)
            conn.execute(
                'INSERT INTO outreach_campaigns (campaign_id, user_prompt, target_roles, company_type) VALUES (?, ?, ?, ?)',
                (campaign_id, user_prompt, roles_str, company_type))

    @sqlite_retry()
    def save_outreach(self, campaign_id: str, company_name: str, company_domain: str,
                      contact_name: str, contact_email: str, targeted_role: str, email_draft: str):
        with self._get_conn() as conn:
            conn.execute('''INSERT INTO outreach_records
                (campaign_id, company_name, company_domain, contact_name, contact_email, targeted_role, email_draft)
                VALUES (?, ?, ?, ?, ?, ?, ?)''',
                (campaign_id, company_name, company_domain, contact_name, contact_email, targeted_role, email_draft))

    def get_campaign_records(self, campaign_id: str) -> List[Dict[str, Any]]:
        with self._get_conn() as conn:
            rows = conn.execute('''SELECT company_name, contact_name, contact_email, targeted_role, email_draft, status, created_at
                FROM outreach_records WHERE campaign_id = ? ORDER BY created_at DESC''', (campaign_id,)).fetchall()
            return [{"company": r[0], "contact": r[1], "email": r[2], "role": r[3],
                     "draft": r[4], "status": r[5], "created_at": r[6]} for r in rows]

    def update_status(self, record_id: int, status: str):
        with self._get_conn() as conn:
            conn.execute('UPDATE outreach_records SET status = ? WHERE id = ?', (status, record_id))

    # ── Company Pool ─────────────────────────────────────────────────────────

    @sqlite_retry()
    def add_company(self, name: str, email: str, focus: str = "", tech: list = None,
                    location: str = "", role: str = "", source: str = "", domain: str = "",
                    contact_name: str = "Hiring Team", batch: str = "") -> bool:
        """Add a company to the master pool. Returns True if added, False if duplicate."""
        try:
            with self._get_conn() as conn:
                conn.execute('''INSERT INTO companies (name, domain, email, contact_name, focus, tech, location, role, source, batch)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)''',
                    (name, domain, email, contact_name, focus, json.dumps(tech or []),
                     location, role, source, batch))
                return True
        except sqlite3.IntegrityError:
            return False

    def bulk_add_companies(self, companies: List[Dict]) -> int:
        """Add multiple companies. Returns count of newly added."""
        added = 0
        for c in companies:
            if self.add_company(
                name=c.get("company", c.get("name", "")),
                email=c.get("email", ""),
                focus=c.get("focus", ""),
                tech=c.get("tech", []),
                location=c.get("location", ""),
                role=c.get("role", ""),
                source=c.get("source", ""),
                domain=c.get("domain", ""),
                contact_name=c.get("contact", c.get("contact_name", "Hiring Team")),
                batch=c.get("batch", ""),
            ):
                added += 1
        return added

    def get_pending_companies(self, limit: int = 15) -> List[Dict]:
        """Get companies not yet contacted. Prioritizes tier-1 companies."""
        with self._get_conn() as conn:
            # Check if priority column exists
            cols = [r[1] for r in conn.execute("PRAGMA table_info(companies)").fetchall()]
            order = "priority ASC, RANDOM()" if "priority" in cols else "RANDOM()"
            # Prioritize companies with named contacts over generic addresses
            # CASE gives 0 to real contacts (sorted first), 1 to generic ones
            personal_contact_sort = """CASE WHEN contact_name NOT IN ('Hiring Team', 'Hiring Manager', 'Engineering Team', 'HR Manager', '') THEN 0 ELSE 1 END"""
            rows = conn.execute(f'''SELECT id, name, domain, email, contact_name, focus, tech, location, role, source
                FROM companies WHERE status = 'pending'
                AND email NOT IN (SELECT email_address FROM bounces)
                AND email NOT IN (SELECT email_address FROM unsubscribes)
                AND name NOT IN (SELECT DISTINCT company_name FROM emails_sent)
                AND email != '' AND email NOT LIKE '%unknown%'
                ORDER BY {personal_contact_sort}, {order} LIMIT ?''', (limit,)).fetchall()
            result = []
            for r in rows:
                result.append({
                    "id": r["id"], "company": r["name"], "domain": r["domain"],
                    "email": r["email"], "contact": r["contact_name"],
                    "focus": r["focus"], "tech": json.loads(r["tech"]) if r["tech"] else [],
                    "location": r["location"], "role": r["role"], "source": r["source"],
                })
            return result

    @sqlite_retry()
    def mark_company_status(self, company_id: int, status: str):
        with self._get_conn() as conn:
            conn.execute('UPDATE companies SET status = ? WHERE id = ?', (status, company_id))

    @sqlite_retry()
    def update_company_email(self, company_id: int, email: str, contact_name: str):
        with self._get_conn() as conn:
            conn.execute('UPDATE companies SET email = ?, contact_name = ? WHERE id = ?', (email, contact_name, company_id))

    @sqlite_retry()
    def mark_company_by_name(self, name: str, status: str):
        with self._get_conn() as conn:
            conn.execute('UPDATE companies SET status = ? WHERE name = ?', (status, name))

    def get_contacted_company_names(self) -> set:
        with self._get_conn() as conn:
            rows = conn.execute("SELECT DISTINCT name FROM companies WHERE status != 'pending'").fetchall()
            return {r["name"] for r in rows}

    def get_company_count_by_status(self) -> Dict[str, int]:
        with self._get_conn() as conn:
            rows = conn.execute("SELECT status, COUNT(*) as cnt FROM companies GROUP BY status").fetchall()
            return {r["status"]: r["cnt"] for r in rows}

    @sqlite_retry()
    def reset_old_contacts(self, days: int = 30) -> int:
        """
        Reset 'contacted' companies back to 'pending' if they were last emailed
        more than `days` ago and never replied. Allows re-contacting after a
        cooling-off period once the pool runs dry.

        Returns number of companies reset.
        """
        cutoff = (datetime.now() - timedelta(days=days)).strftime("%Y-%m-%d %H:%M:%S")
        with self._get_conn() as conn:
            # Find contacted companies whose most recent outgoing email is older than cutoff
            # and whose email is not in the replies table (they never wrote back)
            result = conn.execute('''
                UPDATE companies
                SET status = 'pending'
                WHERE status = 'contacted'
                  AND email != '' AND email NOT LIKE '%unknown%'
                  AND email NOT IN (SELECT email_address FROM bounces)
                  AND email NOT IN (SELECT email_address FROM unsubscribes)
                  AND email NOT IN (SELECT from_email FROM replies)
                  AND id IN (
                      SELECT c.id FROM companies c
                      LEFT JOIN emails_sent es ON es.company_id = c.id
                      GROUP BY c.id
                      HAVING MAX(es.sent_at) <= ? OR MAX(es.sent_at) IS NULL
                  )
            ''', (cutoff,))
            return result.rowcount

    # ── Emails Sent ──────────────────────────────────────────────────────────

    @sqlite_retry()
    def log_email_sent(self, company_name: str, to_email: str, subject: str, body: str,
                       email_type: str = "initial", ab_variant: str = None,
                       company_id: int = None, gmail_message_id: str = None,
                       gmail_thread_id: str = None) -> int:
        with self._get_conn() as conn:
            cur = conn.execute('''INSERT INTO emails_sent
                (company_id, company_name, to_email, subject, body, email_type, ab_variant, gmail_message_id, gmail_thread_id)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)''',
                (company_id, company_name, to_email, subject, body, email_type, ab_variant,
                 gmail_message_id, gmail_thread_id))
            return cur.lastrowid

    def get_emails_needing_followup(self, days_after: int = 5) -> List[Dict]:
        """Get initial emails that need follow-up (no reply, no followup sent, N+ days old)."""
        cutoff = (datetime.now() - timedelta(days=days_after)).strftime("%Y-%m-%d %H:%M:%S")
        with self._get_conn() as conn:
            rows = conn.execute('''
                SELECT es.id, es.company_name, es.to_email, es.subject, es.body, es.sent_at, c.contact_name
                FROM emails_sent es
                LEFT JOIN companies c ON es.company_id = c.id
                WHERE es.email_type = 'initial'
                  AND es.sent_at <= ?
                  AND es.to_email NOT IN (SELECT email_address FROM bounces)
                  AND es.to_email NOT IN (SELECT from_email FROM replies)
                  AND es.id NOT IN (SELECT email_sent_id FROM followups WHERE status = 'sent')
                ORDER BY es.sent_at ASC
            ''', (cutoff,)).fetchall()
            return [dict(r) for r in rows]

    def get_all_sent_emails(self) -> List[Dict]:
        with self._get_conn() as conn:
            rows = conn.execute('SELECT * FROM emails_sent ORDER BY sent_at DESC').fetchall()
            return [dict(r) for r in rows]

    def get_sent_email_addresses(self) -> set:
        with self._get_conn() as conn:
            rows = conn.execute('SELECT DISTINCT to_email FROM emails_sent').fetchall()
            return {r["to_email"] for r in rows}

    # ── Bounces ──────────────────────────────────────────────────────────────

    @sqlite_retry()
    def add_bounce(self, email_address: str, company_name: str = "", reason: str = "") -> bool:
        try:
            with self._get_conn() as conn:
                conn.execute('INSERT INTO bounces (email_address, company_name, bounce_reason) VALUES (?, ?, ?)',
                             (email_address.lower(), company_name, reason))
                conn.execute("UPDATE companies SET status = 'bounced' WHERE email = ?", (email_address.lower(),))
                return True
        except sqlite3.IntegrityError:
            return False

    def get_bounced_emails(self) -> set:
        with self._get_conn() as conn:
            rows = conn.execute('SELECT email_address FROM bounces').fetchall()
            return {r["email_address"] for r in rows}

    def get_bounce_count(self) -> int:
        with self._get_conn() as conn:
            return conn.execute('SELECT COUNT(*) FROM bounces').fetchone()[0]

    # ── Replies ──────────────────────────────────────────────────────────────

    @sqlite_retry()
    def add_reply(self, from_email: str, company_name: str = "", gmail_message_id: str = "",
                  gmail_thread_id: str = "", subject: str = "", snippet: str = "",
                  category: str = "unknown", ai_summary: str = "", referred_email: str = "") -> bool:
        with self._get_conn() as conn:
            existing = conn.execute('SELECT id FROM replies WHERE gmail_message_id = ?', (gmail_message_id,)).fetchone()
            if existing:
                return False
            conn.execute('''INSERT INTO replies
                (from_email, company_name, gmail_message_id, gmail_thread_id, subject, snippet, category, ai_summary, referred_email)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)''',
                (from_email, company_name, gmail_message_id, gmail_thread_id, subject, snippet, category, ai_summary, referred_email))
            
            if category == "referral" and referred_email and company_name:
                conn.execute("UPDATE companies SET status = 'referral_pending', email = ? WHERE name = ?", (referred_email.lower(), company_name))
            else:
                conn.execute("UPDATE companies SET status = 'replied' WHERE email = ?", (from_email.lower(),))
            return True

    def get_replies(self, category: str = None) -> List[Dict]:
        with self._get_conn() as conn:
            if category:
                rows = conn.execute('SELECT * FROM replies WHERE category = ? ORDER BY detected_at DESC', (category,)).fetchall()
            else:
                rows = conn.execute('SELECT * FROM replies ORDER BY detected_at DESC').fetchall()
            return [dict(r) for r in rows]

    def get_reply_stats(self) -> Dict[str, int]:
        with self._get_conn() as conn:
            rows = conn.execute('SELECT category, COUNT(*) as cnt FROM replies GROUP BY category').fetchall()
            return {r["category"]: r["cnt"] for r in rows}

    # ── Follow-ups ───────────────────────────────────────────────────────────

    @sqlite_retry()
    def schedule_followup(self, email_sent_id: int, company_name: str, to_email: str,
                          scheduled_date: str):
        with self._get_conn() as conn:
            conn.execute('''INSERT INTO followups (email_sent_id, company_name, to_email, scheduled_date)
                VALUES (?, ?, ?, ?)''', (email_sent_id, company_name, to_email, scheduled_date))

    def get_pending_followups(self, as_of_date: str = None) -> List[Dict]:
        if not as_of_date:
            as_of_date = datetime.now().strftime("%Y-%m-%d")
        with self._get_conn() as conn:
            rows = conn.execute('''SELECT f.*, es.subject as original_subject, es.body as original_body, c.contact_name
                FROM followups f
                LEFT JOIN emails_sent es ON f.email_sent_id = es.id
                LEFT JOIN companies c ON es.company_id = c.id
                WHERE f.status = 'pending' AND f.scheduled_date <= ?
                  AND f.to_email NOT IN (SELECT email_address FROM bounces)
                  AND f.to_email NOT IN (SELECT from_email FROM replies)
                  AND f.to_email NOT IN (SELECT email_address FROM unsubscribes)
            ''', (as_of_date,)).fetchall()
            return [dict(r) for r in rows]

    @sqlite_retry()
    def mark_followup_sent(self, followup_id: int):
        with self._get_conn() as conn:
            conn.execute("UPDATE followups SET status = 'sent', sent_at = ? WHERE id = ?",
                         (datetime.now().strftime("%Y-%m-%d %H:%M:%S"), followup_id))

    @sqlite_retry()
    def mark_followup_skipped(self, followup_id: int):
        with self._get_conn() as conn:
            conn.execute("UPDATE followups SET status = 'skipped' WHERE id = ?", (followup_id,))

    # ── A/B Tests ────────────────────────────────────────────────────────────

    @sqlite_retry()
    def log_ab_test(self, company_name: str, variant_a: str, variant_b: str, chosen: str):
        with self._get_conn() as conn:
            conn.execute('INSERT INTO ab_tests (company_name, variant_a, variant_b, chosen_variant) VALUES (?, ?, ?, ?)',
                         (company_name, variant_a, variant_b, chosen))

    @sqlite_retry()
    def mark_ab_reply(self, company_name: str):
        with self._get_conn() as conn:
            conn.execute('UPDATE ab_tests SET got_reply = 1 WHERE company_name = ?', (company_name,))

    def get_ab_stats(self) -> Dict:
        with self._get_conn() as conn:
            rows = conn.execute('''SELECT chosen_variant, COUNT(*) as total,
                SUM(got_reply) as replies FROM ab_tests GROUP BY chosen_variant''').fetchall()
            stats = {}
            for r in rows:
                stats[r["chosen_variant"]] = {"total": r["total"], "replies": r["replies"]}
            return stats

    # ── Daily Runs ───────────────────────────────────────────────────────────

    @sqlite_retry()
    def log_daily_run(self, companies_sent: int = 0, followups_sent: int = 0,
                      bounces_found: int = 0, replies_found: int = 0, log: str = "") -> int:
        run_date = datetime.now().strftime("%Y-%m-%d")
        with self._get_conn() as conn:
            cur = conn.execute('''INSERT INTO daily_runs
                (run_date, companies_sent, followups_sent, bounces_found, replies_found, log, finished_at)
                VALUES (?, ?, ?, ?, ?, ?, ?)''',
                (run_date, companies_sent, followups_sent, bounces_found, replies_found,
                 log, datetime.now().strftime("%Y-%m-%d %H:%M:%S")))
            return cur.lastrowid

    def already_ran_today(self) -> bool:
        today = datetime.now().strftime("%Y-%m-%d")
        with self._get_conn() as conn:
            row = conn.execute("SELECT id FROM daily_runs WHERE run_date = ? AND status = 'completed'", (today,)).fetchone()
            return row is not None

    def get_daily_runs(self, limit: int = 30) -> List[Dict]:
        with self._get_conn() as conn:
            rows = conn.execute('SELECT * FROM daily_runs ORDER BY started_at DESC LIMIT ?', (limit,)).fetchall()
            return [dict(r) for r in rows]

    # ── Analytics ────────────────────────────────────────────────────────────

    def get_analytics(self) -> Dict:
        with self._get_conn() as conn:
            total_sent = conn.execute('SELECT COUNT(*) FROM emails_sent').fetchone()[0]
            initial_sent = conn.execute("SELECT COUNT(*) FROM emails_sent WHERE email_type='initial'").fetchone()[0]
            followups_sent = conn.execute("SELECT COUNT(*) FROM emails_sent WHERE email_type='followup'").fetchone()[0]
            total_bounces = conn.execute('SELECT COUNT(*) FROM bounces').fetchone()[0]
            total_replies = conn.execute('SELECT COUNT(*) FROM replies').fetchone()[0]
            unique_companies = conn.execute('SELECT COUNT(DISTINCT company_name) FROM emails_sent').fetchone()[0]
            pending_companies = conn.execute("SELECT COUNT(*) FROM companies WHERE status='pending'").fetchone()[0]

            reply_cats = {}
            for row in conn.execute('SELECT category, COUNT(*) as cnt FROM replies GROUP BY category').fetchall():
                reply_cats[row["category"]] = row["cnt"]

            ab_stats = {}
            for row in conn.execute('SELECT chosen_variant, COUNT(*) as total, SUM(got_reply) as replies FROM ab_tests GROUP BY chosen_variant').fetchall():
                ab_stats[row["chosen_variant"]] = {"total": row["total"], "replies": row["replies"]}

            daily_sent = []
            for row in conn.execute('''SELECT DATE(sent_at) as day, COUNT(*) as cnt
                FROM emails_sent GROUP BY DATE(sent_at) ORDER BY day DESC LIMIT 30''').fetchall():
                daily_sent.append({"date": row["day"], "count": row["cnt"]})

            delivered = initial_sent - total_bounces
            response_rate = (total_replies / delivered * 100) if delivered > 0 else 0

            return {
                "total_emails_sent": total_sent,
                "initial_emails": initial_sent,
                "followups_sent": followups_sent,
                "total_bounces": total_bounces,
                "total_replies": total_replies,
                "unique_companies": unique_companies,
                "pending_companies": pending_companies,
                "delivered": max(0, delivered),
                "response_rate": round(response_rate, 1),
                "reply_categories": reply_cats,
                "ab_test_stats": ab_stats,
                "daily_sent": daily_sent,
            }

    # ── Unsubscribes ─────────────────────────────────────────────────────────

    @sqlite_retry()
    def add_unsubscribe(self, email_address: str, company_name: str = "") -> bool:
        try:
            with self._get_conn() as conn:
                conn.execute('INSERT INTO unsubscribes (email_address, company_name) VALUES (?, ?)',
                             (email_address.lower(), company_name))
                return True
        except sqlite3.IntegrityError:
            return False

    def is_unsubscribed(self, email_address: str) -> bool:
        with self._get_conn() as conn:
            row = conn.execute('SELECT id FROM unsubscribes WHERE email_address = ?',
                               (email_address.lower(),)).fetchone()
            return row is not None

    def get_unsubscribed_emails(self) -> set:
        with self._get_conn() as conn:
            rows = conn.execute('SELECT email_address FROM unsubscribes').fetchall()
            return {r["email_address"] for r in rows}

    # ── Failed Sends (Retry Queue) ───────────────────────────────────────────

    @sqlite_retry()
    def add_failed_send(self, company_name: str, to_email: str, subject: str = "",
                        body: str = "", error_message: str = "", company_id: int = None):
        with self._get_conn() as conn:
            conn.execute('''INSERT INTO failed_sends
                (company_id, company_name, to_email, subject, body, error_message)
                VALUES (?, ?, ?, ?, ?, ?)''',
                (company_id, company_name, to_email, subject, body, error_message))

    def get_pending_retries(self, limit: int = 10) -> List[Dict]:
        with self._get_conn() as conn:
            rows = conn.execute('''SELECT * FROM failed_sends
                WHERE status = 'pending' AND retry_count < max_retries
                ORDER BY created_at ASC LIMIT ?''', (limit,)).fetchall()
            return [dict(r) for r in rows]

    @sqlite_retry()
    def mark_retry_success(self, failed_id: int):
        with self._get_conn() as conn:
            conn.execute("UPDATE failed_sends SET status = 'sent', last_retry_at = ? WHERE id = ?",
                         (datetime.now().strftime("%Y-%m-%d %H:%M:%S"), failed_id))

    @sqlite_retry()
    def increment_retry(self, failed_id: int, error: str = ""):
        with self._get_conn() as conn:
            conn.execute("""UPDATE failed_sends
                SET retry_count = retry_count + 1, error_message = ?, last_retry_at = ?
                WHERE id = ?""",
                (error, datetime.now().strftime("%Y-%m-%d %H:%M:%S"), failed_id))
            # Mark as 'exhausted' if max retries reached
            conn.execute("""UPDATE failed_sends SET status = 'exhausted'
                WHERE id = ? AND retry_count >= max_retries""", (failed_id,))

    # ── Daily Send Count Tracking ────────────────────────────────────────────

    @sqlite_retry()
    def increment_daily_send_count(self, count: int = 1) -> int:
        """Atomically increment today's send count. Returns new total."""
        today = datetime.now().strftime("%Y-%m-%d")
        with self._get_conn() as conn:
            conn.execute('''INSERT INTO daily_send_counts (send_date, send_count)
                VALUES (?, ?) ON CONFLICT(send_date) DO UPDATE SET send_count = send_count + ?''',
                (today, count, count))
            row = conn.execute('SELECT send_count FROM daily_send_counts WHERE send_date = ?', (today,)).fetchone()
            return row["send_count"] if row else count

    @sqlite_retry()
    def decrement_daily_send_count(self, count: int = 1) -> int:
        """Decrement today's send count (for rollback on failed sends). Returns new total."""
        today = datetime.now().strftime("%Y-%m-%d")
        with self._get_conn() as conn:
            conn.execute('''UPDATE daily_send_counts SET send_count = MAX(0, send_count - ?)
                WHERE send_date = ?''', (count, today))
            row = conn.execute('SELECT send_count FROM daily_send_counts WHERE send_date = ?', (today,)).fetchone()
            return row["send_count"] if row else 0

    def get_daily_send_count(self) -> int:
        """Get today's total send count."""
        today = datetime.now().strftime("%Y-%m-%d")
        with self._get_conn() as conn:
            row = conn.execute('SELECT send_count FROM daily_send_counts WHERE send_date = ?', (today,)).fetchone()
            return row["send_count"] if row else 0

    # ── Bounce Recovery Tracking ─────────────────────────────────────────────

    def get_bounce_recovery_attempts(self, company_name: str) -> int:
        """Get number of bounce recovery attempts for a company."""
        with self._get_conn() as conn:
            row = conn.execute(
                "SELECT COUNT(*) as cnt FROM emails_sent WHERE company_name = ? AND email_type = 'bounce_recovery'",
                (company_name,)
            ).fetchone()
            return row["cnt"] if row else 0