src/services.py

import os
import json
import io
import hashlib
import threading
import time
import logging
import requests
from requests.auth import HTTPBasicAuth
from huggingface_hub import HfApi, hf_hub_download
from utils import add_script_run_ctx

TARGET_SPACE = os.getenv("TARGET_SPACE")
HF_TOKEN = os.getenv("HF_TOKEN")
DATASET_ID = os.getenv("DATASET_ID")

class TokenManager:
    def __init__(self):
        self.api = HfApi(token=HF_TOKEN)
        self.config_file = "pollinations_tokens.json"
        
        self.tokens = self.load_tokens()
        self._migrate_data_structure()
        
        self.last_save_time = 0
        self.lock = threading.Lock()
        
        self.monitor_thread = threading.Thread(target=self._monitor_loop, daemon=True)
        add_script_run_ctx(self.monitor_thread)
        self.monitor_thread.start()

    def _migrate_data_structure(self):
        dirty = False
        cleaned_tokens = []
        for t in self.tokens:
            if "key" in t and "token" not in t:
                t["token"] = t.pop("key")
                dirty = True
            if "alias" not in t: t["alias"] = "Unknown"
            if "status" not in t: t["status"] = "active"
            if "balance" not in t: t["balance"] = 0.0
            if t.get("token"): cleaned_tokens.append(t)
            else: dirty = True
        self.tokens = cleaned_tokens
        if dirty: self.save_tokens()

    def load_tokens(self):
        try:
            return json.load(open(self.config_file, "r"))
        except:
            try:
                path = hf_hub_download(repo_id=DATASET_ID, filename=self.config_file, repo_type="dataset", token=HF_TOKEN)
                return json.load(open(path, "r"))
            except: return []

    def save_tokens(self):
        # 限制保存频率
        if time.time() - self.last_save_time < 5: return False
        try:
            with open(self.config_file, "w") as f: json.dump(self.tokens, f, indent=2)
            self.api.upload_file(
                path_or_fileobj=io.BytesIO(json.dumps(self.tokens, indent=2).encode("utf-8")), 
                path_in_repo=self.config_file, 
                repo_id=DATASET_ID, 
                repo_type="dataset"
            )
            self.last_save_time = time.time()
            return True
        except: return False

    def add_token(self, alias, token):
        with self.lock:
            existing = next((t for t in self.tokens if t["token"] == token), None)
            if existing:
                existing["alias"] = alias
                existing["status"] = "active"
                logging.info(f"🔄 Token updated: {alias}")
            else:
                self.tokens.append({"alias": alias, "token": token, "status": "active", "balance": 0.0})
                logging.info(f"➕ New Token added: {alias}")
        self.save_tokens()

    def delete_token(self, idx):
        with self.lock:
            if 0 <= idx < len(self.tokens): del self.tokens[idx]
        self.save_tokens()

    def _check_balance_http(self, token):
        """
        查询余额。
        返回: 
          float: 余额数值
          0.0: 明确的余额不足或Token无效(403)
          None: 网络错误，保留原状态
        """
        try:
            url = f"https://gen.pollinations.ai/account/balance?key={token}"
            headers = {
                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
            }
            resp = requests.get(url, headers=headers, timeout=10)
            
            if resp.status_code == 200:
                try: return float(resp.text.strip())
                except: return float(resp.json().get('balance', 0))
            
            # 如果是 403 Forbidden 或 401 Unauthorized，说明 Token 废了，视为余额 0 处理
            if resp.status_code in [401, 403]:
                return 0.0
                
            return None # 其他错误 (500等) 返回 None，不更新状态
        except: return None

    def _monitor_loop(self):
        logging.info("🚀 TokenMonitor thread started.")
        while True:
            changes_detected = False
            # 浅拷贝列表以避免迭代时修改问题
            with self.lock: tokens_to_check = list(enumerate(self.tokens))
            
            for idx, t in tokens_to_check:
                token_str = t.get("token")
                if not token_str: continue 

                bal = self._check_balance_http(token_str)
                
                with self.lock:
                    # 再次确认索引有效性
                    if idx < len(self.tokens) and self.tokens[idx].get("token") == token_str:
                        if bal is not None:
                            # 更新余额显示
                            if abs(self.tokens[idx]["balance"] - bal) > 0.001:
                                self.tokens[idx]["balance"] = bal
                                changes_detected = True
                            
                            # [V30.38] 状态流转核心逻辑
                            # 规则：> 0.1 Active, <= 0.1 Exhausted
                            current_status = self.tokens[idx]["status"]
                            
                            if bal > 0.1:
                                if current_status != "active":
                                    logging.info(f"🔋 Token {self.tokens[idx]['alias']} 余额充足 ({bal})，状态激活 (Active)")
                                    self.tokens[idx]["status"] = "active"
                                    changes_detected = True
                            else:
                                if current_status != "exhausted":
                                    logging.warning(f"🪫 Token {self.tokens[idx]['alias']} 余额不足/无效 ({bal})，标记耗尽 (Exhausted)")
                                    self.tokens[idx]["status"] = "exhausted"
                                    changes_detected = True
            
            if changes_detected: self.save_tokens()
            time.sleep(60) # 60秒轮询一次

    def get_active_tokens(self):
        """
        [V30.38] 获取当前所有可用 Token 的列表 (快照)
        执行层将拿着这个列表去轮询，而不依赖 Manager 每次分配
        """
        with self.lock:
            # 筛选出所有状态为 active 的 token
            # 按余额降序排列，尽可能先用钱多的
            candidates = [t for t in self.tokens if t.get("status") == "active" and t.get("token")]
            candidates.sort(key=lambda x: x.get("balance", 0.0), reverse=True)
            return [t["token"] for t in candidates]

