main.py

import asyncio
import base64
import datetime
import json
import logging
import os
import re
import shutil
import subprocess
import time
from typing import Optional

import requests
from bs4 import BeautifulSoup
from fake_useragent import UserAgent
from fastapi import FastAPI
from huggingface_hub import HfApi, hf_hub_download, login

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class Config:
    """設定用クラス"""
    HUGGINGFACE_API_KEY = os.environ["HUGGINGFACE_API_KEY"]
    CIVITAI_API_TOKEN = os.environ["CIVITAI_API_TOKEN"]
    LOG_FILE = "civitai_backup.log"
    LIST_FILE = "model_list.log"
    REPO_IDS = {
        "log": "ttttdiva/CivitAI_log_test",
        "model_list": "ttttdiva/CivitAI_model_info_test",
        "current": ""
    }
    URLS = {
        "latest": "https://civitai.com/api/v1/models?sort=Newest",
        "modelPage": "https://civitai.com/models/",
        "modelId": "https://civitai.com/api/v1/models/",
    }
    JST = datetime.timezone(datetime.timedelta(hours=9))
    UA = UserAgent()
    HEADERS = {
        "Authorization": f'Bearer {CIVITAI_API_TOKEN}',
        "User-Agent": "civitai-crawler/1.0",
        "Content-Type": "application/json"
    }

    # rclone
    RCLONE_CONF_BASE64 = os.environ.get("RCLONE_CONF_BASE64", "")
    ENCRYPTED_DIR = "/home/user/app/encrypted"


