database.py

import os
import time
import random
from datetime import datetime
from urllib.parse import urlparse, urlunparse

import settings
import psycopg2
import psycopg2.extras

RAW_DATABASE_URL = settings.DATABASE_URL
_DB_OFFLINE_CACHE = False
_DB_LAST_CHECK = 0.0

def optimize_neon_url(url: str | None) -> str | None:
    """Ensure Neon PostgreSQL connection URL uses the pooled connection string (-pooler.neon.tech) for serverless efficiency."""
    if not url or "neon.tech" not in url:
        return url
    if "-pooler" in url:
        return url
    try:
        parsed = urlparse(url)
        if parsed.hostname and "neon.tech" in parsed.hostname:
            host_parts = parsed.hostname.split(".")
            host_parts[0] = f"{host_parts[0]}-pooler"
            new_host = ".".join(host_parts)
            user_pass = f"{parsed.username}:{parsed.password}" if parsed.password else (parsed.username or "")
            netloc_auth = f"{user_pass}@" if user_pass else ""
            port_str = f":{parsed.port}" if parsed.port else ""
            new_netloc = f"{netloc_auth}{new_host}{port_str}"
            return urlunparse((parsed.scheme, new_netloc, parsed.path, parsed.params, parsed.query, parsed.fragment))
    except Exception:
        pass
    return url

DATABASE_URL = optimize_neon_url(RAW_DATABASE_URL)

def mark_db_offline():
    global _DB_OFFLINE_CACHE, _DB_LAST_CHECK
    _DB_OFFLINE_CACHE = True
    _DB_LAST_CHECK = time.time()

def is_db_offline() -> bool:
    global _DB_OFFLINE_CACHE
    # Return cached status without opening a connection on passive page renders.
    # Real DB operations update _DB_OFFLINE_CACHE when connections fail or succeed.
    return _DB_OFFLINE_CACHE

def get_db_connection(max_retries=1, retry_delay=0.5):
    """
    Connect to PostgreSQL with quick retry for Neon cold starts.
    Includes connect_timeout=5 to prevent hanging connections.
    """
    if not DATABASE_URL:
        raise ValueError("DATABASE_URL is not set.")

    for attempt in range(max_retries):
        try:
            conn = psycopg2.connect(DATABASE_URL, connect_timeout=5)
            _DB_OFFLINE_CACHE = False
            return conn
        except Exception as e:
            mark_db_offline()
            if attempt < max_retries - 1:
                time.sleep(retry_delay)
            else:
                raise e

def get_db_cursor(conn):
    return conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)

def insert_and_get_id(conn, query, params):
    cursor = conn.cursor()
    if "RETURNING id" not in query.upper():
        query += " RETURNING id"
    cursor.execute(query, params)
    last_id = cursor.fetchone()[0]
    return last_id

