"""Code generation agentic RL workflow.

Rollout : CodeRolloutProcessor  (multi-turn LangGraph + execute_code tool)
Reward  : CodeRewardProcessor   (format 0.1 + accuracy/pass@1 0.9)
Data    : ablation_taco (4300 examples) / livecodebench_v6_test (182 examples)
"""
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():
    """Main execution workflow"""
    try:
        logger.info("Starting code generation RL workflow")

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

        training_file = data_file("ablation_taco.jsonl")
        validation_file = data_file("livecodebench_v6_test.jsonl")
        ensure_data_files([
            ("sample_ablation_subsets.py", ["ablation_taco.jsonl"]),
            ("download_livecodebench.py", ["livecodebench_v6_test.jsonl"]),
        ])
        batch_size = 256

        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_code_job_name',
            # model="qwen3-4b-instruct-2507",
            model="qwen3.5-9b",
            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,
            )],
            functions=[
                RolloutFunctionComponent(
                    name="rollout-1",
                    # timeout=600,
                    fcmodel=FunctionComponentModel(
                        zipdir=str(SCRIPT_DIR),
                        classpath="functions.rollout.code_rollout.CodeRolloutProcessor"),
                    runtime=FunctionComponentRuntime(**rollout_runtime)),

                RewardFunctionComponent(
                    name="reward-1",
                    weight=1.0,
                    # timeout=120,
                    fcmodel=FunctionComponentModel(
                        zipdir=str(SCRIPT_DIR),
                        classpath="functions.reward.code_reward.CodeRewardProcessor"),
                    runtime=reward_runtime),
            ],
            hyper_parameters={
                "n_epochs": 1,
                "learning_rate": 7e-7,
                "kl_loss_coef": 0.01,
                "max_prompt_length": 12288,
                "max_length": 20480,
                "sequence_parallel_size": 2,
                "ppo_mini_batch_size": 64,
                "batch_size": batch_size,
                "eval_steps": 2,
                "nnodes_training": 2,
                "nnodes_rollout": 1,
                "save_strategy": "steps",
                "save_steps": save_steps,
                "save_total_limit": 3,
                "lr_scheduler_type": "constant",
                "lr_warmup_steps": 3,
                "inference_gpu_memory_utilization": 0.75,
            },
            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=}")

        logger.info("All code generation workflows completed successfully")

    except Exception as e:
        logger.error(f"Code generation workflow terminated: {e}")


if __name__ == "__main__":
    import asyncio

    asyncio.run(main_workflow())
