#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import PurePosixPath
from typing import Any


DEFAULT_DATA_REPO = "git@git.n.xiaomi.com:ai-service/ai-planning.git"
DEFAULT_BRANCH = "autoresearch-v1"
DEFAULT_RECIPE = "zk_sft"


def main() -> int:
    parser = argparse.ArgumentParser(description="Render a lightweight SFT submission plan.")
    parser.add_argument("--input", "-i", help="JSON input file. Reads stdin when omitted.")
    args = parser.parse_args()
    text = open(args.input, encoding="utf-8").read() if args.input else sys.stdin.read()
    payload = json.loads(text or "{}")
    result = render_plan(payload)
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0 if result["ok"] else 1


def render_plan(payload: dict[str, Any]) -> dict[str, Any]:
    missing: list[str] = []
    jupyter_url = clean(payload.get("jupyter_url"))
    data_commit = clean(payload.get("data_commit") or payload.get("commit"))
    data_path = clean(payload.get("data_path"))
    if not jupyter_url:
        missing.append("jupyter_url")
    if not data_commit and not data_path:
        missing.append("data_commit or data_path")

    owner = clean(payload.get("owner")) or "wuyang6"
    run_name = clean(payload.get("run_name")) or default_run_name(data_commit or data_path or "manual")
    workspace = clean(payload.get("workspace")) or str(
        PurePosixPath("/mnt/wangsenhao/autoresearch-zk-users") / owner / run_name
    )
    data_repo = clean(payload.get("data_repo")) or DEFAULT_DATA_REPO
    data_branch = clean(payload.get("data_branch")) or DEFAULT_BRANCH
    recipe = clean(payload.get("recipe")) or DEFAULT_RECIPE

    commands = {
        "prepare_workspace": [
            f"mkdir -p {workspace}/{{scripts,assets,results,output}}",
            f"cd {workspace}",
        ],
        "sync_data": [
            f"[ -d ai-planning/.git ] || git clone -b {data_branch} {data_repo} ai-planning",
            f"git -C ai-planning fetch origin {data_branch}",
            f"git -C ai-planning checkout {data_branch}",
            f"git -C ai-planning reset --hard {data_commit}" if data_commit else f"# use existing data_path: {data_path}",
        ],
        "submit_sft": [
            "source ~/.cloudml-cli/.profile",
            f"export AUTORESEARCH_CHAT_ROOT={workspace}",
            'export AUTORESEARCH_ROOT="$AUTORESEARCH_CHAT_ROOT"',
            f"cd {workspace}",
            'eval "$(./scripts/resolve_run_ids.sh)"',
            './scripts/submit_sft.sh "$SFT_RUNDIC"',
        ],
        "check_status": [
            "cml custom_train describe <JOB_ID>",
            f"test -f {workspace}/sft_output/_SUCCESS",
        ],
    }
    return {
        "ok": not missing,
        "missing": missing,
        "plan": {
            "jupyter_url": jupyter_url,
            "recipe": recipe,
            "workspace": workspace,
            "data_repo": data_repo,
            "data_branch": data_branch,
            "data_commit": data_commit,
            "data_path": data_path,
            "success_marker": f"{workspace}/sft_output/_SUCCESS",
        },
        "commands": commands,
        "required_outputs": [
            "CloudML JobID",
            "CloudML task URL",
            "SFT_RUNDIC",
            "EVAL_RUNDIC",
            "workspace",
            "success_marker",
        ],
    }


def clean(value: Any) -> str:
    return str(value or "").strip()


def default_run_name(seed: str) -> str:
    slug = re.sub(r"[^a-zA-Z0-9_-]+", "_", seed).strip("_")[:32] or "manual"
    return f"manual_sft_{slug}"


if __name__ == "__main__":
    raise SystemExit(main())