import os
import argparse
import re
import json
import random
import threading
from typing import Dict, Tuple, List, Optional, Set

import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry


# ============================================================
# 1. 基础配置
# ============================================================

BASE_URL = "https://modelhub.org.cn"

LOGIN_ENDPOINT = "/adminApi/user/login"
ADD_ADAPT_TASK_ENDPOINT = "/api/adapt/task/add"
MODELHUB_DB_CHECK_ENDPOINT = "/api/computility/models/list/page/vo"

TASK_TYPE = "text-generation"
FRAMEWORK = "llamacpp"

# 按当前策略部署约定，ModelHub 登录凭据固定写入策略代码。
USER_ACCOUNT = "1326356020@qq.com"
USER_PASSWORD = "4pdpassword"

# 不使用手动 token，统一用账号密码登录获取 token
MODELHUB_TOKEN = ""

# 平台环境变量契约：提交任务时必须携带当前策略 ID。
STRATEGY_ID = os.environ["STRATEGY_ID"]


# ============================================================
# 2. Hugging Face / hf-mirror 筛选配置
# ============================================================

ORG_NAME = "mradermacher"
HF_API_URL = "https://hf-mirror.com/api/models"

PER_PAGE = 100
START_PAGE = 6  # 起始页码
END_PAGE = 6    # 结束页码
REQUEST_DELAY = 1  # 每次请求的基础延迟（秒），可根据情况调整
RANDOM_DELAY_RANGE = (0.5, 2)  # 随机延迟范围（秒），避免固定间隔被识别
RETRY_TIMES = 3  # 429错误重试次数
RETRY_DELAY = 5  # 每次重试的等待时间（秒）

MAX_FILE_SIZE_GB = 12
# ============================================================
# 3. 提交与轮询配置
# ============================================================

# 如果所有模型都想强制用某种 GPU，就填：
# FORCE_TARGET_GPU = "Ascend_910-b4"
# 如果保留 None，则 i1/i2/... 模型走 Ascend_910-b4，其他走 Mthreads_s4000
FORCE_TARGET_GPU = None

SUBMIT_INTERVAL_SECONDS = 2

TASK_LIMIT_CODE = 60007
POLL_INTERVAL_SECONDS = 120
MAX_LIMIT_RETRY_TIMES = None  # None 表示一直等到有空位

DEBUG_PRINT_PAYLOAD = False

# 是否开启断点续跑：已经成功提交过的模型，下次运行自动跳过
RESUME = True

# ============================================================
# 断点续跑 / 重跑配置
# ============================================================
# run_mode 说明：
#   resume：优先复用已有步骤产物；缺失时才运行该步骤。适合 Step 4 报错后接着跑。
#   fresh ：从 start_step 开始重跑，并覆盖对应步骤产物；start_step 之前的步骤读取已有产物。
# 命令行参数优先级高于这里的默认值，例如：
#   python modelhub_pipeline_resumable.py --start-step 4 --run-mode resume
DEFAULT_RUN_MODE = os.getenv("PIPELINE_RUN_MODE", "resume")
DEFAULT_START_STEP = int(os.getenv("PIPELINE_START_STEP", "1"))
DEFAULT_FORCE_RERUN_STEPS = os.getenv("PIPELINE_FORCE_RERUN_STEPS", "")

OUTPUT_DIR = "pipeline_outputs"

ALL_MODELS_FILE = os.path.join(OUTPUT_DIR, "01_all_models.txt")
SIZE_FILTERED_FILE = os.path.join(OUTPUT_DIR, "02_size_filtered_models.txt")
NOT_IN_DB_FILE = os.path.join(OUTPUT_DIR, "03_not_in_db_models.txt")
SUBMITTED_FILE = os.path.join(OUTPUT_DIR, "04_submitted_models.txt")
FAILED_FILE = os.path.join(OUTPUT_DIR, "05_failed_models.txt")
SUMMARY_FILE = os.path.join(OUTPUT_DIR, "summary.json")


