autoChatGPTPlus-Paypal/storage.py

"""SQLite 持久化：注册成功的账号、session 快照、CPA 上传记录。"""
from __future__ import annotations

import json
import os
import sqlite3
import threading
import time
from contextlib import contextmanager
from typing import Any, Iterable

DATA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
DB_PATH = os.path.join(DATA_DIR, "accounts.db")
_LOCK = threading.Lock()
TRIAL_ELIGIBLE_NOTE = "trial_eligible"
TRIAL_INELIGIBLE_NOTE = "trial_ineligible"
_TRIAL_EVENT_PATTERNS = ("%免费试用%", "%trial%")


SCHEMA = """
CREATE TABLE IF NOT EXISTS accounts (
    email TEXT PRIMARY KEY,
    password TEXT NOT NULL,
    created_at INTEGER NOT NULL,
    updated_at INTEGER NOT NULL,
    plan_type TEXT,
    final_status TEXT,         -- registered/paid/plus/cpa_uploaded/failed/stopped
    last_error TEXT,
    long_link TEXT,
    cpa_file_name TEXT,
    cpa_uploaded_at INTEGER,
    initial_session_json TEXT, -- 注册成功时拉到的 /api/auth/session
    plus_session_json TEXT,    -- 付款后拉到的 plus session
    notes TEXT
);

CREATE INDEX IF NOT EXISTS idx_accounts_status ON accounts(final_status);
CREATE INDEX IF NOT EXISTS idx_accounts_created ON accounts(created_at);

CREATE TABLE IF NOT EXISTS account_events (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    email TEXT NOT NULL,
    ts INTEGER NOT NULL,
    stage TEXT NOT NULL,
    status TEXT NOT NULL,        -- info/ok/warn/error
    detail TEXT,
    payload_json TEXT
);

CREATE INDEX IF NOT EXISTS idx_events_email ON account_events(email);
CREATE INDEX IF NOT EXISTS idx_events_ts ON account_events(ts);

CREATE TABLE IF NOT EXISTS tasks (
    task_id TEXT PRIMARY KEY,
    mode TEXT NOT NULL,           -- full | pay_only
    status TEXT NOT NULL,         -- queued | running | success | failed | cancelled
    stage TEXT,
    attempts INTEGER DEFAULT 0,
    max_attempts INTEGER DEFAULT 3,
    params_json TEXT,
    result_json TEXT,
    last_error TEXT,
    email TEXT,
    plan_type TEXT,
    cpa_file_name TEXT,
    created_at INTEGER NOT NULL,
    updated_at INTEGER NOT NULL,
    started_at INTEGER,
    finished_at INTEGER
);

CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status);
CREATE INDEX IF NOT EXISTS idx_tasks_created ON tasks(created_at);
"""


def _now_ms() -> int:
    return int(time.time() * 1000)


def _ensure_dir():
    os.makedirs(DATA_DIR, exist_ok=True)


def init_db():
    _ensure_dir()
    with _LOCK:
        conn = sqlite3.connect(DB_PATH)
        try:
            conn.executescript(SCHEMA)
            conn.commit()
        finally:
            conn.close()


@contextmanager
def _conn():
    _ensure_dir()
    with _LOCK:
        c = sqlite3.connect(DB_PATH)
        c.row_factory = sqlite3.Row
        try:
            yield c
            c.commit()
        finally:
            c.close()


def _dump(value: Any) -> str | None:
    if value is None:
        return None
    try:
        return json.dumps(value, ensure_ascii=False)
    except Exception:
        return str(value)