def init_db():
    # Cleanup local SQLite leftover database file if present
    sqlite_file = "d:/crew/experiment/sysex_vault.db"
    if os.path.exists(sqlite_file):
        try:
            os.remove(sqlite_file)
        except Exception as e:
            print(f"Failed to delete SQLite leftovers: {e}")
            
    # Cleanup unused, leftover templates
    deprecated_templates = [
        "d:/crew/experiment/templates/changelog.html",
        "d:/crew/experiment/templates/about.html",
        "d:/crew/experiment/templates/library.html",
        "d:/crew/experiment/templates/roadmap.html"
    ]
    for temp_path in deprecated_templates:
        if os.path.exists(temp_path):
            try:
                os.remove(temp_path)
            except Exception as e:
                print(f"Failed to delete leftover template {temp_path}: {e}")

    try:
        conn = get_db_connection()
    except Exception as e:
        print(f"Warning: Skipping init_db because database is offline/suspended: {e}")
        return

    try:
        cursor = conn.cursor()
        
        # Create users
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS users (
            id SERIAL PRIMARY KEY,
            email VARCHAR(255) UNIQUE NOT NULL,
            hashed_password VARCHAR(255) NOT NULL,
            tier VARCHAR(50) NOT NULL DEFAULT 'free',
            stripe_customer_id VARCHAR(255),
            created_at VARCHAR(100) NOT NULL
        )
        """)
        # Create banks
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS banks (
            id SERIAL PRIMARY KEY,
            user_id INTEGER REFERENCES users(id) ON DELETE CASCADE,
            name VARCHAR(255) NOT NULL,
            synth_model VARCHAR(255) NOT NULL,
            sysex_hex TEXT NOT NULL,
            created_at VARCHAR(100) NOT NULL
        )
        """)
        # Create patches
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS patches (
            id SERIAL PRIMARY KEY,
            bank_id INTEGER REFERENCES banks(id) ON DELETE CASCADE,
            patch_index INTEGER NOT NULL,
            name VARCHAR(255) NOT NULL
        )
        """)
        # Create subscribers
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS subscribers (
            id SERIAL PRIMARY KEY,
            email VARCHAR(255) UNIQUE NOT NULL,
            created_at VARCHAR(100) NOT NULL
        )
        """)
        # Pending premiums: holds paid-but-not-yet-registered users
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS pending_premiums (
            id SERIAL PRIMARY KEY,
            email VARCHAR(255) UNIQUE NOT NULL,
            stripe_customer_id VARCHAR(255),
            created_at VARCHAR(100) NOT NULL,
            plan VARCHAR(50) DEFAULT 'personal'
        )
        """)
        
        # Create unsubscribed_emails table for tracking opt-outs
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS unsubscribed_emails (
            id SERIAL PRIMARY KEY,
            email VARCHAR(255) UNIQUE NOT NULL,
            unsubscribed_at VARCHAR(100) NOT NULL
        )
        """)
        
        # Create rate_limits table for atomic distributed rate limiting across serverless instances
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS rate_limits (
            key VARCHAR(128) PRIMARY KEY,
            count INT NOT NULL DEFAULT 1,
            window_start TIMESTAMPTZ NOT NULL DEFAULT NOW()
        );
        CREATE INDEX IF NOT EXISTS idx_rate_limits_window ON rate_limits (window_start);
        """)

        # Create user_sessions table for stateful session validation and instant revocation
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS user_sessions (
            id SERIAL PRIMARY KEY,
            user_id INTEGER REFERENCES users(id) ON DELETE CASCADE,
            session_token VARCHAR(64) UNIQUE NOT NULL,
            created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
            last_seen_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
        );
        CREATE INDEX IF NOT EXISTS idx_user_sessions_token ON user_sessions (session_token);
        CREATE INDEX IF NOT EXISTS idx_user_sessions_user_id ON user_sessions (user_id);
        CREATE INDEX IF NOT EXISTS idx_banks_user_id ON banks (user_id);
        CREATE INDEX IF NOT EXISTS idx_patches_bank_id ON patches (bank_id);
        """)
        
        # Add drip_email_sent, reengagement_email_sent, and last_marketing_email_sent tracking to users table
        cursor.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS login_count INTEGER DEFAULT 1;")
        cursor.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS last_login_at VARCHAR(100);")
        cursor.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS drip_email_sent BOOLEAN DEFAULT FALSE;")
        cursor.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS reengagement_email_sent BOOLEAN DEFAULT FALSE;")
        cursor.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS abandoned_checkout_sent BOOLEAN DEFAULT FALSE;")
        cursor.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS last_checkout_initiated_at VARCHAR(100);")
        cursor.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS last_marketing_email_sent VARCHAR(100);")
        cursor.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS marketing_email_count INTEGER DEFAULT 0;")
        cursor.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS plan VARCHAR(50);")
        cursor.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS picture_url TEXT;")
        cursor.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS onboarding_completed BOOLEAN DEFAULT FALSE;")
        cursor.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS synths_owned TEXT DEFAULT '';")
        cursor.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS workflow_goal TEXT DEFAULT '';")
        cursor.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS inactivity_warning_sent BOOLEAN DEFAULT FALSE;")
        cursor.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS is_archived BOOLEAN DEFAULT FALSE;")
        cursor.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS archived_at VARCHAR(100);")
        cursor.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS second_login_email_sent BOOLEAN DEFAULT FALSE;")
        cursor.execute("ALTER TABLE pending_premiums ADD COLUMN IF NOT EXISTS plan VARCHAR(50) DEFAULT 'personal';")
        cursor.execute("ALTER TABLE banks ADD COLUMN IF NOT EXISTS share_token VARCHAR(255) UNIQUE;")
        
        # Absorb mobile studio-link leads & newsletter subscribers into users table as free leads
        cursor.execute("""
            INSERT INTO users (email, hashed_password, tier, created_at)
            SELECT email, 'studio_link_lead', 'free', created_at
            FROM subscribers
            ON CONFLICT (email) DO NOTHING;
        """)
        
        conn.commit()
    except Exception as e:
        print(f"Warning: Database error during table migrations: {e}")
    finally:
        try:
            conn.close()
        except Exception:
            pass