BASE_HEADERS = {
    "Content-Type": "application/json",
    "Accept": "application/json, text/plain, */*",
    "Origin": BASE_URL,
    "Referer": BASE_URL + "/",
    "User-Agent": "Mozilla/5.0",
}

_shutdown_event = threading.Event()


class ShutdownRequested(Exception):
    """平台请求策略进程优雅停止。"""


def request_shutdown() -> None:
    _shutdown_event.set()


def reset_shutdown() -> None:
    _shutdown_event.clear()


def check_shutdown() -> None:
    if _shutdown_event.is_set():
        raise ShutdownRequested("收到停机信号")


def interruptible_sleep(seconds: float) -> None:
    if _shutdown_event.wait(max(0, seconds)):
        raise ShutdownRequested("收到停机信号")


# ============================================================
# 4. 通用工具
# ============================================================

def ensure_output_dir() -> None:
    os.makedirs(OUTPUT_DIR, exist_ok=True)


def write_model_list(path: str, models: List[str]) -> None:
    with open(path, "w", encoding="utf-8") as f:
        for model in models:
            f.write(f'"{model}",\n')


def append_model(path: str, model: str) -> None:
    with open(path, "a", encoding="utf-8") as f:
        f.write(f'"{model}",\n')


def load_model_set(path: str) -> Set[str]:
    if not os.path.exists(path):
        return set()

    with open(path, "r", encoding="utf-8") as f:
        content = f.read()

    quoted_items = re.findall(r'["\']([^"\']+)["\']', content)
    if quoted_items:
        return set(x.strip() for x in quoted_items if x.strip())

    result = set()
    for line in content.splitlines():
        item = line.strip().strip(",").strip('"').strip("'").strip()
        if item:
            result.add(item)
    return result


def load_model_list(path: str) -> List[str]:
    """读取模型列表，保持顺序并去重。"""
    if not os.path.exists(path):
        raise FileNotFoundError(
            f"找不到断点文件: {path}。如果要从中间步骤继续，请先确认前置步骤已成功生成该文件。"
        )

    with open(path, "r", encoding="utf-8") as f:
        content = f.read()

    quoted_items = re.findall(r'["\']([^"\']+)["\']', content)
    if quoted_items:
        raw_items = [x.strip() for x in quoted_items if x.strip()]
    else:
        raw_items = []
        for line in content.splitlines():
            item = line.strip().strip(",").strip('"').strip("'").strip()
            if item:
                raw_items.append(item)

    seen = set()
    result = []
    for item in raw_items:
        if item not in seen:
            result.append(item)
            seen.add(item)
    return result


def parse_bool_env(value: str, default: bool = False) -> bool:
    if value is None or value == "":
        return default
    return value.strip().lower() in {"1", "true", "yes", "y", "on"}


