services/register_service.py

from __future__ import annotations

import json
import threading
import time
import uuid
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
from datetime import datetime, timezone
from pathlib import Path

from services.account_service import account_service
from services import codex_register_service
from services.config import DATA_DIR
from services.register import openai_register


REGISTER_FILE = DATA_DIR / "register.json"
HERO_SMS_MAX_WAIT_TIMEOUT = 45
HERO_SMS_MIN_WAIT_TIMEOUT = 45
HERO_SMS_DEFAULT_SEND_RETRY_ATTEMPTS = 3
HERO_SMS_MAX_SEND_RETRY_ATTEMPTS = 5


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _default_config() -> dict:
    return {
        **openai_register.config,
        "mode": "total",
        "target_quota": 100,
        "target_available": 10,
        "check_interval": 5,
        "enabled": False,
        "stats": {
            "success": 0,
            "fail": 0,
            "done": 0,
            "running": 0,
            "threads": openai_register.config["threads"],
            "elapsed_seconds": 0,
            "avg_seconds": 0,
            "success_rate": 0,
            "current_quota": 0,
            "current_available": 0,
        },
    }


def _positive_int(value, default: int) -> int:
    try:
        parsed = int(value)
    except Exception:
        return default
    return parsed if parsed > 0 else default


def _positive_float(value, default: float) -> float:
    try:
        parsed = float(value)
    except Exception:
        return default
    return parsed if parsed > 0 else default


def _non_negative_float(value, default: float = 0.0) -> float:
    try:
        parsed = float(value)
    except Exception:
        return default
    return parsed if parsed >= 0 else default


def _positive_int_list(value, default: list[int]) -> list[int]:
    items = []

    def add(raw) -> None:
        try:
            parsed = int(raw)
        except Exception:
            return
        if parsed > 0 and parsed not in items:
            items.append(parsed)

    if isinstance(value, str):
        for item in value.replace(";", ",").replace("\n", ",").split(","):
            add(item.strip())
    elif isinstance(value, (list, tuple)):
        for item in value:
            add(item)
    return items or list(default)


def _normalize_hero_sms(raw: object) -> dict:
    defaults = dict(getattr(openai_register, "default_hero_sms_config", {}) or {})
    source = raw if isinstance(raw, dict) else {}
    service = str(source.get("service") or defaults.get("service") or "dr").strip().lower() or "dr"
    operator = str(source.get("operator") or defaults.get("operator") or "any").strip().lower() or "any"
    default_country_blacklist = _positive_int_list(defaults.get("country_blacklist"), [16, 10, 4])
    country_blacklist = _positive_int_list(source.get("country_blacklist"), default_country_blacklist)
    country_blacklist = list(dict.fromkeys([*default_country_blacklist, *country_blacklist]))
    default_country_pool = _positive_int_list(defaults.get("country_pool"), [6, 117, 31, 33, 2, 39, 48, 37, 13, 40, 15, 8, 129, 32, 86, 173, 43, 49, 34, 7, 85, 27, 172, 63, 56, 177, 54, 24, 1, 46, 175, 14, 67, 83, 59, 187, 36])
    default_country_pool = [item for item in default_country_pool if item not in country_blacklist] or [6, 117, 31, 33, 2, 39, 48, 37, 13, 40, 15, 8, 129, 32, 86, 173, 43, 49, 34, 7, 85, 27, 172, 63, 56, 177, 54, 24, 1, 46, 175, 14, 67, 83, 59, 187, 36]
    country = _positive_int(source.get("country"), int(defaults.get("country") or default_country_pool[0]))
    country_pool = _positive_int_list(source.get("country_pool"), default_country_pool)
    if "country_pool" not in source and "country" in source:
        country_pool = [country, *[item for item in country_pool if item != country]]
    country_pool = [item for item in country_pool if item not in country_blacklist]
    if not country_pool:
        country_pool = [item for item in default_country_pool if item not in country_blacklist] or list(default_country_pool)
    if country in country_blacklist or country not in country_pool:
        country = country_pool[0]
    max_price_usd = _positive_float(source.get("max_price_usd"), float(defaults.get("max_price_usd") or 0.1))
    min_price_usd = _non_negative_float(source.get("min_price_usd"), float(defaults.get("min_price_usd") or 0.045))
    min_price_usd = min(min_price_usd, max_price_usd)
    return {
        "enabled": bool(source.get("enabled") if "enabled" in source else defaults.get("enabled", False)),
        "api_key": str(source.get("api_key") or defaults.get("api_key") or "").strip(),
        "service": service,
        "country": country,
        "country_pool": country_pool,
        "country_blacklist": country_blacklist,
        "operator": operator,
        "wait_timeout": max(
            HERO_SMS_MIN_WAIT_TIMEOUT,
            min(
                HERO_SMS_MAX_WAIT_TIMEOUT,
                _positive_int(source.get("wait_timeout"), int(defaults.get("wait_timeout") or HERO_SMS_MAX_WAIT_TIMEOUT)),
            ),
        ),
        "poll_interval": min(5, _positive_int(source.get("poll_interval"), int(defaults.get("poll_interval") or 5))),
        "reuse_activation_id": "",
        "reuse_phone": "",
        "auto_buy": True,
        "min_price_usd": min_price_usd,
        "max_price_usd": max_price_usd,
        "send_retry_attempts": min(
            HERO_SMS_MAX_SEND_RETRY_ATTEMPTS,
            _positive_int(source.get("send_retry_attempts"), int(defaults.get("send_retry_attempts") or HERO_SMS_DEFAULT_SEND_RETRY_ATTEMPTS)),
        ),
        "cancel_on_send_fail": True,
    }


