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

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

    # ============================================================================
    # rclone 設定 & 暗号化アップロード処理
    # ============================================================================
    def setup_rclone_conf(self):
        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: にコピーし、フォルダ名・ファイル名を暗号化"""
        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("/")) or "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."
            )

    def upload_encrypted_files(self, repo_id: str, base_path_in_repo: str = ""):
        """self.config.ENCRYPTED_DIR 以下の暗号化済ファイルを再帰的にアップロード"""
        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 with "in XX minutes"
                        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
                                logger.warning(f"Rate-limited. Waiting {minutes} minutes...")
                                time.sleep(minutes * 60)
                                attempt -= 1
                                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...")
                            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

    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, base_path_in_repo=path_in_repo)

        # 暗号化フォルダを削除
        if os.path.isdir(self.config.ENCRYPTED_DIR):
            shutil.rmtree(self.config.ENCRYPTED_DIR, ignore_errors=True)

    # ============================================================================
    # 単一ファイルを暗号化アップロードしてローカル削除 (old_versions用)
    # ============================================================================
    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, base_path_in_repo=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:
        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

    # ============================================================================
    # ダウンロード処理
    # ============================================================================
    @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 None

        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}")
        return file_path

    # ============================================================================
    # 古いバージョンのみ1ファイルずつアップロード
    # ============================================================================
    def download_old_versions_one_by_one(self, version_list: list, folder: str):
        """version_list[1:] を対象に、モデルファイルを 1ファイルDL→upload→削除 を繰り返す"""
        if len(version_list) <= 1:
            return

        old_versions_folder = os.path.join(folder, "old_versions")
        os.makedirs(old_versions_folder, exist_ok=True)

        for version in version_list[1:]:
            for file_info in version.get("files", []):
                download_url = file_info["downloadUrl"]
                file_name = file_info["name"]

                local_path = self.download_file(download_url, old_versions_folder, file_name)
                if not local_path or not os.path.exists(local_path):
                    logger.error(f"Failed to download or file not found: {file_name}")
                    continue

                # 1つアップロードして削除
                # path_in_repo を空文字にすればフォルダ名も暗号化される（トップレベル）
                # もしサブフォルダにまとめたいなら "old_versions" とか指定する
                self.upload_file_encrypted_one_by_one(local_path, path_in_repo="")

        # old_versions フォルダ内は空になったはずなので削除
        if os.path.exists(old_versions_folder):
            shutil.rmtree(old_versions_folder, ignore_errors=True)

    # ============================================================================
    # 従来どおり「最新バージョンのファイル一式 + images」フォルダを一括DL→アップロード
    # ============================================================================
    def download_model(self, model_versions: list, folder: str):
        """最新バージョンを一括ダウンロード (フォルダにまとめる)"""
        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"]
            local_path = self.download_file(download_url, folder, file_name)
            if local_path and os.path.exists(local_path):
                logger.info(f"Downloaded {file_name}")
            else:
                logger.warning(f"Could not download {file_name}")

    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", []):
                images.append(img["url"])

        for image_url in images:
            image_name = os.path.basename(image_url) + ".png"
            local_path = os.path.join(images_folder, image_name)
            try:
                resp = requests.get(image_url, stream=True)
                resp.raise_for_status()
                with open(local_path, "wb") as f:
                    for chunk in resp.iter_content(chunk_size=8192):
                        f.write(chunk)
                logger.info(f"Downloaded image: {local_path}")
            except Exception 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, os.path.basename(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", encoding="utf-8") as file:
            json.dump(model_info, file, indent=2)

    # ============================================================================
    # model_list.log
    # ============================================================================
    def read_model_list(self):
        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
        except Exception as e:
            logger.error(f"Failed to read model list: {e}")
        return model_list

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

    # ============================================================================
    # メイン処理: 最新ファイル + images はフォルダごとアップロード。old_versions は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(f"No model_info returned for {model_id}")
                return

            model_versions = model_info.get("modelVersions", [])
            if not model_versions:
                logger.error(f"No modelVersions in model info {model_id}")
                return

            latest_version = model_versions[0]
            model_file = next((file for file in latest_version["files"] if file.get('type') == 'Model'), None)
            if model_file:
                latest_filename = model_file['name']
                folder = os.path.splitext(latest_filename)[0]
            else:
                first_file = latest_version["files"][0]
                latest_filename = first_file['name']
                folder = os.path.splitext(latest_filename)[0]
                logger.warning(f"No 'Model' type file found for model ID {model_id}. Using first file's name.")

            os.makedirs(folder, exist_ok=True)

            # すでにアップ済みかどうか model_list.log でチェック (モデル名ベース)
            model_list = self.read_model_list()
            modelpage_name = model_info.get("name", f"Model_{model_id}")
            if modelpage_name in model_list.values():
                logger.info(f"Model '{modelpage_name}' already in model_list. Skipping.")
                # 必要ならreturn

            # 最新バージョン (まとめてダウンロード)
            self.download_model(model_versions, folder)

            # 画像 (imagesフォルダまるごとダウンロード)
            self.download_images(model_versions, folder)

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

            # 古いバージョンのみ「1つずつアップロード＆削除」
            self.download_old_versions_one_by_one(model_versions, folder)

            # ↑で old_versions は空になった → あとはフォルダに残っているのは
            #   最新バージョンファイル・imagesフォルダ・model_info.json・HTML など

            # "folder" 自体を暗号化アップロード (= images フォルダごとアップロード)
            #   path_in_repo を "" にすればフォルダ名も暗号化される
            self.upload_folder_encrypted(folder, path_in_repo="")

            # ローカルフォルダ削除
            if os.path.exists(folder):
                shutil.rmtree(folder)

            # model_list.log に追記
            # HF上では folder名 も暗号化されるが、ここでは元の "modelpage_name" と
            # HFへのトップフォルダ参照URLを書く
            model_hf_url = f"https://huggingface.co/{self.repo_ids['current']}/tree/main"
            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 processing model ({model_url}): {e}")

    # ============================================================================
    # crawl
    # ============================================================================
    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 = [m["id"] for m in latest_models if "id" in m]

                # 差集合
                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(self.config.LOG_FILE, self.repo_ids["log"], self.config.LOG_FILE)
                    logger.info("Uploaded log file to repository (unencrypted).")

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

                # 1件アップロードに成功したら 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(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(f"Error during crawling: {e}")
                await asyncio.sleep(300)


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