# --- User Table operations ---
def create_user(email: str, hashed_password: str) -> int:
    try:
        conn = get_db_connection()
        created_at = datetime.now().isoformat()
        try:
            user_id = insert_and_get_id(
                conn,
                "INSERT INTO users (email, hashed_password, tier, created_at) VALUES (%s, %s, 'free', %s)",
                (email.lower().strip(), hashed_password, created_at)
            )
            conn.commit()
            return user_id
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in create_user: {e}")
        raise RuntimeError("Database is undergoing monthly quota reset. Registration will resume tomorrow!")

def save_user_onboarding(email: str, synths_owned: str = "", workflow_goal: str = "") -> bool:
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            cursor.execute(
                """
                UPDATE users
                SET onboarding_completed = TRUE,
                    synths_owned = %s,
                    workflow_goal = %s
                WHERE email = %s
                """,
                (synths_owned, workflow_goal, email.lower().strip())
            )
            conn.commit()
            return True
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in save_user_onboarding: {e}")
        return False

def get_user_by_email(email: str) -> dict | None:
    try:
        conn = get_db_connection()
        try:
            cursor = get_db_cursor(conn)
            cursor.execute("SELECT * FROM users WHERE email = %s", (email.lower().strip(),))
            row = cursor.fetchone()
            return dict(row) if row else None
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in get_user_by_email: {e}")
        return None

def get_user_by_id(user_id: int) -> dict | None:
    try:
        conn = get_db_connection()
        try:
            cursor = get_db_cursor(conn)
            cursor.execute("SELECT * FROM users WHERE id = %s", (user_id,))
            row = cursor.fetchone()
            return dict(row) if row else None
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in get_user_by_id: {e}")
        return None

def get_user_by_customer_id(stripe_customer_id: str) -> dict | None:
    try:
        conn = get_db_connection()
        try:
            cursor = get_db_cursor(conn)
            cursor.execute("SELECT * FROM users WHERE stripe_customer_id = %s", (stripe_customer_id,))
            row = cursor.fetchone()
            return dict(row) if row else None
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in get_user_by_customer_id: {e}")
        return None

def update_user_tier(email: str, tier: str, stripe_customer_id: str = None, plan: str = None):
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            email = email.lower().strip()
            if stripe_customer_id and plan:
                query = "UPDATE users SET tier = %s, stripe_customer_id = %s, plan = %s WHERE email = %s"
                params = (tier, stripe_customer_id, plan, email)
            elif stripe_customer_id:
                query = "UPDATE users SET tier = %s, stripe_customer_id = %s WHERE email = %s"
                params = (tier, stripe_customer_id, email)
            elif plan:
                query = "UPDATE users SET tier = %s, plan = %s WHERE email = %s"
                params = (tier, plan, email)
            else:
                query = "UPDATE users SET tier = %s WHERE email = %s"
                params = (tier, email)
                
            cursor.execute(query, params)
            conn.commit()
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in update_user_tier: {e}")

def update_user_picture(email: str, picture_url: str):
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("UPDATE users SET picture_url = %s WHERE email = %s", (picture_url, email.lower().strip()))
            conn.commit()
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in update_user_picture: {e}")
        return None

def update_user_password(email: str, hashed_password: str):
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("UPDATE users SET hashed_password = %s WHERE email = %s", (hashed_password, email.lower().strip()))
            conn.commit()
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in update_user_password: {e}")
        return None

def delete_user(user_id: int) -> bool:
    """Permanently delete user and all associated soundbank vaults."""
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("DELETE FROM banks WHERE user_id = %s", (user_id,))
            cursor.execute("DELETE FROM users WHERE id = %s", (user_id,))
            conn.commit()
            return True
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in delete_user: {e}")
        return False

def update_user_tier_by_customer_id(stripe_customer_id: str, tier: str):
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("UPDATE users SET tier = %s WHERE stripe_customer_id = %s", (tier, stripe_customer_id))
            conn.commit()
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in update_user_tier_by_customer_id: {e}")
        return None

