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, create_repo, 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/",
        "modelVersionId": "https://civitai.com/api/v1/model-versions/",
        "hash": "https://civitai.com/api/v1/model-versions/by-hash/"
    }
    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", "")
    # 暗号化されたファイルが出力されるローカルディレクトリ（cryptLocal: の実体）
    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

        # rclone のセットアップ
        self.setup_rclone_conf()

        self.setup_routes()

    def setup_routes(self):
        """FastAPIのルーティングを設定する。"""
        @self.app.get("/")
        def read_root():
            now = str(datetime.datetime.now(self.jst))
            description = f"""
            CivitAIを定期的に周回し新規モデルを {self.repo_ids['current']} にバックアップするSpaceです。
            model_list.log や civitai_backup.log は暗号化しないでアップロードします。
            モデルのフォルダやファイルは暗号化してアップロードします。
            Status: {now} + currently running :D
            """
            return description

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

    # =============================================================================
    # rclone の設定・暗号化アップロード処理
    # =============================================================================
    def setup_rclone_conf(self):
        """環境変数 RCLONE_CONF_BASE64 から rclone.conf を生成し、RCLONE_CONFIG 環境変数を設定"""
        if not self.config.RCLONE_CONF_BASE64:
            logger.warning("[WARN] RCLONE_CONF_BASE64 is empty. rclone may fail.")
            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"[INFO] rclone.conf created at: {conf_path}")

    def encrypt_with_rclone(self, local_path: str):
        """
        指定ファイル or ディレクトリを cryptLocal: にコピー。
        フォルダ構造やファイル名を rclone の filename_encryption 設定に応じて暗号化する。
        """
        if not os.path.exists(local_path):
            raise FileNotFoundError(f"[ERROR] Local path not found: {local_path}")

        # 事前に暗号先ディレクトリをクリーンアップ
        if os.path.isdir(self.config.ENCRYPTED_DIR):
            shutil.rmtree(self.config.ENCRYPTED_DIR, ignore_errors=True)

        top_level_name = os.path.basename(local_path.rstrip("/"))
        if not top_level_name:
            top_level_name = "unnamed"

        cmd = ["rclone", "copy", local_path, f"cryptLocal:{top_level_name}", "-v"]
        logger.info(f"[INFO] Running: {' '.join(cmd)}")
        subprocess.run(cmd, check=True)
        logger.info(f"[OK] rclone copy => cryptLocal:{top_level_name}")

        if not os.path.isdir(self.config.ENCRYPTED_DIR):
            raise FileNotFoundError(
                f"[ERROR] {self.config.ENCRYPTED_DIR} not found. Check your rclone config."
            )

    # 例: upload_encrypted_files の中の再試行処理
    def upload_encrypted_files(self, repo_id: str, base_path_in_repo: str = ""):
        max_retries = 5
        for root, dirs, files in os.walk(self.config.ENCRYPTED_DIR):
            for fn in files:
                encrypted_file_path = os.path.join(root, fn)
                if not os.path.isfile(encrypted_file_path):
                    continue

                relative_path = os.path.relpath(encrypted_file_path, self.config.ENCRYPTED_DIR)
                upload_path_in_repo = os.path.join(base_path_in_repo, relative_path)

                attempt = 0
                while attempt < max_retries:
                    try:
                        self.api.upload_file(
                            path_or_fileobj=encrypted_file_path,
                            repo_id=repo_id,
                            path_in_repo=upload_path_in_repo
                        )
                        logger.info(f"[OK] Uploaded => {repo_id}/{upload_path_in_repo}")
                        break

                    except Exception as e:
                        attempt += 1
                        error_message = str(e)

                        # ================================
                        # 429によるrate-limit検出追加
                        # ================================
                        # "You have been rate-limited; you can retry this action in 31 minutes."
                        # のようなメッセージから時間を抽出し、その時間+1分だけ待機後、再試行
                        if "rate-limited" in error_message and "minutes" in error_message:
                            import re
                            match = re.search(r"in (\d+) minutes?", error_message)
                            if match:
                                minutes = int(match.group(1))
                                # +1分して待機
                                minutes += 1
                                logger.warning(f"Rate-limited. Waiting {minutes} minutes before retry...")
                                time.sleep(minutes * 60)
                                attempt -= 1  # 同じ attempt カウントで再試行
                                continue

                        # ================================
                        # すでにある1時間待機処理
                        # ================================
                        if "you can retry this action in about 1 hour" in error_message:
                            logger.warning("Encountered 'retry in 1 hour' error. Waiting 1 hour before retrying...")
                            time.sleep(3600)
                            attempt -= 1  # 再試行回数を増やさずにループを続ける
                            continue

                        if "over the limit of 100000 files" in error_message:
                            logger.warning("Repository file limit exceeded. Creating a new repository...")
                            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(
                                f"Failed to upload {encrypted_file_path}, retry {attempt}/{max_retries}..."
                            )
                        else:
                            logger.error(
                                f"Failed to upload after {max_retries} attempts: {encrypted_file_path}"
                            )
                            raise

    @staticmethod
    def get_filename_from_cd(content_disposition: Optional[str], default_name: str) -> str:
        if content_disposition:
            parts = content_disposition.split(';')
            for part in parts:
                if "filename=" in part:
                    return part.split("=")[1].strip().strip('"')
        return default_name

    def download_file(self, url: str, destination_folder: str, default_name: str):
        try:
            response = requests.get(url, headers=self.config.HEADERS, stream=True)
            response.raise_for_status()
        except requests.RequestException as e:
            logger.error(f"Failed to download file from {url}: {e}")
            return

        filename = self.get_filename_from_cd(response.headers.get('content-disposition'), default_name)
        file_path = os.path.join(destination_folder, filename)

        with open(file_path, 'wb') as file:
            for chunk in response.iter_content(chunk_size=8192):
                file.write(chunk)
        logger.info(f"Download completed: {file_path}")

    def get_model_info(self, model_id: str) -> dict:
        try:
            response = requests.get(self.config.URLS["modelId"] + str(model_id), headers=self.config.HEADERS)
            response.raise_for_status()
            return response.json()
        except requests.RequestException as e:
            logger.error(f"Failed to retrieve model info for ID {model_id}: {e}")

    def download_model(self, model_versions: list, folder: str):
        """
        - 最新バージョンのファイルを「folder」へまとめてダウンロード
        - 古いバージョンはまだダウンロードしない (実は後から個別にダウンロード)
        （※今回は「まとめて old_versions フォルダに入れずに、1ファイルずつ別のメソッドで対処する」流れに）
        """
        latest_version = model_versions[0]
        latest_files = latest_version["files"]

        for file_info in latest_files:
            download_url = file_info["downloadUrl"]
            file_name = file_info["name"]
            login_detected_count = 0

            while login_detected_count < 5:
                try:
                    self.download_file(download_url, folder, file_name)
                except Exception as e:
                    logger.error(f"Exception occurred while downloading {file_name}: {e}")
                    login_detected_count += 1
                    continue

                if "login" in os.listdir(folder):
                    login_detected_count += 1
                    logger.warning(f"'login' file found. Will try again. ({login_detected_count}/5)")
                    os.remove(os.path.join(folder, "login"))
                else:
                    logger.info(f"Successfully downloaded {file_name}")
                    break

            if login_detected_count >= 5:
                dummy_file_name = f"{file_name}.download_failed"
                dummy_file_path = os.path.join(folder, dummy_file_name)
                try:
                    with open(dummy_file_path, "w") as f:
                        f.write("Download failed after 5 attempts.")
                    logger.error(f"Failed to download {file_name}. Created dummy file {dummy_file_name}. URL: {download_url}")
                except Exception as e:
                    logger.error(f"Failed to create dummy file for {file_name}: {e}")

    def download_old_version_file_and_upload(self, file_info, encrypted_top_name: str):
        """
        古いバージョンのファイルを1個ダウンロードしたら即アップロードし、ローカル削除。
        - file_info: CivitAI API で得られる 'files' の要素(dict)
        - encrypted_top_name: すでにアップ済みのフォルダのトップ (ランダム名)
        """
        file_name = file_info["name"]
        download_url = file_info["downloadUrl"]

        # 一時フォルダ作成
        temp_folder = "temp_old_file"
        os.makedirs(temp_folder, exist_ok=True)
        local_path = os.path.join(temp_folder, file_name)

        login_detected_count = 0
        while login_detected_count < 5:
            try:
                self.download_file(download_url, temp_folder, file_name)
            except Exception as e:
                logger.error(f"Exception while downloading old file {file_name}: {e}")
                login_detected_count += 1
                continue

            if "login" in os.listdir(temp_folder):
                login_detected_count += 1
                logger.warning(f"'login' file found while downloading {file_name}. Retry {login_detected_count}/5.")
                os.remove(os.path.join(temp_folder, "login"))
            else:
                logger.info(f"Successfully downloaded old version file: {file_name}")
                break

        if login_detected_count >= 5:
            dummy_file_name = f"{file_name}.download_failed"
            dummy_file_path = os.path.join(temp_folder, dummy_file_name)
            try:
                with open(dummy_file_path, "w") as f:
                    f.write("Download failed after 5 attempts.")
                logger.error(f"Failed to download {file_name} -> created dummy: {dummy_file_name}")
            except Exception as e:
                logger.error(f"Failed to create dummy file for old version {file_name}: {e}")
            # ダウンロード失敗したのでこのファイルはアップロードしない
            shutil.rmtree(temp_folder, ignore_errors=True)
            return

        # ダウンロード成功したので、暗号化してアップロード(既存のトップフォルダに追加)
        try:
            # encrypt_with_rclone(temp_folder) で temp_folder 全体を暗号化
            # → upload_encrypted_files(... base_path_in_repo=encrypted_top_name) で
            #    既存フォルダ(encrypted_top_name) 以下にファイルを追加
            self.encrypt_with_rclone(temp_folder)
            self.upload_encrypted_files(
                repo_id=self.repo_ids["current"],
                base_path_in_repo=encrypted_top_name  # ★ここがポイント★
            )
            logger.info(f"Uploaded old version file: {file_name} into {encrypted_top_name}")

        except Exception as e:
            logger.error(f"Error uploading old version file {file_name}: {e}")

        # 一時フォルダを削除
        shutil.rmtree(temp_folder, ignore_errors=True)

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

        images = []
        for version in model_versions:
            for img in version.get("images", []):
                image_url = img["url"]
                images.append(image_url)

        for image_url in images:
            image_name = image_url.split("/")[-1]
            try:
                response = requests.get(image_url)
                response.raise_for_status()
                with open(os.path.join(images_folder, f"{image_name}.png"), "wb") as file:
                    file.write(response.content)
            except requests.RequestException as e:
                logger.error(f"Error downloading image {image_url}: {e}")

    def save_html_content(self, url: str, folder: str):
        try:
            response = requests.get(url)
            response.raise_for_status()
            html_path = os.path.join(folder, f"{folder}.html")
            with open(html_path, 'w', encoding='utf-8') as file:
                file.write(response.text)
        except Exception as e:
            logger.error(f"Error saving HTML content for URL {url}: {e}")

    @staticmethod
    def save_model_info(model_info: dict, folder: str):
        with open(os.path.join(folder, "model_info.json"), "w") as file:
            json.dump(model_info, file, indent=2)

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

    # =============================================================================
    # 暗号化しないアップロード（ログや model_list.log 用）
    # =============================================================================
    def upload_file_raw(
        self,
        file_path: str,
        repo_id: Optional[str] = None,
        path_in_repo: Optional[str] = None
    ):
        if repo_id is None:
            repo_id = self.repo_ids['current']
        if path_in_repo is None:
            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(f"[OK] Uploaded {file_path} => {repo_id}/{path_in_repo}")
                return
            except Exception as e:
                attempt += 1
                error_message = str(e)
                if "over the limit of 100000 files" in error_message:
                    logger.warning("Repository file limit exceeded, creating a new repository.")
                    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
                elif "you can retry this action in about 1 hour" in error_message:
                    logger.warning("Encountered 'retry in 1 hour' error. Waiting 1 hour before retrying...")
                    time.sleep(3600)
                    attempt -= 1
                else:
                    if attempt < max_retries:
                        logger.warning(f"Failed to upload raw file {file_path}, retry {attempt}/{max_retries}...")
                    else:
                        logger.error(f"Failed to upload raw file after {max_retries} attempts: {file_path}")
                        raise

    # =============================================================================
    # 暗号化してアップロード (単ファイル)
    # =============================================================================
    def upload_file_encrypted(
        self,
        file_path: str,
        repo_id: Optional[str] = None,
        path_in_repo: Optional[str] = None
    ):
        if repo_id is None:
            repo_id = self.repo_ids['current']
        base_path = path_in_repo or ""

        self.encrypt_with_rclone(file_path)
        self.upload_encrypted_files(repo_id=repo_id, base_path_in_repo=base_path)

        if os.path.isdir(self.config.ENCRYPTED_DIR):
            shutil.rmtree(self.config.ENCRYPTED_DIR, ignore_errors=True)

    # =============================================================================
    # 暗号化してアップロード (フォルダ)
    # =============================================================================
    def upload_folder_encrypted(
        self,
        folder_path: str,
        repo_id: Optional[str] = None,
        path_in_repo: Optional[str] = None
    ) -> str:
        if repo_id is None:
            repo_id = self.repo_ids['current']
        base_path = path_in_repo or ""

        self.encrypt_with_rclone(folder_path)

        top_levels = [
            d for d in os.listdir(self.config.ENCRYPTED_DIR)
            if os.path.isdir(os.path.join(self.config.ENCRYPTED_DIR, d))
        ]
        if not top_levels:
            raise RuntimeError("No top-level folder found after rclone encryption.")
        if len(top_levels) > 1:
            logger.warning(f"Multiple top-level folders found after encryption? {top_levels}. Using the first one.")

        encrypted_top_name = top_levels[0]

        self.upload_encrypted_files(repo_id=repo_id, base_path_in_repo=base_path)

        if os.path.isdir(self.config.ENCRYPTED_DIR):
            shutil.rmtree(self.config.ENCRYPTED_DIR, ignore_errors=True)

        return encrypted_top_name

    # =============================================================================
    # model_list.log の読み書きを「model_id: model_hf_url」で扱うよう変更
    # =============================================================================
    def read_model_list(self):
        """
        model_list.log の各行を
          "123456: https://huggingface.co/...encrypted_folder_name"
        の形式で読み込み、 { "123456": "https://huggingface.co/..."} の dict を返す
        """
        model_list = {}
        try:
            with open(self.config.LIST_FILE, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    parts = line.split(": ", 1)
                    if len(parts) == 2:
                        stored_id, stored_url = parts
                        model_list[stored_id] = stored_url
            return model_list
        except Exception as e:
            logger.error(f"Failed to read model list: {e}")
            return {}

    def process_model(self, model_url: str):
        """
        - 最新バージョンをフォルダに一括DL → フォルダごと暗号化アップロード
        - 古いバージョンはファイル単位で即アップロード
        """
        try:
            model_id = model_url.rstrip("/").split("/")[-1]
            model_info = self.get_model_info(model_id)

            # =====================================
            # 1) フォルダ作成 & 最新バージョンDL
            # =====================================
            model_versions = model_info.get("modelVersions", [])
            if not model_versions:
                logger.warning(f"No versions found for model ID {model_id}")
                return

            # 例: 最新バージョン用フォルダ名を適当に
            folder = f"model_{model_id}_latest"
            os.makedirs(folder, exist_ok=True)

            # download_model() で最新バージョンだけダウンロード
            self.download_model(model_versions, folder)

            # 画像ダウンロード & HTML保存 & info.json保存（従来通り）
            self.download_images(model_versions, folder)
            self.save_html_content(model_url, folder)
            self.save_model_info(model_info, folder)

            # =====================================
            # 2) 最新バージョンのフォルダを丸ごと暗号化 & アップロード
            #    → Hugging Face 側には"ランダムフォルダ名"が作られる
            # =====================================
            encrypted_top_name = self.upload_folder_encrypted(folder)
            logger.info(f"[MAIN] Uploaded latest version folder => {encrypted_top_name}")

            # ローカル削除
            shutil.rmtree(folder, ignore_errors=True)

            # =====================================
            # 3) 古いバージョンのファイルを1つずつDL → 同じフォルダに追加
            # =====================================
            if len(model_versions) > 1:
                for version in model_versions[1:]:
                    for file_info in version["files"]:
                        self.download_old_version_file_and_upload(file_info, encrypted_top_name)

            # =====================================
            # 4) model_list.log への登録やその他の処理
            # =====================================
            modelpage_name = model_info.get("name", f"Model_{model_id}")
            model_hf_url = f"https://huggingface.co/{self.repo_ids['current']}/tree/main/{encrypted_top_name}"
            with open(self.config.LIST_FILE, "a", encoding="utf-8") as f:
                f.write(f"{modelpage_name}: {model_hf_url}\n")

        except Exception as e:
            logger.error(f"Unexpected error in process_model ({model_url}): {e}")

    async def crawl(self):
        """モデルを定期的にチェックし、更新を行う。"""
        while True:
            try:
                login(token=self.config.HUGGINGFACE_API_KEY, add_to_git_credential=True)

                # model_list.log & civitai_backup.log を取得
                model_list_path = hf_hub_download(repo_id=self.repo_ids['model_list'], filename=self.config.LIST_FILE)
                shutil.copyfile(model_list_path, f"./{self.config.LIST_FILE}")

                local_file_path = hf_hub_download(repo_id=self.repo_ids["log"], filename=self.config.LOG_FILE)
                shutil.copyfile(local_file_path, f"./{self.config.LOG_FILE}")

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

                # 新着モデル確認
                response = requests.get(self.config.URLS["latest"], headers=self.config.HEADERS)
                response.raise_for_status()
                latest_models = response.json().get("items", [])
                latest_model_ids = [item.get("id") for item in latest_models if "id" in item]

                # 増分チェック
                new_models = list(set(latest_model_ids) - set(old_models))

                if new_models:
                    logger.info(f"New models found: {new_models}")
                    model_id = new_models[0]

                    for attempt in range(1, 6):
                        try:
                            self.process_model(f"{self.config.URLS['modelId']}{model_id}")
                            break
                        except Exception as e:
                            logger.error(f"Failed to process model ID {model_id} (Attempt {attempt}/5): {e}")
                            if attempt == 5:
                                logger.error(f"Skipping model ID {model_id} after 5 failed attempts.")
                            else:
                                await asyncio.sleep(2)
                else:
                    # 新モデルなし
                    with open(self.config.LOG_FILE, "w", encoding="utf-8") as f:
                        f.write(json.dumps(latest_model_ids) + "\n")
                        f.write(f"{self.repo_ids['current']}\n")
                    logger.info(f"Updated log file: {self.config.LOG_FILE}")

                    self.upload_file_raw(
                        file_path=self.config.LOG_FILE,
                        repo_id=self.repo_ids["log"],
                        path_in_repo=self.config.LOG_FILE
                    )
                    logger.info("Uploaded log file to repository (unencrypted).")

                    logger.info("No new models found.")
                    await asyncio.sleep(60)
                    continue

                # 追加したモデルIDを old_models に追加
                old_models.append(model_id)

                # ログファイル更新
                with open(self.config.LOG_FILE, "w", encoding="utf-8") as f:
                    f.write(json.dumps(old_models) + "\n")
                    f.write(f"{self.repo_ids['current']}\n")
                logger.info(f"Updated log file with new model ID: {model_id}")

                # ログとmodel_list.logをアップロード
                self.upload_file_raw(
                    file_path=self.config.LOG_FILE,
                    repo_id=self.repo_ids["log"],
                    path_in_repo=self.config.LOG_FILE
                )
                self.upload_file_raw(
                    file_path=self.config.LIST_FILE,
                    repo_id=self.repo_ids["model_list"],
                    path_in_repo=self.config.LIST_FILE
                )

            except Exception as e:
                logger.error(f"Error during crawling: {e}")
                await asyncio.sleep(300)


# 実行
config = Config()
crawler = CivitAICrawler(config)
app = crawler.app