def _normalize(raw: dict) -> dict:
    cfg = _default_config()
    cfg.update({k: v for k, v in raw.items() if k not in {"stats", "logs"}})
    cfg["total"] = max(1, int(cfg.get("total") or 1))
    cfg["threads"] = max(1, int(cfg.get("threads") or 1))
    cfg["mode"] = str(cfg.get("mode") or "total").strip() if str(cfg.get("mode") or "total").strip() in {"total", "quota", "available"} else "total"
    cfg["target_quota"] = max(1, int(cfg.get("target_quota") or 1))
    cfg["target_available"] = max(1, int(cfg.get("target_available") or 1))
    cfg["check_interval"] = max(1, int(cfg.get("check_interval") or 5))
    cfg["proxy"] = str(cfg.get("proxy") or "").strip()
    cfg["hero_sms"] = _normalize_hero_sms(cfg.get("hero_sms"))
    cfg["enabled"] = bool(cfg.get("enabled"))
    cfg.pop("cpa_auto_import", None)
    stats = {**_default_config()["stats"], **(raw.get("stats") if isinstance(raw.get("stats"), dict) else {}),
             "threads": cfg["threads"]}
    cfg["stats"] = stats
    return cfg


def _merge_config_update(current: dict, updates: dict) -> dict:
    merged = {**current, **updates}
    for key in ("mail", "hero_sms"):
        if isinstance(current.get(key), dict) and isinstance(updates.get(key), dict):
            merged[key] = {**current[key], **updates[key]}
    return merged