def upsert_pending_premium(email: str, stripe_customer_id: str = None, plan: str = "personal"):
    """Park a premium grant for an email that hasn't registered yet."""
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            created_at = datetime.now().isoformat()
            cursor.execute(
                """
                INSERT INTO pending_premiums (email, stripe_customer_id, created_at, plan)
                VALUES (%s, %s, %s, %s)
                ON CONFLICT (email) DO UPDATE SET
                    stripe_customer_id = EXCLUDED.stripe_customer_id,
                    plan = EXCLUDED.plan
                """,
                (email.lower().strip(), stripe_customer_id, created_at, plan)
            )
            conn.commit()
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in upsert_pending_premium: {e}")
        return None

def consume_pending_premium(email: str) -> dict | None:
    """Check if there's a pending premium for this email. Returns the record and deletes it."""
    try:
        conn = get_db_connection()
        try:
            cursor = get_db_cursor(conn)
            cursor.execute("SELECT * FROM pending_premiums WHERE email = %s", (email.lower().strip(),))
            row = cursor.fetchone()
            if row:
                dc = conn.cursor()
                dc.execute("DELETE FROM pending_premiums WHERE email = %s", (email.lower().strip(),))
                conn.commit()
            return dict(row) if row else None
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in consume_pending_premium: {e}")
        return None

def get_pending_premium(email: str) -> dict | None:
    try:
        conn = get_db_connection()
        try:
            cursor = get_db_cursor(conn)
            cursor.execute("SELECT * FROM pending_premiums WHERE email = %s", (email.lower().strip(),))
            row = cursor.fetchone()
            return dict(row) if row else None
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in get_pending_premium: {e}")
        return None

def delete_pending_premium(email: str) -> bool:
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("DELETE FROM pending_premiums WHERE email = %s", (email.lower().strip(),))
            conn.commit()
            return True
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in delete_pending_premium: {e}")
        return False

# --- Bank & Patch scoped operations ---
def save_bank(name: str, synth_model: str, sysex_hex: str, patch_names: list[str], user_id: int) -> int | None:
    try:
        conn = get_db_connection()
        try:
            created_at = datetime.now().isoformat()
            
            bank_id = insert_and_get_id(
                conn,
                "INSERT INTO banks (name, synth_model, sysex_hex, created_at, user_id) VALUES (%s, %s, %s, %s, %s)",
                (name, synth_model, sysex_hex, created_at, user_id)
            )
            
            cursor = conn.cursor()
            for idx, patch_name in enumerate(patch_names):
                cursor.execute(
                    "INSERT INTO patches (bank_id, patch_index, name) VALUES (%s, %s, %s)",
                    (bank_id, idx, patch_name)
                )
                
            conn.commit()
            return bank_id
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in save_bank: {e}")
        return None

def get_all_banks(user_id: int) -> list[dict]:
    try:
        conn = get_db_connection()
        try:
            cursor = get_db_cursor(conn)
            cursor.execute("SELECT * FROM banks WHERE user_id = %s ORDER BY id DESC", (user_id,))
            rows = cursor.fetchall()
            
            banks = []
            for row in rows:
                cursor.execute("SELECT COUNT(*) as count FROM patches WHERE bank_id = %s", (row['id'],))
                count = cursor.fetchone()['count']
                
                dt = datetime.fromisoformat(row['created_at'])
                formatted_date = dt.strftime("%b %d, %Y %I:%M %p")
                
                banks.append({
                    "id": row['id'],
                    "name": row['name'],
                    "synth_model": row['synth_model'],
                    "sysex_hex": row['sysex_hex'],
                    "created_at": formatted_date,
                    "patch_count": count
                })
                
            return banks
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in get_all_banks: {e}")
        return []

def get_bank(bank_id: int, user_id: int) -> dict | None:
    try:
        conn = get_db_connection()
        try:
            cursor = get_db_cursor(conn)
            cursor.execute("SELECT * FROM banks WHERE id = %s AND user_id = %s", (bank_id, user_id))
            row = cursor.fetchone()
            if not row:
                return None
                
            cursor.execute("SELECT * FROM patches WHERE bank_id = %s ORDER BY patch_index ASC", (bank_id,))
            patch_rows = cursor.fetchall()
            
            dt = datetime.fromisoformat(row['created_at'])
            formatted_date = dt.strftime("%b %d, %Y %I:%M %p")
            
            return {
                "id": row['id'],
                "name": row['name'],
                "synth_model": row['synth_model'],
                "sysex_hex": row['sysex_hex'],
                "created_at": formatted_date,
                "patches": [{"index": p['patch_index'], "name": p['name']} for p in patch_rows]
            }
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in get_bank: {e}")
        return None

