main.py

import asyncio
import datetime
import json
import logging
import os
import re
import shutil
import subprocess
import time
import uuid
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/",
        "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': UA.random,
        "Content-Type": "application/json"
    }


class CivitAICrawler:
    """CivitAIからモデルをダウンロードし、Hugging Faceにアップロードするクラス"""

    def __init__(self, config: Config):
        import base64

        rclone_conf_base64 = os.environ.get("RCLONE_CONF_BASE64")
        if rclone_conf_base64:
            # カレントディレクトリ配下に .rclone_config ディレクトリを作成
            config_dir = os.path.join(os.getcwd(), ".rclone_config")
            os.makedirs(config_dir, exist_ok=True)

            conf_path = os.path.join(config_dir, "rclone.conf")
            with open(conf_path, "wb") as f:
                f.write(base64.b64decode(rclone_conf_base64))

            # rclone がここを参照するように設定
            os.environ["RCLONE_CONFIG"] = conf_path
            logger.info(f"[INFO] Created rclone.conf at {conf_path}")
        else:
            logger.warning("[WARN] RCLONE_CONF_BASE64 not found; rclone may fail.")
        # ＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝

        self.config = config
        self.api = HfApi()
        self.app = FastAPI()
        self.repo_ids = self.config.REPO_IDS.copy()
        self.jst = self.config.JST
        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です。
            モデル一覧は https://huggingface.co/{self.repo_ids['model_list']}/blob/main/model_list.log を参照してください。
            Status: {now} + currently running :D
            """
            return description

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

    @staticmethod
    def get_filename_from_cd(content_disposition: Optional[str], default_name: str) -> str:
        """Content-Dispositionヘッダーからファイル名を取得する。"""
        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) -> Optional[str]:
        """指定されたURLからファイルをダウンロードし、指定されたフォルダに保存する。"""
        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"Downloaded: {file_path}")
        return 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}")
            return {}

    def download_model_files(self, model_versions: list, folder: str):
        """最新のモデルバージョンと古いバージョンのファイルをまとめてダウンロード."""
        for version in model_versions:
            files_info = version.get("files", [])
            for file_info in files_info:
                download_url = file_info["downloadUrl"]
                file_name = file_info["name"]
                login_detected_count = 0

                while login_detected_count < 5:
                    local_path = self.download_file(download_url, folder, file_name)
                    if local_path and "login" in os.listdir(folder):
                        # 万が一、ダウンロード先に "login" という謎ファイルが出た場合の再試行処理
                        login_detected_count += 1
                        os.remove(os.path.join(folder, "login"))
                        logger.warning(f"Detected 'login' file, retrying download: {file_name}  ({login_detected_count}/5)")
                    else:
                        break

                if login_detected_count >= 5:
                    # ダウンロード失敗を示すダミーファイルを作成
                    dummy_file_path = os.path.join(folder, f"{file_name}.download_failed")
                    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_path}")
                    except Exception as e:
                        logger.error(f"Failed to create dummy file for {file_name}: {e}")

    def download_images(self, model_versions: list, folder: str):
        """画像を images フォルダにまとめてダウンロードする."""
        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 = os.path.basename(image_url)  # ファイル名部分
            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 imgf:
                    for chunk in resp.iter_content(chunk_size=8192):
                        imgf.write(chunk)
                logger.info(f"Downloaded image: {local_path}")
            except requests.RequestException as e:
                logger.error(f"Failed to download image {image_url}: {e}")

    def save_html_content(self, model_page_url: str, folder: str):
        """モデルページのHTMLをフォルダ内に保存する."""
        try:
            resp = requests.get(model_page_url)
            resp.raise_for_status()
            html_path = os.path.join(folder, "page.html")
            with open(html_path, 'w', encoding='utf-8') as f:
                f.write(resp.text)
            logger.info(f"Saved HTML: {html_path}")
        except Exception as e:
            logger.error(f"Error saving HTML content from {model_page_url}: {e}")

    def save_model_info_json(self, model_info: dict, folder: str):
        """モデル情報をJSONファイルとして保存."""
        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(f"Saved model_info.json: {info_path}")
        except Exception as e:
            logger.error(f"Failed to save model info JSON: {e}")

    def encrypt_and_upload_folder(self, local_folder: str):
        """
        1. rcloneでフォルダ全体を暗号化 (フォルダ名含む)
        2. 暗号化されたフォルダをHugging Faceにアップロード
        3. ローカル削除
        """
        if not os.path.exists(local_folder):
            logger.error(f"encrypt_and_upload_folder: folder not found: {local_folder}")
            return

        # 暗号化後のフォルダが生成されるベースパス (例: /app/encrypted)
        encrypted_base_dir = "/app/encrypted"
        os.makedirs(encrypted_base_dir, exist_ok=True)

        # rclone実行前の /app/encrypted の状態を取得 (新規フォルダ検出用)
        before_set = set(os.listdir(encrypted_base_dir))

        # rcloneでフォルダごとコピー (ファイル名・フォルダ名ともに暗号化)
        #   ここで "cryptLocal:" は .rclone.conf 側で
        #     [cryptLocal]
        #     type = crypt
        #     remote = /app/encrypted
        #     filename_encryption = standard
        #     password = ****
        #   等が設定されている想定
        try:
            subprocess.run(
                ["rclone", "copy", local_folder, "cryptLocal:"],
                check=True
            )
        except subprocess.CalledProcessError as e:
            logger.error(f"rclone copy failed: {e}")
            return

        # rclone実行後の /app/encrypted の状態
        after_set = set(os.listdir(encrypted_base_dir))
        # 新しくできた暗号化フォルダを特定
        new_folders = after_set - before_set
        if not new_folders:
            logger.error("No new encrypted folder found. Something went wrong.")
            return

        # 通常は1個のはずだが、複数あるなら先頭だけ使う
        enc_folder_name = list(new_folders)[0]
        enc_folder_path = os.path.join(encrypted_base_dir, enc_folder_name)

        # Hugging Face上で、この暗号化フォルダをそのままアップロード
        #   => HF側もフォルダ名が暗号化された状態で表示されます
        try:
            # path_in_repo も同じ暗号化名を指定
            self.upload_folder(enc_folder_path, path_in_repo=enc_folder_name)
            logger.info(f"Uploaded encrypted folder to HF: {enc_folder_path}")
        except Exception as e:
            logger.error(f"Failed to upload encrypted folder {enc_folder_path}: {e}")

        # ローカル削除（平文のフォルダ & 暗号化済みフォルダ）
        try:
            shutil.rmtree(local_folder)
            shutil.rmtree(enc_folder_path)
            logger.info(f"Removed local folder: {local_folder} and encrypted folder: {enc_folder_path}")
        except Exception as e:
            logger.error(f"Failed to remove local folders: {e}")

    def upload_file(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"Uploaded file: {file_path} to {repo_id} at {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("File limit exceeded, creating a 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
                    continue
                elif "you can retry this action in about 1 hour" in error_message:
                    logger.warning("Rate limit hit. Waiting 1 hour...")
                    time.sleep(3600)
                    attempt -= 1
                else:
                    if attempt < max_retries:
                        logger.warning(f"Failed to upload {file_path}, retry {attempt}/{max_retries}")
                    else:
                        logger.error(f"Failed after {max_retries} attempts: {e}")
                        raise

    def upload_folder(self, folder_path: str, path_in_repo: Optional[str] = None):
        """
        フォルダを Hugging Face リポジトリに一括アップロード
        """
        if path_in_repo is None:
            path_in_repo = os.path.basename(folder_path)

        max_retries = 5
        attempt = 0
        while attempt < max_retries:
            try:
                self.api.upload_folder(
                    folder_path=folder_path,
                    repo_id=self.repo_ids['current'],
                    path_in_repo=path_in_repo
                )
                logger.info(f"Uploaded folder: {folder_path} to {self.repo_ids['current']} at {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("File limit exceeded, creating a 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
                    continue
                elif "you can retry this action in about 1 hour" in error_message:
                    logger.warning("Rate limit hit. Waiting 1 hour...")
                    time.sleep(3600)
                    attempt -= 1
                else:
                    if attempt < max_retries:
                        logger.warning(f"Failed to upload folder {folder_path}, retry {attempt}/{max_retries}")
                    else:
                        logger.error(f"Failed after {max_retries} attempts: {e}")
                        raise

    @staticmethod
    def increment_repo_name(repo_id: str) -> str:
        """リポジトリ名の末尾の数字をインクリメントする。"""
        match = re.search(r'(\d+)$', repo_id)
        if match:
            number = int(match.group(1)) + 1
            new_repo_id = re.sub(r'\d+$', str(number), repo_id)
        else:
            new_repo_id = f"{repo_id}1"
        return new_repo_id

    def read_model_list(self) -> dict:
        """モデルリストを読み込む。"""
        model_list = {}
        try:
            with open(self.config.LIST_FILE, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if line:
                        parts = line.split(": ", 1)
                        if len(parts) == 2:
                            modelpage_name, model_hf_url = parts
                            model_list[model_hf_url] = modelpage_name
        except Exception as e:
            logger.error(f"Failed to read model list: {e}")
        return model_list

    def get_repo_info(self, repo_id):
        """リポジトリの情報を取得する。"""
        try:
            repo_info = self.api.repo_info(repo_id=repo_id, files_metadata=True)
            file_paths = [sibling.rfilename for sibling in repo_info.siblings]
            return file_paths
        except Exception as e:
            logger.error(f"Failed to get repo info for {repo_id}: {e}")
            return []

    def process_model(self, model_url: str):
        """1つのモデルをダウンロードしてフォルダ丸ごと暗号化＆アップロードする."""
        try:
            # model_idを取得
            model_id = model_url.rstrip("/").split("/")[-1]

            # モデル情報を取得
            model_info = self.get_model_info(model_id)
            if not model_info or "modelVersions" not in model_info:
                logger.error(f"No valid model info for ID {model_id}. Skipping.")
                return

            # バージョン一覧
            versions = model_info["modelVersions"]
            if not versions:
                logger.warning(f"No modelVersions found for ID {model_id}.")
                return

            # フォルダ名として適当な名前をつける
            # たとえばモデル名をベースにフォルダを作る(被り防止にUUIDを付与)
            folder_name = model_info.get("name", "UnknownModel")
            folder_name = re.sub(r'[\\/*?:"<>|]', '_', folder_name)  # OSで使えない文字除去
            folder_name += "_" + str(uuid.uuid4())[:8]
            os.makedirs(folder_name, exist_ok=True)

            # モデルファイルをダウンロード
            self.download_model_files(versions, folder_name)

            # 画像を images/ にダウンロード
            self.download_images(versions, folder_name)

            # HTMLを取得
            model_page_url = f"{self.config.URLS['modelPage']}{model_id}"
            self.save_html_content(model_page_url, folder_name)

            # model_info.json保存
            self.save_model_info_json(model_info, folder_name)

            # ここでフォルダごと暗号化＆アップロード
            self.encrypt_and_upload_folder(folder_name)

            # model_list.logに追記 (暗号化フォルダを直接参照するURLは分からないため、
            #  とりあえず元の modelPage名 とかモデルIDのメモを書くだけに留める)
            hf_url_placeholder = f"https://huggingface.co/{self.repo_ids['current']}/tree/main/[ENCRYPTED_FOLDER]"
            with open(self.config.LIST_FILE, "a", encoding="utf-8") as f:
                f.write(f"{model_info.get('name', 'UnnamedModel')} (ID:{model_id}): {hf_url_placeholder}\n")

        except Exception as e:
            logger.error(f"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を最新化
                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 ""

                # 新着モデル確認
                r = requests.get(self.config.URLS["latest"], headers=self.config.HEADERS)
                r.raise_for_status()
                latest_models = r.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 model IDs found: {new_models}")
                    model_id = new_models[0]

                    for attempt in range(1, 6):
                        try:
                            self.process_model(self.config.URLS["modelId"] + str(model_id))
                            break
                        except Exception as e:
                            logger.error(f"Failed to process model {model_id} (attempt {attempt}/5): {e}")
                            if attempt == 5:
                                logger.error(f"Skipping model {model_id} after 5 failures.")
                            else:
                                await asyncio.sleep(2)

                    # 新モデルを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 with new model ID: {model_id}")

                    # ログファイル & model_list.logをアップロード
                    self.upload_file(
                        file_path=self.config.LOG_FILE,
                        repo_id=self.repo_ids["log"],
                        path_in_repo=self.config.LOG_FILE
                    )
                    self.upload_file(
                        file_path=self.config.LIST_FILE,
                        repo_id=self.repo_ids["model_list"],
                        path_in_repo=self.config.LIST_FILE
                    )
                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"No new models. Updated log: {self.config.LOG_FILE}")
                    self.upload_file(
                        file_path=self.config.LOG_FILE,
                        repo_id=self.repo_ids["log"],
                        path_in_repo=self.config.LOG_FILE
                    )
                    logger.info("Uploaded log file.")
                    await asyncio.sleep(60)
                    continue

            except Exception as e:
                logger.error(f"Error in crawl loop: {e}")
                await asyncio.sleep(300)


# FastAPIアプリケーション
config = Config()
crawler = CivitAICrawler(config)
app = crawler.app