def upsert_account(email: str, password: str, *, fields: dict | None = None) -> dict:
    """Insert or update by email. fields 中只更新非 None 字段。"""
    init_db()
    fields = dict(fields or {})
    now = _now_ms()
    with _conn() as c:
        row = c.execute("SELECT email FROM accounts WHERE email = ?", (email,)).fetchone()
        if row is None:
            c.execute(
                """
                INSERT INTO accounts (email, password, created_at, updated_at, plan_type,
                  final_status, last_error, long_link, cpa_file_name, cpa_uploaded_at,
                  initial_session_json, plus_session_json, notes)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    email, password, now, now,
                    fields.get("plan_type"),
                    fields.get("final_status") or "registered",
                    fields.get("last_error"),
                    fields.get("long_link"),
                    fields.get("cpa_file_name"),
                    fields.get("cpa_uploaded_at"),
                    _dump(fields.get("initial_session")) if "initial_session" in fields else fields.get("initial_session_json"),
                    _dump(fields.get("plus_session")) if "plus_session" in fields else fields.get("plus_session_json"),
                    fields.get("notes"),
                ),
            )
        else:
            sets = ["updated_at = ?"]
            args: list[Any] = [now]
            for col in ("plan_type", "final_status", "last_error", "long_link",
                        "cpa_file_name", "cpa_uploaded_at", "notes"):
                if col in fields and fields[col] is not None:
                    sets.append(f"{col} = ?")
                    args.append(fields[col])
            if "initial_session" in fields:
                sets.append("initial_session_json = ?")
                args.append(_dump(fields["initial_session"]))
            elif "initial_session_json" in fields and fields["initial_session_json"] is not None:
                sets.append("initial_session_json = ?")
                args.append(fields["initial_session_json"])
            if "plus_session" in fields:
                sets.append("plus_session_json = ?")
                args.append(_dump(fields["plus_session"]))
            elif "plus_session_json" in fields and fields["plus_session_json"] is not None:
                sets.append("plus_session_json = ?")
                args.append(fields["plus_session_json"])
            if password:
                sets.append("password = ?")
                args.append(password)
            args.append(email)
            c.execute(f"UPDATE accounts SET {', '.join(sets)} WHERE email = ?", args)

        return _row_to_dict(c.execute("SELECT * FROM accounts WHERE email = ?", (email,)).fetchone())


def add_event(email: str, stage: str, status: str = "info",
              detail: str | None = None, payload: Any = None) -> int:
    init_db()
    with _conn() as c:
        cur = c.execute(
            "INSERT INTO account_events (email, ts, stage, status, detail, payload_json) VALUES (?, ?, ?, ?, ?, ?)",
            (email or "", _now_ms(), stage, status, detail, _dump(payload)),
        )
        return cur.lastrowid


def list_accounts(limit: int = 200, status: str | None = None, offset: int = 0) -> list[dict]:
    init_db()
    with _conn() as c:
        if status == "trial":
            rows = c.execute(
                "SELECT * FROM accounts ORDER BY created_at DESC",
            ).fetchall()
        elif status:
            rows = c.execute(
                "SELECT * FROM accounts WHERE final_status = ? ORDER BY created_at DESC LIMIT ? OFFSET ?",
                (status, limit, offset),
            ).fetchall()
        else:
            rows = c.execute(
                "SELECT * FROM accounts ORDER BY created_at DESC LIMIT ? OFFSET ?",
                (limit, offset),
            ).fetchall()
        accounts = _annotate_accounts([_row_to_dict(r) for r in rows], c)
        if status == "trial":
            accounts = [account for account in accounts if account_matches_status(account, "trial")]
            return accounts[offset:offset + limit]
        return accounts


def count_accounts(status: str | None = None) -> int:
    init_db()
    with _conn() as c:
        if status == "trial":
            rows = c.execute("SELECT * FROM accounts").fetchall()
            accounts = _annotate_accounts([_row_to_dict(r) for r in rows], c)
            return sum(1 for account in accounts if account_matches_status(account, "trial"))
        if status:
            row = c.execute("SELECT COUNT(*) FROM accounts WHERE final_status = ?", (status,)).fetchone()
        else:
            row = c.execute("SELECT COUNT(*) FROM accounts").fetchone()
        return row[0] if row else 0


def get_account(email: str) -> dict | None:
    init_db()
    with _conn() as c:
        r = c.execute("SELECT * FROM accounts WHERE email = ?", (email,)).fetchone()
        if not r:
            return None
        return _annotate_accounts([_row_to_dict(r)], c)[0]


def list_events(email: str, limit: int = 100) -> list[dict]:
    init_db()
    with _conn() as c:
        rows = c.execute(
            "SELECT * FROM account_events WHERE email = ? ORDER BY ts DESC LIMIT ?",
            (email, limit),
        ).fetchall()
        return [_event_row(r) for r in rows]


def _row_to_dict(row: sqlite3.Row | None) -> dict | None:
    if row is None:
        return None
    d = {k: row[k] for k in row.keys()}
    # session JSON 反序列化但保留 raw 副本
    for k in ("initial_session_json", "plus_session_json"):
        raw = d.get(k)
        if raw:
            try:
                d[k.replace("_json", "")] = json.loads(raw)
            except Exception:
                d[k.replace("_json", "")] = None
    return d


def account_matches_status(account: dict | None, status: str | None) -> bool:
    if not account:
        return False
    if not status:
        return True
    if status == "trial":
        return bool(account.get("is_trial_account") or (account.get("final_status") or "") == "trial")
    return (account.get("final_status") or "") == status


def derive_account_flags(account: dict | None, *, has_trial_event: bool = False) -> dict | None:
    if account is None:
        return None

    annotated = dict(account)
    final_status = str(annotated.get("final_status") or "")
    note_flags = _parse_note_flags(annotated.get("notes"))
    trial_state = "unknown"
    if TRIAL_ELIGIBLE_NOTE in note_flags:
        trial_state = "eligible"
    elif TRIAL_INELIGIBLE_NOTE in note_flags:
        trial_state = "ineligible"
    elif has_trial_event:
        trial_state = "eligible"

    is_trial_account = trial_state == "eligible"

    annotated["trial_state"] = trial_state
    annotated["trial_eligible"] = is_trial_account
    annotated["is_trial_account"] = is_trial_account
    annotated["can_retry_payment"] = bool(_session_access_token(annotated) and (is_trial_account or final_status == "trial"))
    return annotated


def set_account_trial_eligibility(email: str, password: str, eligible: bool) -> dict | None:
    existing = get_account(email) or {}
    notes = _update_note_flags(
        existing.get("notes"),
        add=[TRIAL_ELIGIBLE_NOTE if eligible else TRIAL_INELIGIBLE_NOTE],
        remove=[TRIAL_INELIGIBLE_NOTE if eligible else TRIAL_ELIGIBLE_NOTE],
    )
    return upsert_account(email, password, fields={"notes": notes})


def _annotate_accounts(accounts: Iterable[dict | None], conn: sqlite3.Connection) -> list[dict]:
    items = [account for account in accounts if account]
    if not items:
        return []

    trial_emails = _load_trial_event_emails(conn, [str(account.get("email") or "") for account in items])
    out: list[dict] = []
    for account in items:
        annotated = derive_account_flags(account, has_trial_event=str(account.get("email") or "") in trial_emails)
        if annotated:
            out.append(annotated)
    return out


def _load_trial_event_emails(conn: sqlite3.Connection, emails: Iterable[str]) -> set[str]:
    email_list = [email for email in emails if email]
    if not email_list:
        return set()
    placeholders = ",".join("?" for _ in email_list)
    rows = conn.execute(
        f"""
        SELECT DISTINCT email
        FROM account_events
        WHERE email IN ({placeholders})
          AND stage = 'paypal'
          AND detail IS NOT NULL
          AND (detail LIKE ? OR detail LIKE ?)
        """,
        [*email_list, *_TRIAL_EVENT_PATTERNS],
    ).fetchall()
    return {str(row[0]) for row in rows}


def _session_access_token(account: dict) -> str:
    for key in ("plus_session", "initial_session"):
        session = account.get(key)
        if isinstance(session, dict):
            token = session.get("accessToken")
            if token:
                return str(token)
    return ""


def _parse_note_flags(notes: Any) -> set[str]:
    raw = str(notes or "").replace(",", " ")
    return {part.strip() for part in raw.split() if part.strip()}


def _update_note_flags(notes: Any, *, add: Iterable[str] = (), remove: Iterable[str] = ()) -> str | None:
    flags = _parse_note_flags(notes)
    flags.difference_update({flag for flag in remove if flag})
    flags.update({flag for flag in add if flag})
    if not flags:
        return None
    return " ".join(sorted(flags))


def _event_row(row: sqlite3.Row) -> dict:
    d = {k: row[k] for k in row.keys()}
    raw = d.get("payload_json")
    if raw:
        try:
            d["payload"] = json.loads(raw)
        except Exception:
            d["payload"] = None
    return d


# ------------------------- tasks -------------------------

def create_task(task_id: str, mode: str, params: dict, max_attempts: int = 3) -> dict:
    init_db()
    now = _now_ms()
    with _conn() as c:
        c.execute(
            """
            INSERT INTO tasks (task_id, mode, status, stage, attempts, max_attempts,
                params_json, result_json, last_error, email, plan_type, cpa_file_name,
                created_at, updated_at, started_at, finished_at)
            VALUES (?, ?, 'queued', '', 0, ?, ?, NULL, NULL, NULL, NULL, NULL, ?, ?, NULL, NULL)
            """,
            (task_id, mode, max_attempts, _dump(params or {}), now, now),
        )
        return _task_row(c.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,)).fetchone())


def update_task(task_id: str, fields: dict) -> dict | None:
    init_db()
    if not fields:
        return get_task(task_id)
    sets = ["updated_at = ?"]
    args: list[Any] = [_now_ms()]
    for col in (
        "status", "stage", "attempts", "max_attempts", "last_error",
        "email", "plan_type", "cpa_file_name", "started_at", "finished_at",
    ):
        if col in fields and fields[col] is not None:
            sets.append(f"{col} = ?")
            args.append(fields[col])
    if "result" in fields:
        sets.append("result_json = ?")
        args.append(_dump(fields["result"]))
    elif "result_json" in fields and fields["result_json"] is not None:
        sets.append("result_json = ?")
        args.append(fields["result_json"])
    if "params" in fields:
        sets.append("params_json = ?")
        args.append(_dump(fields["params"]))
    args.append(task_id)
    with _conn() as c:
        c.execute(f"UPDATE tasks SET {', '.join(sets)} WHERE task_id = ?", args)
        row = c.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,)).fetchone()
        return _task_row(row) if row else None


def get_task(task_id: str) -> dict | None:
    init_db()
    with _conn() as c:
        r = c.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,)).fetchone()
        return _task_row(r) if r else None


def list_tasks(limit: int = 100, status: str | None = None) -> list[dict]:
    init_db()
    with _conn() as c:
        if status:
            rows = c.execute(
                "SELECT * FROM tasks WHERE status = ? ORDER BY created_at DESC LIMIT ?",
                (status, limit),
            ).fetchall()
        else:
            rows = c.execute(
                "SELECT * FROM tasks ORDER BY created_at DESC LIMIT ?",
                (limit,),
            ).fetchall()
        return [_task_row(r) for r in rows]


def count_tasks_by_mode_statuses(mode: str, statuses: tuple[str, ...]) -> int:
    init_db()
    status_values = tuple(str(status) for status in statuses if status)
    if not status_values:
        return 0
    placeholders = ",".join("?" for _ in status_values)
    with _conn() as c:
        row = c.execute(
            f"SELECT COUNT(*) FROM tasks WHERE mode = ? AND status IN ({placeholders})",
            (mode, *status_values),
        ).fetchone()
        return int(row[0]) if row else 0


def get_next_queued_task() -> dict | None:
    """取一个最早的 queued 任务（FIFO）。"""
    init_db()
    with _conn() as c:
        r = c.execute(
            "SELECT * FROM tasks WHERE status = 'queued' ORDER BY created_at ASC LIMIT 1"
        ).fetchone()
        return _task_row(r) if r else None


def _task_row(row: sqlite3.Row | None) -> dict | None:
    if row is None:
        return None
    d = {k: row[k] for k in row.keys()}
    for k in ("params_json", "result_json"):
        raw = d.get(k)
        if raw:
            try:
                d[k.replace("_json", "")] = json.loads(raw)
            except Exception:
                d[k.replace("_json", "")] = None
    return d