def generate_bank_share_token(bank_id: int, user_id: int) -> str | None:
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("SELECT share_token FROM banks WHERE id = %s AND user_id = %s", (bank_id, user_id))
            row = cursor.fetchone()
            if not row:
                return None
            if row[0]:
                return row[0]
            
            import secrets
            token = secrets.token_urlsafe(16)
            cursor.execute("UPDATE banks SET share_token = %s WHERE id = %s AND user_id = %s", (token, bank_id, user_id))
            conn.commit()
            return token
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in generate_bank_share_token: {e}")
        return None

def get_bank_by_share_token(token: str) -> dict | None:
    try:
        conn = get_db_connection()
        try:
            cursor = get_db_cursor(conn)
            cursor.execute("SELECT * FROM banks WHERE share_token = %s", (token,))
            row = cursor.fetchone()
            if not row:
                return None
                
            bank_id = row['id']
            cursor.execute("SELECT * FROM patches WHERE bank_id = %s ORDER BY patch_index ASC", (bank_id,))
            patch_rows = cursor.fetchall()
            
            dt = datetime.fromisoformat(row['created_at'])
            formatted_date = dt.strftime("%b %d, %Y %I:%M %p")
            
            return {
                "id": row['id'],
                "name": row['name'],
                "synth_model": row['synth_model'],
                "sysex_hex": row['sysex_hex'],
                "created_at": formatted_date,
                "patches": [{"index": p['patch_index'], "name": p['name']} for p in patch_rows],
                "user_id": row['user_id']
            }
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in get_bank_by_share_token: {e}")
        return None

def delete_bank(bank_id: int, user_id: int):
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            
            cursor.execute("SELECT id FROM banks WHERE id = %s AND user_id = %s", (bank_id, user_id))
            if not cursor.fetchone():
                return
                
            cursor.execute("DELETE FROM banks WHERE id = %s", (bank_id,))
            cursor.execute("DELETE FROM patches WHERE bank_id = %s", (bank_id,))
            conn.commit()
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in delete_bank: {e}")
        return

def create_subscriber(email: str) -> bool:
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            created_at = datetime.now().isoformat()
            email_clean = email.lower().strip()
            cursor.execute(
                "INSERT INTO subscribers (email, created_at) VALUES (%s, %s) ON CONFLICT (email) DO NOTHING",
                (email_clean, created_at)
            )
            # Automatically absorb into users table as a free lead to feed the 20-campaign engine!
            cursor.execute(
                """
                INSERT INTO users (email, hashed_password, tier, created_at)
                VALUES (%s, 'studio_link_lead', 'free', %s)
                ON CONFLICT (email) DO NOTHING
                """,
                (email_clean, created_at)
            )
            # Automatically remove from unsubscribed_emails so resubscribed users receive emails again
            cursor.execute("DELETE FROM unsubscribed_emails WHERE email = %s", (email_clean,))
            conn.commit()
            return True
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in create_subscriber: {e}")
        return True

def get_all_users() -> list[dict]:
    try:
        conn = get_db_connection()
        try:
            cursor = get_db_cursor(conn)
            cursor.execute("SELECT * FROM users")
            rows = cursor.fetchall()
            return [dict(r) for r in rows] if rows else []
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in get_all_users: {e}")
        return []

def get_pending_drip_users() -> list[dict]:
    try:
        conn = get_db_connection()
        try:
            cursor = get_db_cursor(conn)
            cursor.execute("SELECT id, email, created_at FROM users WHERE tier = 'free' AND drip_email_sent = FALSE")
            rows = cursor.fetchall()
            return rows
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in get_pending_drip_users: {e}")
        return []

def mark_drip_sent(user_id: int):
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("UPDATE users SET drip_email_sent = TRUE WHERE id = %s", (user_id,))
            conn.commit()
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in mark_drip_sent: {e}")
        return

def get_pending_reengagement_users() -> list[dict]:
    try:
        conn = get_db_connection()
        try:
            cursor = get_db_cursor(conn)
            cursor.execute("SELECT id, email, created_at FROM users WHERE reengagement_email_sent = FALSE")
            rows = cursor.fetchall()
            return rows
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in get_pending_reengagement_users: {e}")
        return []