class AuthManager:
    def __init__(self): 
        self.api = HfApi(token=HF_TOKEN)
        self.user_file = "users_config.json"
        
    def load_users(self):
        try:
            if os.path.exists(self.user_file):
                mtime = os.path.getmtime(self.user_file)
                if time.time() - mtime < 600: 
                    return json.load(open(self.user_file, "r"))
            path = hf_hub_download(repo_id=DATASET_ID, filename=self.user_file, repo_type="dataset", token=HF_TOKEN, local_dir=".", force_download=True)
            return json.load(open(path, "r"))
        except: 
            return {"admin": {"password": hashlib.sha256("admin".encode()).hexdigest(), "must_change_password": True}}
            
    def save_users(self, u):
        try: 
            self.api.upload_file(path_or_fileobj=io.BytesIO(json.dumps(u, indent=2).encode("utf-8")), path_in_repo=self.user_file, repo_id=DATASET_ID, repo_type="dataset")
            with open(self.user_file, "w") as f: json.dump(u, f, indent=2)
        except: pass
        
    def login(self, u, p): 
        d = self.load_users()
        if u in d and hashlib.sha256(p.encode()).hexdigest() == d[u]["password"]:
            return (True, d[u].get("must_change_password", False))
        else: return (False, False)
            
    def update_password(self, u, np): 
        d = self.load_users()
        d[u]["password"] = hashlib.sha256(np.encode()).hexdigest()
        d[u]["must_change_password"] = False
        self.save_users(d)
        return True

class WebDAVManager:
    def __init__(self):
        self.api = HfApi(token=HF_TOKEN)
        self.config_file = "webdav_config.json"
        self.configs = self.load_configs()

    def load_configs(self):
        try: 
            if os.path.exists(self.config_file): return json.load(open(self.config_file, "r"))
            path = hf_hub_download(repo_id=DATASET_ID, filename=self.config_file, repo_type="dataset", token=HF_TOKEN)
            return json.load(open(path, "r"))
        except: return [] 

    def save_configs(self):
        try: 
            self.api.upload_file(path_or_fileobj=io.BytesIO(json.dumps(self.configs, indent=2).encode("utf-8")), path_in_repo=self.config_file, repo_id=DATASET_ID, repo_type="dataset")
            with open(self.config_file, "w") as f: json.dump(self.configs, f, indent=2)
        except: pass
        
    def test_connection(self, h, u, p, r):
        h = h.rstrip("/")
        h = "https://" + h if not h.startswith("http") else h
        try:
            res = requests.request("PROPFIND", h + r, auth=HTTPBasicAuth(u, p), timeout=15)
            return (True, "✅") if res.status_code in [200, 207] else (False, f"HTTP {res.status_code}")
        except Exception as e: return False, str(e)

    def check_exists(self, url, auth):
        try:
            return requests.head(url, auth=auth, timeout=10).status_code == 200
        except: return False

    def upload_file(self, conn_id, file_bytes, filename):
        target_conf = next((c for c in self.configs if c["id"] == conn_id), None)
        if not target_conf: return False, "Config Not Found"
        opts = target_conf["options"]
        base_url = ("https://" if not opts["webdav_hostname"].startswith("http") else "") + opts["webdav_hostname"].rstrip("/")
        try:
            from utils import ensure_png_bytes
            url = f"{base_url}{opts.get('webdav_root', '/').rstrip('/')}/{filename.lstrip('/')}"
            auth = HTTPBasicAuth(opts["webdav_login"], opts["webdav_password"])
            res = requests.put(url, data=ensure_png_bytes(file_bytes), headers={'Content-Type': 'image/png'}, auth=auth, timeout=180)
            return (True, "Success") if res.status_code in [200, 201, 204] else (False, f"HTTP {res.status_code}")
        except Exception as e: return False, str(e)
    
    def upload_to_all_nodes(self, file_bytes, filename):
        results = []
        for c in self.configs:
            if not c.get("enabled", True): continue
            ok, msg = self.upload_file(c["id"], file_bytes, filename)
            results.append(f"{'✅' if ok else '❌'} {c['name']}")
        return " | ".join(results)

    def add_connection(self, n, h, u, p, r="/", enabled=True):
        self.configs.append({"id": int(time.time()), "name": n, "enabled": enabled, "options": {"webdav_hostname": h, "webdav_login": u, "webdav_password": p, "webdav_root": r}})
        self.save_configs()

    def delete_connection(self, conn_id):
        self.configs = [c for c in self.configs if c["id"] != conn_id]
        self.save_configs()
    
    def update_connection(self, conn_id, n, h, u, p, r, enabled):
        for c in self.configs:
            if c["id"] == conn_id:
                c["name"] = n
                c["enabled"] = enabled
                c["options"] = {"webdav_hostname": h, "webdav_login": u, "webdav_password": p, "webdav_root": r}
                break
        self.save_configs()