"""
Hendrycks MATH agentic RL 训练任务提交脚本。

整体架构：
  本脚本是 agentic RL 训练的入口，负责组装训练配置并提交任务。
  训练框架会自动协调以下流程：
    1. 从训练数据中采样 prompt
    2. Rollout 阶段：用当前策略模型 + 工具（python_interpreter）生成完整对话轨迹
    3. Reward 阶段：对轨迹打分（格式分 0.2 + 正确性分 0.8）
    4. 策略优化：基于 reward 信号更新模型参数

组件说明：
  Rollout : CalcXRolloutProcessor  — 基于 LangGraph + MCP 的 agent 推理，工具为 python_interpreter
  Reward  : MathRewardProcessor    — 格式奖励 (±1, 权重 0.2) + 正确性奖励 (0/1, 权重 0.8)
  Data    : hendrycks_math 数据集（训练集 + 测试集）

最终效果demo：ft-202605131703-f963
"""
import json
import subprocess
import sys
from pathlib import Path


SCRIPT_DIR = Path(__file__).resolve().parent
DATA_DIR = SCRIPT_DIR / "data"
SCRIPTS_DIR = SCRIPT_DIR / "scripts"


def data_file(filename: str) -> str:
    return str(DATA_DIR / filename)


def ensure_data_files(prep_steps: list[tuple[str, list[str]]]) -> None:
    """确保提交任务所需的数据文件存在；缺失时运行对应生成脚本。"""
    DATA_DIR.mkdir(parents=True, exist_ok=True)

    for script_name, target_names in prep_steps:
        missing = [name for name in target_names if not (DATA_DIR / name).is_file()]
        if not missing:
            logger.info(f"Data files already exist, skip {script_name}: {', '.join(target_names)}")
            continue

        script_path = SCRIPTS_DIR / script_name
        if not script_path.is_file():
            raise FileNotFoundError(f"Data preparation script not found: {script_path}")

        logger.info(f"Missing data files: {', '.join(missing)}. Running {script_path} ...")
        subprocess.run([sys.executable, str(script_path)], cwd=SCRIPT_DIR, check=True)

        remaining = [name for name in target_names if not (DATA_DIR / name).is_file()]
        if remaining:
            raise FileNotFoundError(
                f"Data preparation finished but target files are still missing: {', '.join(remaining)}"
            )


def count_jsonl_lines(file_path: str) -> int:
    """统计 JSONL 文件的行数（即样本数）。"""
    path = Path(file_path)
    if not path.exists():
        raise FileNotFoundError(f"Training file not found: {file_path}")
    with open(path) as f:
        return sum(1 for line in f if line.strip())


def calculate_save_steps(training_files: list[str], batch_size: int) -> int:
    """计算每 epoch 的步数，用于将 save_strategy='steps' 模拟 per-epoch 保存。

    steps_per_epoch = total_samples // batch_size
    """
    total = sum(count_jsonl_lines(f) for f in training_files)
    steps = total // batch_size
    if steps < 1:
        steps = 1
    return steps

from dashscope.finetune.reinforcement import (logger,
                                              DataSourceType, TrainingDataset, ValidationDataset,
                                              AgenticRLTuning, RolloutFunctionComponent, RewardFunctionComponent,
                                              FunctionComponentModel, FunctionComponentRuntime,
                                              RolloutInput, RewardInput, RewardOutput, RolloutOutput,
                                              AgenticRLFunctionComponent, FunctionType)
from dashscope.finetune.agentic_rl import AgenticRL
from dashscope.finetune.finetunes import FineTunes