def mark_reengagement_sent(user_id: int):
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("UPDATE users SET reengagement_email_sent = TRUE WHERE id = %s", (user_id,))
            conn.commit()
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in mark_reengagement_sent: {e}")
        return

def mark_abandoned_checkout_sent(user_id: int):
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("UPDATE users SET abandoned_checkout_sent = TRUE WHERE id = %s", (user_id,))
            conn.commit()
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in mark_abandoned_checkout_sent: {e}")
        return

def record_checkout_initiated(user_id: int):
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            from datetime import datetime
            now_str = datetime.now().isoformat()
            cursor.execute("""
                UPDATE users 
                SET last_checkout_initiated_at = %s, abandoned_checkout_sent = FALSE 
                WHERE id = %s
            """, (now_str, user_id))
            conn.commit()
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in record_checkout_initiated: {e}")
        return

def get_pending_marketing_users(days_interval: int = 2, max_emails: int = 20) -> list[dict]:
    try:
        conn = get_db_connection()
        try:
            cursor = get_db_cursor(conn)
            # Find free tier users who have never received a marketing email or received one older than days_interval.
            # Capped at max_emails (default 20: exactly 1 full loop of the 20-campaign engine) to ensure 100% fresh content and zero repetition.
            query = """
                SELECT id, email, created_at, last_marketing_email_sent, COALESCE(marketing_email_count, 0) AS marketing_email_count 
                FROM users 
                WHERE (tier = 'free' OR tier IS NULL)
                  AND (last_marketing_email_sent IS NULL 
                       OR (CAST(last_marketing_email_sent AS TIMESTAMP) < NOW() - INTERVAL '%s days'))
                  AND email NOT IN (SELECT email FROM unsubscribed_emails)
                  AND COALESCE(marketing_email_count, 0) < %s
            """
            cursor.execute(query % (int(days_interval), int(max_emails)))
            rows = cursor.fetchall()
            return dict(rows) if isinstance(rows, dict) else [dict(r) for r in rows]
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in get_pending_marketing_users: {e}")
        return []

def mark_marketing_sent(user_id: int):
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            now_iso = datetime.now().isoformat()
            cursor.execute("""
                UPDATE users 
                SET last_marketing_email_sent = %s,
                    marketing_email_count = COALESCE(marketing_email_count, 0) + 1 
                WHERE id = %s
            """, (now_iso, user_id))
            conn.commit()
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in mark_marketing_sent: {e}")
        return

def get_pending_second_login_users() -> list[dict]:
    try:
        conn = get_db_connection()
        try:
            cursor = get_db_cursor(conn)
            query = """
                SELECT id, email, created_at, login_count, last_login_at 
                FROM users 
                WHERE (tier = 'free' OR tier IS NULL)
                  AND COALESCE(login_count, 0) >= 2
                  AND (second_login_email_sent IS FALSE OR second_login_email_sent IS NULL)
                  AND email NOT IN (SELECT email FROM unsubscribed_emails)
            """
            cursor.execute(query)
            rows = cursor.fetchall()
            return [dict(r) for r in rows]
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in get_pending_second_login_users: {e}")
        return []

def mark_second_login_sent(user_id: int):
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("UPDATE users SET second_login_email_sent = TRUE WHERE id = %s", (user_id,))
            conn.commit()
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in mark_second_login_sent: {e}")
        return

def add_to_unsubscribed(email: str):
    clean_email = email.lower().strip() if email else ""
    if not clean_email or "{{" in clean_email or "%7b" in clean_email or "@" not in clean_email or "." not in clean_email:
        return
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            unsubscribed_at = datetime.now().isoformat()
            cursor.execute(
                "INSERT INTO unsubscribed_emails (email, unsubscribed_at) VALUES (%s, %s) ON CONFLICT (email) DO NOTHING",
                (clean_email, unsubscribed_at)
            )
            conn.commit()
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in add_to_unsubscribed: {e}")
        return

def remove_from_unsubscribed(email: str) -> bool:
    clean_email = email.lower().strip() if email else ""
    if not clean_email or "@" not in clean_email or "." not in clean_email:
        return False
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("DELETE FROM unsubscribed_emails WHERE email = %s", (clean_email,))
            conn.commit()
            return True
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in remove_from_unsubscribed: {e}")
        return False