def parse_step_set(value: str) -> Set[int]:
    result: Set[int] = set()
    if not value:
        return result

    for item in value.split(","):
        item = item.strip()
        if not item:
            continue
        try:
            step = int(item)
        except ValueError:
            raise ValueError(f"PIPELINE_FORCE_RERUN_STEPS 里包含非法步骤: {item}")
        if step not in {1, 2, 3}:
            raise ValueError("可强制重跑的步骤只能是 1、2、3；Step 4 每次都会按断点提交剩余模型。")
        result.add(step)
    return result


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="HF / hf-mirror -> ModelHub 断点续跑脚本"
    )
    parser.add_argument(
        "--run-mode",
        choices=["resume", "fresh"],
        default=DEFAULT_RUN_MODE,
        help="resume=复用已有步骤结果；fresh=从 start-step 开始重跑。默认读取 PIPELINE_RUN_MODE 或 resume。",
    )
    parser.add_argument(
        "--start-step",
        type=int,
        choices=[1, 2, 3, 4],
        default=DEFAULT_START_STEP,
        help="从第几步开始。Step 4 报错后继续跑用 --start-step 4。默认读取 PIPELINE_START_STEP 或 1。",
    )
    parser.add_argument(
        "--force-rerun-steps",
        default=DEFAULT_FORCE_RERUN_STEPS,
        help="在 resume 模式下强制重跑指定前置步骤，例如 '2,3'。Step 4 不需要写入这里。",
    )
    parser.add_argument(
        "--reset-submitted",
        action="store_true",
        help="清空 04_submitted_models.txt。谨慎使用：会导致已成功提交过的模型不再被跳过。",
    )
    parser.add_argument(
        "--reset-failed",
        action="store_true",
        help="清空 05_failed_models.txt。适合修复参数后重新统计失败模型。",
    )
    parser.add_argument(
        "--skip-failed",
        action="store_true",
        help="跳过 05_failed_models.txt 中记录过失败的模型。默认不跳过，便于修复问题后自动重试失败模型。",
    )
    args = parser.parse_args()
    args.force_rerun_steps = parse_step_set(args.force_rerun_steps)
    return args


def should_run_step(step: int, output_file: str, args: argparse.Namespace) -> bool:
    """
    判断 Step 1-3 是否需要执行。

    resume：如果产物存在就读取，不存在才执行；force-rerun-steps 可强制执行。
    fresh ：start_step 之前读取已有产物，从 start_step 开始重新执行并覆盖产物。
    """
    if step in args.force_rerun_steps:
        return True

    if args.run_mode == "fresh":
        return step >= args.start_step

    # resume mode
    if step < args.start_step:
        return False

    return not os.path.exists(output_file)


def print_step_loaded(step: int, path: str, models: List[str]) -> None:
    print(f"Step {step}/4: 跳过执行，读取已有断点文件: {path}")
    print(f"读取到 {len(models)} 个模型")


def remove_file_if_exists(path: str) -> None:
    if os.path.exists(path):
        os.remove(path)
        print(f"已清空: {path}")


def create_hf_session() -> requests.Session:
    session = requests.Session()

    retry_strategy = Retry(
        total=RETRY_TIMES,
        backoff_factor=RETRY_DELAY,
        status_forcelist=[429, 500, 502, 503, 504],
        allowed_methods=["GET"],
    )

    adapter = HTTPAdapter(max_retries=retry_strategy)
    session.mount("https://", adapter)
    session.mount("http://", adapter)

    session.headers.update({
        "User-Agent": (
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
            "AppleWebKit/537.36 (KHTML, like Gecko) "
            "Chrome/120.0.0.0 Safari/537.36"
        )
    })

    return session


def get_model_filename(model_id: str) -> str:
    """
    从 model_id 生成 GGUF 文件名。

    mradermacher/Qwen3-8B-xxx-i1-GGUF
      -> Qwen3-8B-xxx.i1-Q4_0.gguf

    QuantFactory/Apollo2-9B-GGUF
      -> Apollo2-9B.Q8_0.gguf
    """
    base_name = model_id.split("/")[-1]

    if "_-_" in base_name:
        base_name = base_name.split("_-_")[-1]

    if base_name.lower().endswith("-gguf"):
        base_name = base_name[:-5]

    match = re.search(r"-i(\d+)$", base_name)
    if match:
        number = match.group(1)
        return base_name[:match.start()] + f".i{number}-Q4_0.gguf"

    return base_name + ".Q8_0.gguf"


# ============================================================
# 5. Step 1：从 hf-mirror 获取组织模型
# ============================================================