async def main_workflow():
    """提交并监控一个 agentic RL 训练任务。

    流程：构建 Rollout/Reward 函数组件的运行时配置 → 调用 AgenticRL.run() 提交训练任务 → 查询任务状态。
    """
    try:
        logger.info("Starting main tests workflow")

        # ================================================================ #
        #  函数组件运行时配置（Rollout & Reward 共用类似配置）                   #
        #  ---------------------------------------------------------------- #
        #  cpu / memory_size / disk_size : 单实例资源规格                      #
        #  concurrency      : 每个实例的并发请求数上限                          #
        #  capacity          : 初始实例数                                      #
        #  min/max_capacity  : 弹性伸缩的实例数范围                             #
        #  memory_scale_threshold / concurrency_scale_threshold :             #
        #                     内存/并发利用率超过此阈值时触发扩容                 #
        #  env.ENABLE_TRAJECTORY : 开启后会记录完整对话轨迹，用于可观测性追踪     #
        # ================================================================ #
        rollout_runtime = {"cpu": 2, "memory_size": 4096, "disk_size": 10240, "concurrency": 30,
                           "capacity": 30, "min_capacity": 30, "max_capacity": 60,
                           "memory_scale_threshold": 0.6, "concurrency_scale_threshold": 0.6,
                           "env": {"ENABLE_TRAJECTORY": "true"}
        }

        reward_runtime = FunctionComponentRuntime(
            cpu=2,
            memory_size=4096,
            disk_size=10240,
            concurrency=30,
            env={"ENABLE_TRAJECTORY": "true"},
            capacity=30,
            min_capacity=30,
            max_capacity=60,
            memory_scale_threshold=0.6,
            concurrency_scale_threshold=0.6,
        )

        # ================================================================ #
        #  提交 Agentic RL 训练任务                                          #
        # ================================================================ #
        training_file = data_file("hendrycks_math_train_lv45.jsonl")
        validation_file = data_file("hendrycks_math_test_lv45.jsonl")
        ensure_data_files([
            ("download_hendrycks_math.py", [
                "hendrycks_math_train_lv45.jsonl",
                "hendrycks_math_test_lv45.jsonl",
            ]),
        ])
        batch_size = 128

        save_steps = calculate_save_steps([training_file], batch_size)
        logger.info(f"Calculated save_steps={save_steps} (1 epoch) from {training_file}, batch_size={batch_size}")

        client = AgenticRL()
        result = await client.run(
            job_name='your_math_job_name',
            # model="qwen3-4b-instruct-2507",
            model="qwen3.5-9b",
            # 训练集和验证集，JSONL 格式，每行包含 messages + ground_truth
            training_datasets=[TrainingDataset(
                data_source_type=DataSourceType.FILE_ID,
                file_name=training_file,
            )],
            validation_datasets=[ValidationDataset(
                data_source_type=DataSourceType.FILE_ID,
                file_name=validation_file,
            )],

            # ------------------------------------------------------------ #
            #  函数组件定义 — Rollout 和 Reward                               #
            #  ------------------------------------------------------------ #
            #  每个组件需要指定：                                              #
            #    name      : 组件名称（日志追踪用）                            #
            #    fcmodel   : 处理器类路径（classpath），框架会自动实例化         #
            #    runtime   : 运行时资源配置                                    #
            #    timeout   : 函数执行超时时间（秒）                             #
            #    weight    : 该 reward 组件的全局权重（可配置多个 reward 组件）   #
            # ------------------------------------------------------------ #
            functions=[
                RolloutFunctionComponent(
                    name="rollout-1",
                    timeout=600,
                    fcmodel=FunctionComponentModel(
                        zipdir=str(SCRIPT_DIR),
                        classpath="functions.rollout.math_rollout.CalcXRolloutProcessor"),
                    runtime=FunctionComponentRuntime(**rollout_runtime)),

                RewardFunctionComponent(
                    name="reward-1",
                    weight=1.0,   # 此 reward 组件的全局权重；若有多个 reward 组件可按权重加权
                    timeout=120,
                    fcmodel=FunctionComponentModel(
                        zipdir=str(SCRIPT_DIR),
                        classpath="functions.reward.math_reward.MathRewardProcessor"),
                    runtime=reward_runtime),
            ],

            # ------------------------------------------------------------ #
            #  训练超参数                                                     #
            #  ------------------------------------------------------------ #
            #  n_epochs          : 训练轮数                                   #
            #  learning_rate     : 策略优化学习率，RL 训练通常用较小的 lr       #
            #  max_prompt_length : prompt（含 system prompt）的最大 token 数   #
            #  max_length        : prompt + response 的最大总 token 数          #
            #  batch_size        : 每批采样的 prompt 数（越大方差越小但更费算力） #
            #  eval_steps        : 每隔多少步在验证集上评估一次                  #
            #  inference_gpu_memory_utilization : 推理引擎的 GPU 显存利用率上限 #
            #  save_strategy     : checkpoint 保存策略（"epoch" = 每轮保存）   #
            #  save_total_limit  : 最多保留的 checkpoint 数量                  #
            # ------------------------------------------------------------ #
            hyper_parameters={
                "n_epochs": 1,
                "learning_rate": 1e-6,
                "max_prompt_length": 8192,
                "max_length": 16384,
                "batch_size": batch_size,
                "eval_steps": 2,
                # 以下参数可按需取消注释以自定义分布式训练拓扑：
                "nnodes_training": 2,      # 训练节点数
                # "n_gpus_training": 8,      # 每节点训练 GPU 数
                "nnodes_rollout": 1,       # 推理节点数
                # "n_gpus_rollout": 8,       # 每节点推理 GPU 数
                # "inference_gpu_memory_utilization": 0.5,
                "save_strategy": "steps",  # 目前仅支持"steps"，"epoch" 将在后续支持
                "save_steps": save_steps//2,  # 预先计算好每 epoch 的步数，用于将 save_strategy='steps' 模拟 per-epoch 保存。
                "save_total_limit": 2,
                "lr_scheduler_type":"linear",
                # "lr_warmup_steps": 8
            },

            # ------------------------------------------------------------ #
            #  计算资源配置                                                   #
            #  ------------------------------------------------------------ #
            #  charge_type   : 计费模式（mtu_postpaid = 按量后付费）            #
            #  mtu_spec_code : GPU 规格（MTU4 对应特定算力卡型）               #
            #  mtu_capacity  : 申请的 MTU 算力单元数量                         #
            # ------------------------------------------------------------ #
            resources={
                "charge_type": "mtu_postpaid",
                "mtu_spec_code": "MTU4",
                "mtu_capacity": 24,
            },
        )

        # ================================================================ #
        #  任务提交结果处理                                                   #
        # ================================================================ #
        if result.status_code == 200:
            job_id = result.output.job_id
            logger.info(f"agentic rl submit: {job_id=}, {result=}")
        else:
            raise ValueError(f"agentic rl submit: {result}")

        # 查询任务状态（可用于轮询直到完成）
        result = AgenticRL.get(job_id=job_id)
        status = result.output.status
        logger.info(f"agentic rl get: {job_id=}, {status=}, {result=}")

        # 取消任务（如需要，取消注释以下代码）
        # result = AgenticRL.cancel(job_id=job_id)
        # logger.info(f"agentic rl cancel: {job_id=}, {result=}")

        logger.info("All tests workflows completed successfully")

    except Exception as e:
        logger.error(f"Main execution flow terminated: {e}")


if __name__ == "__main__":
    import asyncio

    asyncio.run(main_workflow())