class RegisterService:
    def __init__(self, store_file: Path):
        self._store_file = store_file
        self._lock = threading.RLock()
        self._runner: threading.Thread | None = None
        self._logs: list[dict] = []
        openai_register.register_log_sink = self._append_log
        self._config = self._load()
        if self._config["enabled"]:
            self.start()

    def _load(self) -> dict:
        try:
            return _normalize(json.loads(self._store_file.read_text(encoding="utf-8")))
        except Exception:
            return _normalize({})

    def _save(self) -> None:
        self._store_file.parent.mkdir(parents=True, exist_ok=True)
        self._store_file.write_text(json.dumps(self._config, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

    def get(self) -> dict:
        with self._lock:
            return json.loads(json.dumps({**self._config, "logs": self._logs[-300:]}, ensure_ascii=False))

    def update(self, updates: dict) -> dict:
        with self._lock:
            self._config = _normalize(_merge_config_update(self._config, updates))
            openai_register.config.update({k: self._config[k] for k in ("mail", "proxy", "total", "threads", "hero_sms")})
            self._save()
            return self.get()

    def start(self) -> dict:
        with self._lock:
            if self._runner and self._runner.is_alive():
                self._config["enabled"] = True
                self._save()
                return self.get()
            self._config["enabled"] = True
            self._logs = []
            metrics = self._pool_metrics()
            self._config["stats"] = {"job_id": uuid.uuid4().hex, "success": 0, "fail": 0, "done": 0, "running": 0, "threads": self._config["threads"], **metrics, "started_at": _now(), "updated_at": _now()}
            openai_register.config.update({k: self._config[k] for k in ("mail", "proxy", "total", "threads", "hero_sms")})
            with openai_register.stats_lock:
                openai_register.stats.update({"done": 0, "success": 0, "fail": 0, "start_time": time.time()})
            self._save()
            self._runner = threading.Thread(target=self._run, daemon=True, name="openai-register")
            self._runner.start()
            self._append_log(f"注册任务启动，模式={self._config['mode']}，线程数={self._config['threads']}", "yellow")
            return self.get()

    def start_codex(self) -> dict:
        with self._lock:
            if self._runner and self._runner.is_alive():
                self._config["enabled"] = True
                self._save()
                return self.get()
            self._config["enabled"] = True
            self._logs = []
            self._config["stats"] = {
                "job_id": uuid.uuid4().hex,
                "task_type": "codex",
                "success": 0,
                "fail": 0,
                "done": 0,
                "running": 0,
                "threads": self._config["threads"],
                **self._pool_metrics(),
                "started_at": _now(),
                "updated_at": _now(),
            }
            openai_register.config.update({k: self._config[k] for k in ("mail", "proxy", "total", "threads", "hero_sms")})
            with openai_register.stats_lock:
                openai_register.stats.update({"done": 0, "success": 0, "fail": 0, "start_time": time.time()})
            self._save()
            self._runner = threading.Thread(target=self._run_codex, daemon=True, name="codex-cpa-register")
            self._runner.start()
            self._append_log(
                f"Codex CPA 注册任务启动，总数={self._config['total']}，线程数={self._config['threads']}，HeroSMS auto_buy={bool(self._config['hero_sms'].get('auto_buy'))}",
                "yellow",
            )
            return self.get()

    def stop(self) -> dict:
        with self._lock:
            self._config["enabled"] = False
            self._config["stats"]["updated_at"] = _now()
            self._save()
            self._append_log("已请求停止注册任务，正在等待当前运行任务结束", "yellow")
            return self.get()

    def reset(self) -> dict:
        with self._lock:
            self._logs = []
            self._config["stats"] = {"success": 0, "fail": 0, "done": 0, "running": 0, "threads": self._config["threads"], "elapsed_seconds": 0, "avg_seconds": 0, "success_rate": 0, **self._pool_metrics(), "updated_at": _now()}
            with openai_register.stats_lock:
                openai_register.stats.update({"done": 0, "success": 0, "fail": 0, "start_time": 0.0})
            self._save()
            return self.get()

    def _append_log(self, text: str, color: str = "") -> None:
        with self._lock:
            self._logs.append({"time": _now(), "text": str(text), "level": str(color or "info")})
            self._logs = self._logs[-300:]

    def _pool_metrics(self) -> dict:
        items = account_service.list_accounts()
        normal = [item for item in items if item.get("status") == "正常"]
        return {
            "current_quota": sum(int(item.get("quota") or 0) for item in normal if not item.get("image_quota_unknown")),
            "current_available": len(normal),
        }

    def _target_reached(self, cfg: dict, submitted: int) -> bool:
        mode = str(cfg.get("mode") or "total")
        metrics = self._pool_metrics()
        self._bump(**metrics)
        if mode == "quota":
            reached = metrics["current_quota"] >= int(cfg.get("target_quota") or 1)
            self._append_log(f"检查号池：当前正常账号={metrics['current_available']}，当前剩余额度={metrics['current_quota']}，目标额度={cfg.get('target_quota')}，{'跳过注册' if reached else '继续注册'}", "yellow")
            return reached
        if mode == "available":
            reached = metrics["current_available"] >= int(cfg.get("target_available") or 1)
            self._append_log(f"检查号池：当前正常账号={metrics['current_available']}，目标账号={cfg.get('target_available')}，当前剩余额度={metrics['current_quota']}，{'跳过注册' if reached else '继续注册'}", "yellow")
            return reached
        return submitted >= int(cfg.get("total") or 1)

    def _bump(self, **updates) -> None:
        with self._lock:
            self._config["stats"].update(updates)
            stats = self._config["stats"]
            started_at = str(stats.get("started_at") or "")
            if started_at:
                try:
                    elapsed = max(0.0, (datetime.now(timezone.utc) - datetime.fromisoformat(started_at)).total_seconds())
                except Exception:
                    elapsed = 0.0
                done = int(stats.get("done") or 0)
                success = int(stats.get("success") or 0)
                fail = int(stats.get("fail") or 0)
                stats["elapsed_seconds"] = round(elapsed, 1)
                stats["avg_seconds"] = round(elapsed / success, 1) if success else 0
                stats["success_rate"] = round(success * 100 / max(1, success + fail), 1)
            self._config["stats"]["updated_at"] = _now()
            self._save()

    def _run(self) -> None:
        threads = int(self.get()["threads"])
        submitted, done, success, fail = 0, 0, 0, 0
        with ThreadPoolExecutor(max_workers=threads) as executor:
            futures = set()
            while True:
                cfg = self.get()
                while self.get()["enabled"] and not self._target_reached(cfg, submitted) and len(futures) < threads:
                    submitted += 1
                    futures.add(executor.submit(openai_register.worker, submitted))
                self._bump(running=len(futures), done=done, success=success, fail=fail)
                if not futures and (not self.get()["enabled"] or str(cfg.get("mode") or "total") == "total"):
                    break
                if not futures:
                    time.sleep(max(1, int(cfg.get("check_interval") or 5)))
                    continue
                finished, futures = wait(futures, return_when=FIRST_COMPLETED)
                for future in finished:
                    done += 1
                    try:
                        result = future.result()
                        success += 1 if result.get("ok") else 0
                        fail += 0 if result.get("ok") else 1
                    except Exception:
                        fail += 1
        self._bump(running=0, done=done, success=success, fail=fail, finished_at=_now())
        with self._lock:
            self._config["enabled"] = False
            self._save()
        self._append_log(f"注册任务结束，成功{success}，失败{fail}", "yellow")

    def _run_codex(self) -> None:
        threads = int(self.get()["threads"])
        total = int(self.get().get("total") or 1)
        submitted, done, success, fail = 0, 0, 0, 0
        with ThreadPoolExecutor(max_workers=threads) as executor:
            futures = set()
            while True:
                while self.get()["enabled"] and submitted < total and len(futures) < threads:
                    submitted += 1
                    futures.add(executor.submit(codex_register_service.run_codex_registration, submitted))
                self._bump(running=len(futures), done=done, success=success, fail=fail)
                if not futures and (not self.get()["enabled"] or submitted >= total):
                    break
                if not futures:
                    time.sleep(1)
                    continue
                finished, futures = wait(futures, return_when=FIRST_COMPLETED)
                for future in finished:
                    done += 1
                    try:
                        result = future.result()
                        success += 1 if result.get("ok") else 0
                        fail += 0 if result.get("ok") else 1
                    except Exception as exc:
                        fail += 1
                        reputation = openai_register._record_mail_failure(exc)
                        if reputation.get("disabled_changed") and isinstance(exc, openai_register.RegisterAttemptError):
                            self._append_log(f"邮箱域名已自动拉黑: {exc.mail_domain}，原因: {reputation.get('bucket')}", "yellow")
                        self._append_log(f"Codex CPA 任务失败: {openai_register.redact_sensitive_text(exc)}", "red")
        self._bump(running=0, done=done, success=success, fail=fail, finished_at=_now())
        with self._lock:
            self._config["enabled"] = False
            self._save()
        self._append_log(f"Codex CPA 注册任务结束，成功{success}，失败{fail}", "yellow")


register_service = RegisterService(REGISTER_FILE)