def get_org_models(session: requests.Session, org_name: str) -> List[str]:
    models: List[str] = []

    for page in range(START_PAGE, END_PAGE + 1):
        try:
            interruptible_sleep(random.uniform(*RANDOM_DELAY_RANGE))

            params = {
                "author": org_name,
                "page": page,
                "perPage": PER_PAGE,
                "sort": "lastModified",
                "direction": "-1",
            }

            response = session.get(HF_API_URL, params=params, timeout=20)
            response.raise_for_status()
            data = response.json()

            if not data:
                print(f"第 {page} 页无数据，停止获取")
                break

            page_models = []
            for item in data:
                model_id = item.get("modelId")
                if model_id:
                    page_models.append(model_id)
                    models.append(model_id)

            print(f"成功获取第 {page} 页，共 {len(page_models)} 个模型")

        except ShutdownRequested:
            raise
        except Exception as e:
            print(f"获取第 {page} 页失败: {e}")
            continue

    # 去重但保持顺序
    seen = set()
    unique_models = []
    for model in models:
        if model not in seen:
            unique_models.append(model)
            seen.add(model)

    return unique_models


# ============================================================
# 6. Step 2：按目标 GGUF 文件大小筛选
# ============================================================

def extract_gguf_file_size(
    session: requests.Session,
    model_id: str,
    target_filename: str,
) -> float:
    """
    返回 GB。
    找不到或失败返回 -1。
    """
    interruptible_sleep(REQUEST_DELAY + random.uniform(*RANDOM_DELAY_RANGE))

    files_api_url = f"https://hf-mirror.com/api/models/{model_id}/tree/main"

    for retry in range(RETRY_TIMES + 1):
        try:
            response = session.get(files_api_url, timeout=20)
            response.raise_for_status()
            files_data = response.json()

            for file_item in files_data:
                if (
                    file_item.get("type") == "file"
                    and file_item.get("path") == target_filename
                ):
                    return file_item.get("size", 0) / (1024 ** 3)

            print(f"未找到文件: {model_id} -> {target_filename}")
            return -1

        except requests.exceptions.HTTPError as e:
            status_code = getattr(e.response, "status_code", None)

            if status_code == 429 and retry < RETRY_TIMES:
                wait_time = RETRY_DELAY * (retry + 1)
                print(
                    f"触发 HF 限流: {model_id}，等待 {wait_time} 秒后重试 "
                    f"({retry + 1}/{RETRY_TIMES})"
                )
                interruptible_sleep(wait_time)
                continue

            print(f"提取文件大小失败: {model_id} -> {target_filename}: {e}")
            return -1

        except Exception as e:
            print(f"提取文件大小失败: {model_id} -> {target_filename}: {e}")
            return -1

    return -1


def filter_models_by_gguf_size(
    session: requests.Session,
    models: List[str],
    max_size_gb: float,
) -> List[str]:
    filtered: List[str] = []

    for index, model in enumerate(models, start=1):
        target_file = get_model_filename(model)
        file_size = extract_gguf_file_size(session, model, target_file)

        prefix = f"[{index}/{len(models)}]"

        if file_size != -1 and file_size <= max_size_gb:
            filtered.append(model)
            print(
                f"{prefix} 符合大小条件: {model} -> {target_file} "
                f"({file_size:.2f} GB)"
            )
        else:
            size_str = f"{file_size:.2f} GB" if file_size != -1 else "未知"
            print(
                f"{prefix} 不符合大小条件: {model} -> {target_file} "
                f"({size_str})"
            )

    return filtered


# ============================================================
# 7. Step 3：ModelHub 查重，过滤已入库模型
# ============================================================