def is_unsubscribed(email: str) -> bool:
    try:
        conn = get_db_connection()
        try:
            cursor = get_db_cursor(conn)
            cursor.execute("SELECT 1 FROM unsubscribed_emails WHERE email = %s", (email.lower().strip(),))
            return cursor.fetchone() is not None
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in is_unsubscribed: {e}")
        return False

def get_all_newsletter_recipients() -> list[str]:
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            # Fetch all user emails
            cursor.execute("SELECT email FROM users")
            user_emails = [row[0].lower().strip() for row in cursor.fetchall()]
            
            # Fetch all newsletter subscribers
            cursor.execute("SELECT email FROM subscribers")
            sub_emails = [row[0].lower().strip() for row in cursor.fetchall()]
            
            # Fetch all unsubscribed emails
            cursor.execute("SELECT email FROM unsubscribed_emails")
            unsub_emails = set(row[0].lower().strip() for row in cursor.fetchall())
            
            # Combine, de-duplicate and filter
            all_emails = set(user_emails + sub_emails)
            valid_recipients = [email for email in all_emails if email and email not in unsub_emails]
            return valid_recipients
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in get_all_newsletter_recipients: {e}")
        return []

def record_user_login(user_id: int):
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            now_str = datetime.now().isoformat()
            cursor.execute(
                """
                UPDATE users 
                SET login_count = COALESCE(login_count, 0) + 1, 
                    last_login_at = %s, 
                    inactivity_warning_sent = FALSE, 
                    is_archived = FALSE, 
                    archived_at = NULL 
                WHERE id = %s
                RETURNING id, email, tier, login_count, second_login_email_sent
                """,
                (now_str, user_id)
            )
            row = cursor.fetchone()
            conn.commit()

            if row:
                u_id, email, tier, login_count, second_sent = row
                if (tier or 'free').lower() == 'free' and login_count == 2 and not second_sent:
                    import threading
                    def _async_second_login_trigger():
                        try:
                            if is_unsubscribed(email):
                                return
                            from scripts.send_second_login_nudge import (
                                render_second_login_template, generate_plain_text,
                                extract_first_name, send_resend, SMTP_FROM
                            )
                            u = {"id": u_id, "email": email}
                            first_name = extract_first_name(email)
                            html = render_second_login_template(u)
                            text = generate_plain_text(first_name, email)
                            ok, _ = send_resend(
                                to=email,
                                subject="quick question about your synth setup (from Max)",
                                body=text,
                                html=html,
                                from_addr=SMTP_FROM,
                                reply_to="support@bipluk.com"
                            )
                            if ok:
                                mark_second_login_sent(u_id)
                        except Exception as ex:
                            print(f"Error in automatic second-login dispatch: {ex}")
                    threading.Thread(target=_async_second_login_trigger, daemon=True).start()
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in record_user_login: {e}")

def get_icp_user_segmentation() -> dict[str, list[dict]]:
    """Analyzes PostgreSQL database and groups users into actionable ICP segments."""
    segments: dict[str, list[dict]] = {
        "LURKER_ZERO_USAGE": [],       # Registered, 0 banks, <= 2 logins (High churn risk -> needs demo bank)
        "ENGAGED_FREE_USER": [],       # 1 bank uploaded, >= 2 logins, Free tier (Prime upgrade candidate!)
        "BOUNCED_ONETIMER": [],        # 1 login, 0 banks, registered > 48h ago (Bounced prospect)
        "HAPPY_CUSTOMERS": [],         # Paid plan (Personal, Studio, VIP)
    }
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
            cursor.execute("""
                SELECT u.id, u.email, u.tier, u.plan, u.created_at, u.login_count, u.last_login_at,
                       COALESCE(b.bank_count, 0) as bank_count
                FROM users u
                LEFT JOIN (
                    SELECT user_id, COUNT(*) as bank_count FROM banks GROUP BY user_id
                ) b ON u.id = b.user_id
                ORDER BY u.created_at DESC;
            """)
            users = cursor.fetchall()
            now = datetime.now()
            for u in users:
                user_data = dict(u)
                tier = (u.get("tier") or "free").lower()
                plan = (u.get("plan") or "").lower()
                bank_count = u.get("bank_count", 0)
                login_count = u.get("login_count") or 1

                created_at_raw = str(u.get("created_at", ""))
                try:
                    created_at = datetime.fromisoformat(created_at_raw)
                    age_hours = (now - created_at).total_seconds() / 3600
                except Exception:
                    age_hours = 0

                if tier == "premium" or tier != "free":
                    segments["HAPPY_CUSTOMERS"].append(user_data)
                elif bank_count >= 1:
                    segments["ENGAGED_FREE_USER"].append(user_data)
                elif login_count == 1 and age_hours >= 48:
                    segments["BOUNCED_ONETIMER"].append(user_data)
                else:
                    segments["LURKER_ZERO_USAGE"].append(user_data)
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in get_icp_user_segmentation: {e}")
    return segments