class CivitAICrawler:
    """CivitAIからダウンロード & Hugging Faceにアップロード"""

    def __init__(self, config: Config):
        self.config = config
        self.api = HfApi()
        self.app = FastAPI()
        self.repo_ids = self.config.REPO_IDS.copy()
        self.jst = self.config.JST

        self.setup_rclone_conf()
        self.setup_routes()

    def setup_routes(self):
        @self.app.get("/")
        def root():
            now = datetime.datetime.now(self.jst)
            return f"Status: {now} -- current repo: {self.repo_ids['current']}"

        @self.app.on_event("startup")
        async def on_startup():
            asyncio.create_task(self.crawl())

    def setup_rclone_conf(self):
        if not self.config.RCLONE_CONF_BASE64:
            logger.warning("No RCLONE_CONF_BASE64 found.")
            return
        os.makedirs(".rclone_config", exist_ok=True)
        conf_path = os.path.join(".rclone_config", "rclone.conf")
        with open(conf_path, "wb") as f:
            f.write(base64.b64decode(self.config.RCLONE_CONF_BASE64))
        os.environ["RCLONE_CONFIG"] = conf_path
        logger.info(f"rclone.conf created at: {conf_path}")

    # ============================================================================
    # rcloneで暗号化アップロード
    # ============================================================================
    def encrypt_with_rclone(self, local_path: str):
        if not os.path.exists(local_path):
            raise FileNotFoundError(f"Local path not found: {local_path}")
        if os.path.isdir(self.config.ENCRYPTED_DIR):
            shutil.rmtree(self.config.ENCRYPTED_DIR, ignore_errors=True)

        top_name = os.path.basename(local_path.rstrip("/")) or "unnamed"

        cmd = [
            "rclone", "copy",
            local_path,
            f"cryptLocal:{top_name}",
            "-v"
        ]
        logger.info("Running: %s", " ".join(cmd))
        subprocess.run(cmd, check=True)

        if not os.path.isdir(self.config.ENCRYPTED_DIR):
            raise FileNotFoundError("Encrypted dir not found after rclone copy.")

    def upload_encrypted_files(self, repo_id: str, path_in_repo: str = ""):
        max_retries = 5
        for root, dirs, files in os.walk(self.config.ENCRYPTED_DIR):
            for fn in files:
                enc_file_path = os.path.join(root, fn)
                rel_path = os.path.relpath(enc_file_path, self.config.ENCRYPTED_DIR)
                upload_path = os.path.join(path_in_repo, rel_path)

                attempt = 0
                while attempt < max_retries:
                    try:
                        self.api.upload_file(
                            path_or_fileobj=enc_file_path,
                            repo_id=repo_id,
                            path_in_repo=upload_path
                        )
                        logger.info("[OK] Uploaded => %s/%s", repo_id, upload_path)
                        break
                    except Exception as e:
                        attempt += 1
                        msg = str(e)
                        # rate-limit / file-limit handle
                        if "rate-limited" in msg and "minutes" in msg:
                            import re
                            m = re.search(r"in (\d+) minutes?", msg)
                            if m:
                                mins = int(m.group(1)) + 1
                                logger.warning("Rate-limited. Wait %d minutes.", mins)
                                time.sleep(mins * 60)
                                attempt -= 1
                                continue
                        if "you can retry this action in about 1 hour" in msg:
                            logger.warning("Encountered retry in 1 hour error, waiting..")
                            time.sleep(3600)
                            attempt -= 1
                            continue
                        if "over the limit of 100000 files" in msg:
                            logger.warning("File limit exceeded. Creating new repo..")
                            self.repo_ids['current'] = self.increment_repo_name(self.repo_ids['current'])
                            self.api.create_repo(repo_id=self.repo_ids['current'], private=True)
                            attempt = 0
                            repo_id = self.repo_ids['current']
                            continue
                        if attempt < max_retries:
                            logger.warning("Failed to upload, retry %d/%d..", attempt, max_retries)
                        else:
                            logger.error("Failed after %d attempts: %s", max_retries, enc_file_path)
                            raise

    def upload_folder_encrypted(self, folder_path: str, repo_id: Optional[str] = None, path_in_repo: str = ""):
        if not repo_id:
            repo_id = self.repo_ids['current']
        self.encrypt_with_rclone(folder_path)
        self.upload_encrypted_files(repo_id, path_in_repo=path_in_repo)
        if os.path.isdir(self.config.ENCRYPTED_DIR):
            shutil.rmtree(self.config.ENCRYPTED_DIR, ignore_errors=True)

    def upload_file_encrypted_one_by_one(self, file_path: str, repo_id: Optional[str] = None, path_in_repo: str = ""):
        """古いバージョン用: ファイルを暗号化アップロード → ローカル削除。"""
        if not repo_id:
            repo_id = self.repo_ids['current']

        self.encrypt_with_rclone(file_path)
        self.upload_encrypted_files(repo_id, path_in_repo)
        if os.path.isdir(self.config.ENCRYPTED_DIR):
            shutil.rmtree(self.config.ENCRYPTED_DIR, ignore_errors=True)
        if os.path.exists(file_path):
            os.remove(file_path)

    @staticmethod
    def increment_repo_name(repo_id: str) -> str:
        m = re.search(r'(\d+)$', repo_id)
        if m:
            num = int(m.group(1)) + 1
            return re.sub(r'\d+$', str(num), repo_id)
        else:
            return repo_id + "1"

    # ============================================================================
    # 生ファイルアップロード(ログやmodel_listなど)
    # ============================================================================
    def upload_file_raw(self, file_path: str, repo_id: Optional[str] = None, path_in_repo: Optional[str] = None):
        if not repo_id:
            repo_id = self.repo_ids['current']
        if not path_in_repo:
            path_in_repo = os.path.basename(file_path)

        max_retries = 5
        attempt = 0
        while attempt < max_retries:
            try:
                self.api.upload_file(
                    path_or_fileobj=file_path,
                    repo_id=repo_id,
                    path_in_repo=path_in_repo
                )
                logger.info("[OK] Uploaded raw %s => %s/%s", file_path, repo_id, path_in_repo)
                return
            except Exception as e:
                attempt += 1
                msg = str(e)
                if "over the limit of 100000 files" in msg:
                    logger.warning("File limit exceeded, creating new repo..")
                    self.repo_ids['current'] = self.increment_repo_name(self.repo_ids['current'])
                    self.api.create_repo(repo_id=self.repo_ids['current'], private=True)
                    attempt = 0
                    repo_id = self.repo_ids['current']
                    continue
                if "you can retry this action in about 1 hour" in msg:
                    logger.warning("Rate-limited 1h, waiting..")
                    time.sleep(3600)
                    attempt -= 1
                else:
                    if attempt < max_retries:
                        logger.warning("Failed raw upload attempt %d/%d..", attempt, max_retries)
                    else:
                        logger.error("Failed raw upload after %d attempts: %s", max_retries, file_path)
                        raise

    # ============================================================================
    # ダウンロード
    # ============================================================================
    @staticmethod
    def get_filename_from_cd(cd: Optional[str], default_name: str) -> str:
        if cd:
            parts = cd.split(";")
            for p in parts:
                if "filename=" in p:
                    return p.split("=")[1].strip().strip('"')
        return default_name

    def download_file(self, url: str, dest_folder: str, default_name: str):
        try:
            resp = requests.get(url, headers=self.config.HEADERS, stream=True)
            resp.raise_for_status()
        except requests.RequestException as e:
            logger.error("Failed to DL from %s: %s", url, e)
            return None

        filename = self.get_filename_from_cd(resp.headers.get('content-disposition'), default_name)
        local_path = os.path.join(dest_folder, filename)
        os.makedirs(dest_folder, exist_ok=True)

        with open(local_path, "wb") as f:
            for chunk in resp.iter_content(chunk_size=8192):
                f.write(chunk)

        logger.info("Downloaded: %s", local_path)
        return local_path

    # ============================================================================
    # 旧バージョンファイルだけ1ファイルずつアップロード→削除
    # ============================================================================
    def download_old_versions_one_by_one(self, model_versions: list, folder: str):
        """model_versions[1:] を対象にファイルDL→暗号化UL→削除."""
        if len(model_versions) <= 1:
            return
        old_dir = os.path.join(folder, "old_versions")
        os.makedirs(old_dir, exist_ok=True)

        for ver in model_versions[1:]:
            for f_info in ver.get("files", []):
                url = f_info["downloadUrl"]
                fname = f_info["name"]
                local_path = self.download_file(url, old_dir, fname)
                if local_path and os.path.exists(local_path):
                    self.upload_file_encrypted_one_by_one(local_path)
        if os.path.exists(old_dir):
            shutil.rmtree(old_dir, ignore_errors=True)

    # ============================================================================
    # 通常の(最新)モデルファイル & images フォルダを一括DL→アップロード
    # ============================================================================
    def download_latest_files(self, model_versions: list, folder: str):
        latest_files = model_versions[0].get("files", [])
        os.makedirs(folder, exist_ok=True)
        # 最新ファイル
        for f_info in latest_files:
            url = f_info["downloadUrl"]
            fname = f_info["name"]
            self.download_file(url, folder, fname)

    def download_images(self, model_versions: list, folder: str):
        images_folder = os.path.join(folder, "images")
        os.makedirs(images_folder, exist_ok=True)

        for ver in model_versions:
            for img in ver.get("images", []):
                img_url = img["url"]
                # 例: 49945407.jpeg --> "49945407.jpeg.png"
                image_name = os.path.basename(img_url) + ".png"
                self.download_file(img_url, images_folder, image_name)

    def save_html_content(self, url: str, folder: str):
        try:
            resp = requests.get(url)
            resp.raise_for_status()
            html_name = os.path.basename(folder) + ".html"
            html_path = os.path.join(folder, html_name)
            with open(html_path, "w", encoding="utf-8") as f:
                f.write(resp.text)
            logger.info("Saved HTML => %s", html_path)
        except Exception as e:
            logger.error("Failed to save HTML from %s: %s", url, e)

    def save_model_info(self, model_info: dict, folder: str):
        info_path = os.path.join(folder, "model_info.json")
        try:
            with open(info_path, "w", encoding="utf-8") as f:
                json.dump(model_info, f, indent=2)
            logger.info("Saved model_info => %s", info_path)
        except Exception as e:
            logger.error("Failed to save model_info: %s", e)

    # ============================================================================
    # model_list.log の読み書き
    # ============================================================================
    def read_model_list(self):
        ret = {}
        if not os.path.exists(self.config.LIST_FILE):
            return ret
        try:
            with open(self.config.LIST_FILE, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    # 例: "111123: https://huggingface.co/xxx"
                    parts = line.split(": ", 1)
                    if len(parts) == 2:
                        key, val = parts
                        ret[key] = val
        except Exception as e:
            logger.error("Failed to read model_list.log: %s", e)
        return ret

    def append_model_list(self, key: str, val: str):
        """'key: val' を model_list.log の末尾に追加."""
        with open(self.config.LIST_FILE, "a", encoding="utf-8") as f:
            f.write(f"{key}: {val}\n")
        logger.info("Appended to model_list.log => '%s: %s'", key, val)

    # ============================================================================
    # get_model_info
    # ============================================================================
    def get_model_info(self, model_id: str):
        url = self.config.URLS["modelId"] + str(model_id)
        try:
            resp = requests.get(url, headers=self.config.HEADERS)
            resp.raise_for_status()
            return resp.json()
        except Exception as e:
            logger.error("Failed to get model info for %s: %s", model_id, e)
            return {}

    # ============================================================================
    # メイン: 1モデルを処理
    # ============================================================================
    def process_model(self, model_url: str):
        try:
            model_id = model_url.rstrip("/").split("/")[-1]
            model_info = self.get_model_info(model_id)
            if not model_info:
                logger.error("No model_info for %s", model_id)
                return
            vers = model_info.get("modelVersions", [])
            if not vers:
                logger.error("No modelVersions for %s", model_id)
                return

            # 既にアップしたかどうか model_list.log で判定(例: model_id で判定)
            model_list = self.read_model_list()
            if model_id in model_list:
                logger.info("Model ID %s is already in model_list. Skipping..", model_id)
                return

            # まずフォルダ名
            latest_files = vers[0].get("files", [])
            if latest_files:
                main_file = next((f for f in latest_files if f.get("type") == "Model"), latest_files[0])
                main_name = main_file["name"]
                folder = os.path.splitext(main_name)[0]
            else:
                folder = f"model_{model_id}"

            # 1) 最新バージョン & images をまとめて download
            os.makedirs(folder, exist_ok=True)
            self.download_latest_files(vers, folder)
            self.download_images(vers, folder)

            # HTML & model_info.json
            self.save_html_content(self.config.URLS["modelPage"] + str(model_id), folder)
            self.save_model_info(model_info, folder)

            # 2) old_versions を1ファイルずつアップロード
            self.download_old_versions_one_by_one(vers, folder)

            # 3) 上記が終わると old_versions は消えている。
            #    あとはフォルダに最新バージョンのファイル＆images etc.が残っているので
            #    フォルダ一括で暗号化アップロード
            encrypted_top_name = self.upload_folder_encrypted(folder)
            shutil.rmtree(folder, ignore_errors=True)

            # model_list.log に書き込み (例: "111123: https://huggingface.co/xxx/tree/main")
            #   → 実際にフォルダ暗号化でトップレベル名が変わってるはずなので、
            #      ざっくり "/tree/main" を書く
            hf_url = f"https://huggingface.co/{self.repo_ids['current']}/tree/main/{encrypted_top_name}"
            self.append_model_list(model_id, hf_url)

        except Exception as e:
            logger.error("Unexpected error in process_model(%s): %s", model_url, e)

    # ============================================================================
    # メインループ
    # ============================================================================
    async def crawl(self):
        while True:
            try:
                login(token=self.config.HUGGINGFACE_API_KEY, add_to_git_credential=True)
                # ダウンロードしてローカルにコピー
                ml_path = hf_hub_download(repo_id=self.repo_ids["model_list"], filename=self.config.LIST_FILE)
                shutil.copyfile(ml_path, self.config.LIST_FILE)

                log_path = hf_hub_download(repo_id=self.repo_ids["log"], filename=self.config.LOG_FILE)
                shutil.copyfile(log_path, self.config.LOG_FILE)

                with open(self.config.LOG_FILE, "r", encoding="utf-8") as f:
                    lines = f.read().splitlines()
                    old_models = json.loads(lines[0]) if len(lines) > 0 else []
                    self.repo_ids["current"] = lines[1] if len(lines) > 1 else ""

                # 最新モデルチェック
                r = requests.get(self.config.URLS["latest"], headers=self.config.HEADERS)
                r.raise_for_status()
                items = r.json().get("items", [])
                latest_ids = [it["id"] for it in items if "id" in it]

                new_ids = list(set(latest_ids) - set(old_models))
                if new_ids:
                    logger.info("New models found: %s", new_ids)
                    target_id = new_ids[0]
                    for attempt in range(1, 6):
                        try:
                            self.process_model(f"{self.config.URLS['modelId']}{target_id}")
                            break
                        except Exception as e:
                            logger.error("Failed to process ID %s (attempt %d/5): %s", target_id, attempt, e)
                            if attempt == 5:
                                logger.error("Skipping ID %s after 5 attempts", target_id)
                            else:
                                await asyncio.sleep(2)
                else:
                    # 新モデルなし
                    with open(self.config.LOG_FILE, "w", encoding="utf-8") as f:
                        f.write(json.dumps(latest_ids) + "\n")
                        f.write(self.repo_ids["current"] + "\n")
                    self.upload_file_raw(self.config.LOG_FILE, self.repo_ids["log"], self.config.LOG_FILE)
                    logger.info("No new models found. Sleep 60s..")
                    await asyncio.sleep(60)
                    continue

                # 成功モデルIDを old_models に追加 & civitai_backup.log 書き込み
                old_models.append(target_id)
                with open(self.config.LOG_FILE, "w", encoding="utf-8") as f:
                    f.write(json.dumps(old_models) + "\n")
                    f.write(self.repo_ids["current"] + "\n")
                logger.info("Updated log with new ID: %s", target_id)

                # logs をアップ
                self.upload_file_raw(self.config.LOG_FILE, self.repo_ids["log"], self.config.LOG_FILE)
                self.upload_file_raw(self.config.LIST_FILE, self.repo_ids["model_list"], self.config.LIST_FILE)

            except Exception as e:
                logger.error("Error in crawl: %s", e)
                await asyncio.sleep(300)


config = Config()
crawler = CivitAICrawler(config)
app = crawler.app