def check_model_in_modelhub_db(model_id: str) -> Optional[bool]:
    """
    返回：
      True  = 已入库
      False = 未入库
      None  = 查询异常
    """
    url = BASE_URL + MODELHUB_DB_CHECK_ENDPOINT

    payload = {
        "current": 1,
        "pageSize": 20,
        "searchText": model_id,
    }

    try:
        response = requests.post(
            url,
            json=payload,
            headers=BASE_HEADERS,
            timeout=15,
        )

        if response.status_code != 200:
            print(f"查重 HTTP 异常: {model_id} -> {response.status_code}")
            return None

        resp_json = response.json()
        code = resp_json.get("code")
        records = resp_json.get("data", {}).get("records", [])

        if code != 0:
            print(
                f"查重接口业务异常: {model_id} -> "
                f"{resp_json.get('message', 'unknown')}"
            )
            return None

        if len(records) == 0:
            print(f"未入库: {model_id}")
            return False

        first_id = records[0].get("id")
        print(f"已入库: {model_id}，记录 ID: {first_id}")
        return True

    except requests.exceptions.Timeout:
        print(f"查重超时: {model_id}")
        return None
    except requests.exceptions.RequestException as e:
        print(f"查重网络错误: {model_id}: {e}")
        return None
    except Exception as e:
        print(f"查重未知错误: {model_id}: {e}")
        return None


def filter_models_not_in_db(models: List[str]) -> List[str]:
    not_in_db: List[str] = []

    for index, model in enumerate(models, start=1):
        print(f"[{index}/{len(models)}] 查重: {model}")

        result = check_model_in_modelhub_db(model)

        if result is False:
            not_in_db.append(model)
        elif result is True:
            pass
        else:
            # 查询异常时，默认不提交，避免重复入库。
            print(f"查重异常，跳过提交: {model}")

        interruptible_sleep(0.3)

    return not_in_db


# ============================================================
# 8. 登录
# ============================================================

def login() -> str:
    check_shutdown()

    if not USER_ACCOUNT or not USER_PASSWORD:
        raise RuntimeError(
            "没有设置账号密码，请在脚本中填写 USER_ACCOUNT / USER_PASSWORD。"
        )

    payload = {
        "userAccount": USER_ACCOUNT,
        "userPassword": USER_PASSWORD,
    }

    print("正在登录 ModelHub...")

    resp = requests.post(
        BASE_URL + LOGIN_ENDPOINT,
        headers=BASE_HEADERS,
        json=payload,
        timeout=20,
    )

    if resp.status_code != 200:
        raise RuntimeError(f"HTTP 登录失败: {resp.status_code} - {resp.text}")

    data = resp.json()
    if data.get("code") != 0:
        raise RuntimeError(f"业务登录失败: {data.get('message') or data}")

    token = data["data"]["token"]
    print("登录成功")
    return token


def get_token() -> str:
    return login()


# ============================================================
# 9. Step 4：提交验证任务
# ============================================================

def is_i_variant_model(model_filename: str) -> bool:
    return re.search(r"\.i\d+-Q4_0\.gguf$", model_filename) is not None