def check_and_increment_rate_limit(key: str, max_requests: int, window_seconds: int) -> tuple[bool, int]:
    """
    Atomic distributed rate limiter backed by PostgreSQL.
    Returns (is_allowed, retry_after_seconds).
    Uses row-level locking on conflict to ensure atomic increment across all serverless Lambdas/ASGI workers.
    """
    conn = get_db_connection(max_retries=1, retry_delay=0.2)
    try:
        cursor = conn.cursor()

        # Opportunistic cleanup of expired windows (approx 1 in 50 calls)
        if random.random() < 0.02:
            try:
                cursor.execute("DELETE FROM rate_limits WHERE window_start < NOW() - INTERVAL '1 hour';")
                conn.commit()
            except Exception:
                pass

        query = """
        INSERT INTO rate_limits (key, count, window_start)
        VALUES (%s, 1, NOW())
        ON CONFLICT (key) DO UPDATE
        SET count = CASE
            WHEN rate_limits.window_start < NOW() - (%s || ' seconds')::INTERVAL THEN 1
            ELSE rate_limits.count + 1
        END,
        window_start = CASE
            WHEN rate_limits.window_start < NOW() - (%s || ' seconds')::INTERVAL THEN NOW()
            ELSE rate_limits.window_start
        END
        RETURNING count, GREATEST(1, CEIL(EXTRACT(EPOCH FROM (window_start + (%s || ' seconds')::INTERVAL - NOW()))))::INT AS retry_after;
        """
        cursor.execute(query, (key, str(window_seconds), str(window_seconds), str(window_seconds)))
        row = cursor.fetchone()
        conn.commit()
        if row:
            curr_count, retry_after = row
            retry_after = max(1, retry_after if retry_after is not None else window_seconds)
            return (curr_count <= max_requests, retry_after)
        return (True, 0)
    finally:
        try:
            conn.close()
        except Exception:
            pass


def create_user_session(user_id: int, session_token: str) -> bool:
    """Register a new active session token for a user in PostgreSQL."""
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            cursor.execute(
                "INSERT INTO user_sessions (user_id, session_token) VALUES (%s, %s) ON CONFLICT (session_token) DO NOTHING;",
                (user_id, session_token)
            )
            conn.commit()
            return True
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in create_user_session: {e}")
        return False


def get_user_by_session(email: str, session_token: str) -> dict | None:
    """Validate that a session token exists in PostgreSQL and matches the user email."""
    if not email or not session_token:
        return None
    try:
        conn = get_db_connection()
        try:
            cursor = get_db_cursor(conn)
            cursor.execute(
                """
                SELECT u.* FROM users u
                JOIN user_sessions s ON u.id = s.user_id
                WHERE s.session_token = %s AND LOWER(u.email) = LOWER(%s);
                """,
                (session_token, email.strip())
            )
            row = cursor.fetchone()
            return dict(row) if row else None
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in get_user_by_session: {e}")
        return None


def delete_user_session(session_token: str) -> bool:
    """Delete a single session token on logout."""
    if not session_token:
        return False
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("DELETE FROM user_sessions WHERE session_token = %s;", (session_token,))
            conn.commit()
            return True
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in delete_user_session: {e}")
        return False


def delete_all_user_sessions(user_id: int) -> bool:
    """Revoke all sessions for a user (called on password reset or account suspension)."""
    if not user_id:
        return False
    try:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("DELETE FROM user_sessions WHERE user_id = %s;", (user_id,))
            conn.commit()
            return True
        finally:
            conn.close()
    except Exception as e:
        print(f"Database error in delete_all_user_sessions: {e}")
        return False