database.py

import sqlite3
from datetime import datetime
from typing import Optional, List, Tuple, Dict
import config


class Database:
    def __init__(self, db_path: str = None):
        self.db_path = db_path or config.DATABASE_PATH
        self.init_db()

    def init_db(self):
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        cursor.execute("""
                       CREATE TABLE IF NOT EXISTS accounts
                       (
                           id
                           INTEGER
                           PRIMARY
                           KEY
                           AUTOINCREMENT,
                           email
                           TEXT
                           NOT
                           NULL,
                           password
                           TEXT
                           NOT
                           NULL,
                           status
                           TEXT
                           DEFAULT
                           'pending',
                           refresh_token
                           TEXT,
                           created_at
                           TEXT
                           DEFAULT
                           CURRENT_TIMESTAMP,
                           processed_at
                           TEXT,
                           skip_verification
                           INTEGER
                           DEFAULT
                           0,
                           error_message
                           TEXT
                       )
                       """)

        conn.commit()
        conn.close()

    def add_account(self, email: str, password: str):
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        cursor.execute(
            "INSERT INTO accounts (email, password) VALUES (?, ?)",
            (email, password)
        )
        conn.commit()
        account_id = cursor.lastrowid
        conn.close()
        return account_id

    def update_or_add_account(self, email: str, password: str):
        """如果邮箱存在则更新密码并重置状态，否则新增"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        # 检查邮箱是否存在
        cursor.execute("SELECT id FROM accounts WHERE email = ?", (email,))
        result = cursor.fetchone()

        if result:
            # 邮箱存在，更新密码并重置状态
            account_id = result[0]
            cursor.execute("""
                           UPDATE accounts
                           SET password      = ?,
                               status        = 'pending',
                               refresh_token = NULL,
                               processed_at  = NULL
                           WHERE email = ?
                           """, (password, email))
            conn.commit()
            conn.close()
            return account_id, True  # True表示是更新
        else:
            # 邮箱不存在，新增
            cursor.execute(
                "INSERT INTO accounts (email, password) VALUES (?, ?)",
                (email, password)
            )
            conn.commit()
            account_id = cursor.lastrowid
            conn.close()
            return account_id, False  # False表示是新增

    def get_pending_account(self) -> Optional[Tuple[int, str, str]]:
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        cursor.execute(
            "SELECT id, email, password FROM accounts WHERE status = 'pending' AND skip_verification = 0 LIMIT 1"
        )
        result = cursor.fetchone()
        conn.close()
        return result

    def update_account_status(self, account_id: int, status: str, refresh_token: Optional[str] = None):
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        cursor.execute(
            "UPDATE accounts SET status = ?, refresh_token = ?, processed_at = ? WHERE id = ?",
            (status, refresh_token, datetime.now().isoformat(), account_id)
        )
        conn.commit()
        conn.close()

    def cleanup_processed_accounts(self) -> int:
        """删除所有已处理的账号，只保留pending状态"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        # 删除非pending状态的记录
        cursor.execute("DELETE FROM accounts WHERE status != 'pending'")
        deleted_count = cursor.rowcount

        conn.commit()
        conn.close()

        return deleted_count

    def get_pending_count(self) -> int:
        """获取待处理账号数量"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        cursor.execute("SELECT COUNT(*) FROM accounts WHERE status = 'pending' AND skip_verification = 0")
        count = cursor.fetchone()[0]
        conn.close()
        return count

    def mark_skip_verification_and_error(self, account_id: int, error_message: Optional[str] = None):
        """标记账号跳过验证码处理并设置为错误状态
        
        Args:
            account_id: 账号ID
            error_message: 错误信息（可选）
        """
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        cursor.execute(
            "UPDATE accounts SET skip_verification = 1, status = ?, processed_at = ?, error_message = ? WHERE id = ?",
            ("error", datetime.now().isoformat(), error_message, account_id)
        )
        conn.commit()
        conn.close()

    def batch_reset_accounts_by_emails(self, emails: List[str]) -> int:
        """批量重置账号状态（根据邮箱列表，只更新未标记跳过验证码的账号）
        
        Args:
            emails: 邮箱列表
            
        Returns:
            更新的记录数
        """
        if not emails:
            return 0

        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        # 使用IN子句批量更新，只更新 skip_verification = 0 的账号
        placeholders = ','.join(['?' for _ in emails])
        cursor.execute(f"""
            UPDATE accounts 
            SET status = 'pending', 
                refresh_token = NULL, 
                processed_at = NULL
            WHERE email IN ({placeholders}) AND skip_verification = 0
        """, emails)

        updated_count = cursor.rowcount
        conn.commit()
        conn.close()

        return updated_count

    def batch_reset_accounts_by_patterns(self, patterns: List[str]) -> int:
        """根据邮箱或邮箱后缀批量重置账号状态
        
        支持两种模式：
        1. 完整邮箱匹配：如 'user@example.com'
        2. 域名后缀匹配：如 'example.com' 会匹配所有 @example.com 的邮箱
        
        Args:
            patterns: 邮箱或域名后缀列表
            
        Returns:
            更新的记录数
        """
        if not patterns:
            return 0

        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        # 构建WHERE条件
        conditions = []
        params = []
        
        for pattern in patterns:
            if '@' in pattern:
                # 完整邮箱匹配
                conditions.append("email = ?")
                params.append(pattern)
            else:
                # 域名后缀匹配
                conditions.append("email LIKE ?")
                params.append(f"%@{pattern}")
        
        if not conditions:
            conn.close()
            return 0
        
        # 构建SQL查询
        where_clause = "(" + " OR ".join(conditions) + ") AND skip_verification = 0"
        
        cursor.execute(f"""
            UPDATE accounts 
            SET status = 'pending', 
                refresh_token = NULL, 
                processed_at = NULL
            WHERE {where_clause}
        """, params)

        updated_count = cursor.rowcount
        conn.commit()
        conn.close()

        return updated_count

    def query_accounts(self, search: Optional[str] = None, page: int = 1, page_size: int = 20) -> Dict:
        """查询邮箱信息（支持分页和搜索）
        
        Args:
            search: 搜索关键词（匹配邮箱）
            page: 页码（从1开始）
            page_size: 每页数量
            
        Returns:
            包含accounts列表和total总数的字典
        """
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 构建查询条件
        where_clause = ""
        params = []
        
        if search:
            where_clause = "WHERE email LIKE ?"
            params.append(f"%{search}%")
        
        # 查询总数
        count_query = f"SELECT COUNT(*) FROM accounts {where_clause}"
        cursor.execute(count_query, params)
        total = cursor.fetchone()[0]
        
        # 查询分页数据
        offset = (page - 1) * page_size
        query = f"""
            SELECT id, email, password, status, created_at, processed_at, skip_verification, error_message
            FROM accounts
            {where_clause}
            ORDER BY CASE WHEN processed_at IS NULL THEN 1 ELSE 0 END, processed_at ASC, id ASC
            LIMIT ? OFFSET ?
        """
        cursor.execute(query, params + [page_size, offset])
        rows = cursor.fetchall()
        
        accounts = []
        for row in rows:
            accounts.append({
                "id": row[0],
                "email": row[1],
                "password": row[2],
                "status": row[3],
                "created_at": row[4],
                "processed_at": row[5],
                "skip_verification": bool(row[6]),
                "error_message": row[7]
            })
        
        conn.close()
        
        return {
            "accounts": accounts,
            "total": total,
            "page": page,
            "page_size": page_size,
            "total_pages": (total + page_size - 1) // page_size
        }