def build_config_params(model_id: str) -> Tuple[str, str]:
    model_filename = get_model_filename(model_id)

    if FORCE_TARGET_GPU:
        target_gpu = FORCE_TARGET_GPU
    elif is_i_variant_model(model_filename):
        target_gpu = "Ascend_910-b4"
    else:
        target_gpu = "Mthreads_s4000"

    if target_gpu == "Ascend_910-b4":
        config_params = f"""framework: llamacpp
api: chat
lang: zh
max_model_len: 4096
max_tokens: 1024
temperature: 0.7
repetition_penalty: 1.1
top_p: 0.9
sut_config:
  gpu_num: 1
  values:
    command: ['/workspace/llama.cpp/build_ascend/bin/llama-server', '--model', '/model/{model_filename}', '--alias', 'llm', '--threads', '16', '--n-gpu-layers', '128', '--prio', '3', '--min_p', '0.01', '--ctx-size', '4096', '--host', '0.0.0.0', '--port', '3316', '--jinja', '--flash-attn', 'off']
ref_config:
  gpu_num: 1
  values:
    command: ['/workspace/llama.cpp/build/bin/llama-server', '--model', '/model/{model_filename}', '--alias', 'llm', '--threads', '16', '--n-gpu-layers', '128', '--prio', '3', '--min_p', '0.01', '--ctx-size', '4096', '--host', '0.0.0.0', '--port', '80', '--jinja']
"""
    elif target_gpu == "Mthreads_s4000":
        config_params = f"""framework: llamacpp
api: completion
max_tokens: 1024
temperature: 0
repetition_penalty: 1.1
top_p: 0.9
max_model_len: 4096
sut_config:
  gpu_num: 1
  values:
    command: ['/app/llama-server', '--model', '/model/{model_filename}', '--alias', 'llm', '--threads', '20', '--n-gpu-layers', '999', '--prio', '3', '--min_p', '0.01', '--ctx-size', '2048', '--host', '0.0.0.0', '--port', '8000', '--jinja', '--flash-attn', 'off']
ref_config:
  gpu_num: 1
  values:
    command: ['/workspace/llama.cpp/build/bin/llama-server', '--model', '/model/{model_filename}', '--alias', 'llm', '--threads', '20', '--n-gpu-layers', '999', '--prio', '3', '--min_p', '0.01', '--ctx-size', '2048', '--host', '0.0.0.0', '--port', '8000', '--jinja', '--flash-attn', 'off']
"""
    else:
        raise ValueError(f"暂不支持的 targetGpu: {target_gpu}")

    return target_gpu, config_params


def is_business_success(data: Dict) -> bool:
    if not isinstance(data, dict):
        return False

    code = data.get("code")
    success = data.get("success")

    if code in (0, "0", 200, "200"):
        return True

    if success is True:
        return True

    return False


def submit_adapt_task(token: str, model_id: str) -> str:
    """
    返回：
      SUCCESS
      LIMIT
      FAILED
      AUTH_FAILED
    """
    check_shutdown()
    target_gpu, config_params = build_config_params(model_id)

    headers = {
        **BASE_HEADERS,
        "Authorization": f"Bearer {token}",
    }

    payload = {
        "modelAddress": model_id,
        "strategyId": STRATEGY_ID,
        "taskType": TASK_TYPE,
        "targetGpu": target_gpu,
        "framework": FRAMEWORK,
        "configParams": config_params,
    }

    print(f"\n提交模型: {model_id}")
    print(f"GPU: {target_gpu}")
    print(f"模型文件名: {get_model_filename(model_id)}")

    if DEBUG_PRINT_PAYLOAD:
        print(json.dumps(payload, ensure_ascii=False, indent=2))

    try:
        resp = requests.post(
            BASE_URL + ADD_ADAPT_TASK_ENDPOINT,
            headers=headers,
            json=payload,
            timeout=20,
        )

        print(f"HTTP 状态码: {resp.status_code}")

        if resp.status_code in (401, 403):
            print(f"鉴权失败: {resp.text}")
            return "AUTH_FAILED"

        try:
            data = resp.json()
        except Exception:
            data = None

        if resp.status_code != 200:
            print(f"提交失败，HTTP 错误: {resp.status_code} - {resp.text}")
            return "FAILED"

        if data is None:
            print(f"提交完成，但返回不是 JSON: {resp.text}")
            return "SUCCESS"

        if is_business_success(data):
            task_id = (
                data.get("data", {}).get("taskId")
                or data.get("data", {}).get("id")
                or data.get("taskId")
                or data.get("id")
            )

            if task_id:
                print(f"提交成功，Task ID: {task_id}")
            else:
                print(f"提交成功，返回: {json.dumps(data, ensure_ascii=False)}")

            return "SUCCESS"

        code = data.get("code")
        message = data.get("message", "")

        if str(code) == str(TASK_LIMIT_CODE):
            print(f"达到任务上限，稍后重试: {json.dumps(data, ensure_ascii=False)}")
            return "LIMIT"

        if "数量已达上限" in message or "当前等待中或运行中" in message:
            print(f"达到任务上限，稍后重试: {json.dumps(data, ensure_ascii=False)}")
            return "LIMIT"

        print(f"提交失败，业务返回: {json.dumps(data, ensure_ascii=False)}")
        return "FAILED"

    except requests.RequestException as e:
        print(f"请求异常: {e}")
        return "FAILED"


