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:
            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))

            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_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", []):
                images.append(img["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, f"{os.path.basename(folder)}.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 download_and_process_versions(self, model_versions: list, folder: str):
        """
        最新バージョン + 古いバージョンをすべて1つのフォルダにダウンロードして
        最後にまとめて暗号化アップロードする。
        """

        # 1) 最新バージョンは folder に
        latest_version = model_versions[0]
        for file_info in latest_version.get("files", []):
            download_url = file_info["downloadUrl"]
            file_name = file_info["name"]
            local_path = self.download_file(download_url, folder, file_name)
            # リトライ処理など省略

        # 2) 古いバージョンを "folder/old_versions" にまとめる
        if len(model_versions) > 1:
            old_versions_folder = os.path.join(folder, "old_versions")
            os.makedirs(old_versions_folder, exist_ok=True)

            for version in model_versions[1:]:
                for file_info in version.get("files", []):
                    file_name = file_info["name"]
                    download_url = file_info["downloadUrl"]
                    local_path = self.download_file(download_url, old_versions_folder, file_name)

    def encrypt_and_upload_folder(self, local_folder: str) -> Optional[str]:
        """
        1. rclone copy local_folder => cryptLocal:   (フォルダ名はrcloneが勝手に暗号化生成)
        2. 差分検知で "./encrypted" に作られた暗号フォルダ名を取得
        3. そのフォルダ名を path_in_repo として Hugging Face にアップロード
        4. ローカル(平文+暗号)フォルダ削除
        5. 戻り値は "実際に作られた暗号フォルダ名"
        """
        if not os.path.isdir(local_folder):
            logger.error(f"[encrypt_and_upload_folder] Not a directory: {local_folder}")
            return None

        encrypted_base_dir = os.path.join(os.getcwd(), "encrypted")
        os.makedirs(encrypted_base_dir, exist_ok=True)

        # 既存の暗号フォルダを削除
        before_set = set(os.listdir(encrypted_base_dir))
        for item in before_set:
            item_path = os.path.join(encrypted_base_dir, item)
            try:
                if os.path.isfile(item_path):
                    os.remove(item_path)
                else:
                    shutil.rmtree(item_path)
                logger.info(f"[CLEANUP] Removed old encrypted item: {item_path}")
            except Exception as e:
                logger.warning(f"[CLEANUP] Failed to remove {item_path}: {e}")

        # === (1) rclone copy local_folder => cryptLocal: ===
        #     これにより /home/user/app/encrypted/<ランダム暗号フォルダ> が作成される
        try:
            subprocess.run(
                ["rclone", "copy", local_folder, "cryptLocal:", "--create-empty-src-dirs"],
                check=True
            )
            logger.info(f"[OK] rclone copy {local_folder} => cryptLocal:")
        except subprocess.CalledProcessError as e:
            logger.error(f"rclone copy failed: {e}")
            return None

        # === (2) 差分検知: 新しく生成された暗号フォルダ名を取得 ===
        after_set = set(os.listdir(encrypted_base_dir))
        diff = after_set - before_set
        if not diff:
            logger.error("[ERROR] No new directory appeared in ./encrypted after rclone copy.")
            return None
        if len(diff) > 1:
            logger.warning(f"[WARN] Multiple new directories created: {diff}")

        enc_folder_name = diff.pop()  # 1個だけ取り出す
        enc_folder_path = os.path.join(encrypted_base_dir, enc_folder_name)

        if not os.path.isdir(enc_folder_path):
            logger.error(f"[ERROR] {enc_folder_path} is not a directory.")
            return None

        # === (3) Hugging Face にアップロード ===
        #     path_in_repo も "enc_folder_name" をそのまま使う
        try:
            self.upload_folder(enc_folder_path, path_in_repo=enc_folder_name)
            logger.info(f"[OK] Uploaded encrypted folder: {enc_folder_path}")
        except Exception as e:
            logger.error(f"Failed to upload encrypted folder {enc_folder_path}: {e}")

        # === (4) ローカル削除 (平文フォルダ + 暗号化フォルダ)
        try:
            shutil.rmtree(local_folder)
            shutil.rmtree(enc_folder_path)
            logger.info(f"[CLEANUP] Removed local folder: {local_folder} & {enc_folder_path}")
        except Exception as e:
            logger.warning(f"[CLEANUP] Could not remove local folders: {e}")

        # === (5) 実際の暗号フォルダ名を返す
        return enc_folder_name

    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):
        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_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

            folder_name = model_info.get("name", "UnnamedModel")
            folder_name = re.sub(r'[\\/*?:"<>|]', '_', folder_name)
            folder_name += "_" + str(uuid.uuid4())[:8]
            os.makedirs(folder_name, exist_ok=True)

            # 最新 + 古いバージョンすべて folder_name 下へ
            self.download_and_process_versions(versions, folder_name)
            self.download_images(versions, folder_name)
            self.save_html_content(f"{self.config.URLS['modelPage']}{model_id}", folder_name)
            self.save_model_info_json(model_info, folder_name)

            # フォルダごと暗号化
            enc_subfolder = self.encrypt_and_upload_folder(folder_name)
            if enc_subfolder is None:
                enc_subfolder = "[ENCRYPT_FAILED]"

            hf_enc_url = f"https://huggingface.co/{self.repo_ids['current']}/tree/main/{enc_subfolder}"
            with open(self.config.LIST_FILE, "a", encoding="utf-8") as f:
                f.write(f"{model_info.get('name', 'UnnamedModel')} (ID:{model_id}): {hf_enc_url}\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.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}")

                    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