def submit_adapt_task_with_polling(token: str, model_id: str) -> Tuple[bool, str]:
    retry_times = 0
    current_token = token

    while True:
        result = submit_adapt_task(current_token, model_id)

        if result == "SUCCESS":
            return True, current_token

        if result == "AUTH_FAILED":
            print("尝试重新登录后重试当前模型...")
            current_token = login()
            interruptible_sleep(2)
            continue

        if result == "FAILED":
            return False, current_token

        if result == "LIMIT":
            retry_times += 1

            if MAX_LIMIT_RETRY_TIMES is not None and retry_times > MAX_LIMIT_RETRY_TIMES:
                print(f"达到最大等待次数，放弃当前模型: {model_id}")
                return False, current_token

            print(
                f"平台等待中/运行中的异步验证任务已满。"
                f"等待 {POLL_INTERVAL_SECONDS} 秒后重试当前模型。"
                f"模型: {model_id}，等待次数: {retry_times}"
            )

            interruptible_sleep(POLL_INTERVAL_SECONDS)
            continue

        print(f"未知提交状态: {result}")
        return False, current_token


# ============================================================
# 10. 主流程：筛选 -> 查重 -> 提交
# ============================================================

def main() -> None:
    ensure_output_dir()
    args = parse_args()
    check_shutdown()

    if args.reset_submitted:
        remove_file_if_exists(SUBMITTED_FILE)
    if args.reset_failed:
        remove_file_if_exists(FAILED_FILE)

    print("=" * 80)
    print("Auto HF -> ModelHub Pipeline")
    print("=" * 80)
    print(f"组织: {ORG_NAME}")
    print(f"页码范围: {START_PAGE} - {END_PAGE}")
    print(f"每页数量: {PER_PAGE}")
    print(f"最大 GGUF 文件大小: {MAX_FILE_SIZE_GB} GB")
    print(f"输出目录: {OUTPUT_DIR}")
    print(f"运行模式: {args.run_mode}")
    print(f"起始步骤: Step {args.start_step}")
    print(f"强制重跑步骤: {sorted(args.force_rerun_steps) if args.force_rerun_steps else '无'}")
    print(f"断点续跑跳过已提交模型: {RESUME}")
    print("=" * 80)

    hf_session: Optional[requests.Session] = None

    try:
        # Step 1: 拉取 HF 模型
        print("\nStep 1/4: 获取 HF / hf-mirror 模型列表")
        if should_run_step(1, ALL_MODELS_FILE, args):
            hf_session = hf_session or create_hf_session()
            all_models = get_org_models(hf_session, ORG_NAME)
            write_model_list(ALL_MODELS_FILE, all_models)
            print(f"共获取到 {len(all_models)} 个模型")
            print(f"已保存: {ALL_MODELS_FILE}")
        else:
            all_models = load_model_list(ALL_MODELS_FILE)
            print_step_loaded(1, ALL_MODELS_FILE, all_models)

        if not all_models:
            print("没有获取到模型，结束")
            return

        # Step 2: 按文件大小筛选
        print("\nStep 2/4: 按 GGUF 文件大小筛选")
        if should_run_step(2, SIZE_FILTERED_FILE, args):
            hf_session = hf_session or create_hf_session()
            size_filtered_models = filter_models_by_gguf_size(
                hf_session,
                all_models,
                MAX_FILE_SIZE_GB,
            )
            write_model_list(SIZE_FILTERED_FILE, size_filtered_models)
            print(f"大小筛选后剩余 {len(size_filtered_models)} 个模型")
            print(f"已保存: {SIZE_FILTERED_FILE}")
        else:
            size_filtered_models = load_model_list(SIZE_FILTERED_FILE)
            print_step_loaded(2, SIZE_FILTERED_FILE, size_filtered_models)

        if not size_filtered_models:
            print("没有符合大小条件的模型，结束")
            return

        # Step 3: ModelHub 查重
        print("\nStep 3/4: 查询 ModelHub 是否已入库")
        if should_run_step(3, NOT_IN_DB_FILE, args):
            not_in_db_models = filter_models_not_in_db(size_filtered_models)
            write_model_list(NOT_IN_DB_FILE, not_in_db_models)
            print(f"未入库模型共 {len(not_in_db_models)} 个")
            print(f"已保存: {NOT_IN_DB_FILE}")
        else:
            not_in_db_models = load_model_list(NOT_IN_DB_FILE)
            print_step_loaded(3, NOT_IN_DB_FILE, not_in_db_models)

        if not not_in_db_models:
            print("没有需要提交的未入库模型，结束")
            return

        # Step 4: 登录并提交
        print("\nStep 4/4: 登录并提交验证任务")
        token = get_token()

        already_submitted = load_model_set(SUBMITTED_FILE) if RESUME else set()
        already_failed = load_model_set(FAILED_FILE) if (RESUME and args.skip_failed) else set()
        skip_models = set(already_submitted) | set(already_failed)

        submit_candidates = [
            model for model in not_in_db_models
            if model not in skip_models
        ]

        if RESUME:
            print(f"断点续跑开启，已成功提交过 {len(already_submitted)} 个模型")
            if args.skip_failed:
                print(f"本次会额外跳过历史失败模型 {len(already_failed)} 个")
            print(f"本次待提交 {len(submit_candidates)} 个模型")

        success_models: List[str] = []
        failed_models: List[str] = []

        for index, model_id in enumerate(submit_candidates, start=1):
            print("\n" + "=" * 80)
            print(f"[{index}/{len(submit_candidates)}] 准备提交: {model_id}")

            ok, token = submit_adapt_task_with_polling(token, model_id)

            if ok:
                success_models.append(model_id)
                append_model(SUBMITTED_FILE, model_id)
            else:
                failed_models.append(model_id)
                append_model(FAILED_FILE, model_id)

            if index < len(submit_candidates):
                interruptible_sleep(SUBMIT_INTERVAL_SECONDS)

        summary = {
            "org_name": ORG_NAME,
            "start_page": START_PAGE,
            "end_page": END_PAGE,
            "max_file_size_gb": MAX_FILE_SIZE_GB,
            "run_mode": args.run_mode,
            "start_step": args.start_step,
            "force_rerun_steps": sorted(args.force_rerun_steps),
            "all_models_count": len(all_models),
            "size_filtered_count": len(size_filtered_models),
            "not_in_db_count": len(not_in_db_models),
            "already_submitted_count": len(already_submitted),
            "submitted_this_run_count": len(success_models),
            "failed_this_run_count": len(failed_models),
            "success_models": success_models,
            "failed_models": failed_models,
        }

        with open(SUMMARY_FILE, "w", encoding="utf-8") as f:
            json.dump(summary, f, ensure_ascii=False, indent=2)

        print("\n" + "=" * 80)
        print("流程完成")
        print(f"HF 获取模型数: {len(all_models)}")
        print(f"大小筛选后: {len(size_filtered_models)}")
        print(f"未入库模型数: {len(not_in_db_models)}")
        print(f"历史已提交跳过: {len(already_submitted)}")
        print(f"本次提交成功: {len(success_models)}")
        print(f"本次提交失败: {len(failed_models)}")
        print(f"汇总文件: {SUMMARY_FILE}")

        if failed_models:
            print("\n失败模型:")
            for model in failed_models:
                print(model)

    finally:
        if hf_session is not None:
            hf_session.close()


if __name__ == "__main__":
    main()