from __future__ import annotations

import argparse
import hashlib
import json
import os
import time
from pathlib import Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Demo background task with progress and checkpointing.")
    parser.add_argument("--steps", type=int, default=20, help="Total number of progress steps.")
    parser.add_argument("--sleep-s", type=float, default=3.0, help="Sleep duration between steps.")
    parser.add_argument(
        "--work-units",
        type=int,
        default=25000,
        help="How many hash rounds to run inside each step.",
    )
    parser.add_argument(
        "--checkpoint",
        default="run_state/demo-progress.json",
        help="Path to the JSON checkpoint file.",
    )
    parser.add_argument(
        "--resume",
        action="store_true",
        help="Resume from the checkpoint file if it exists.",
    )
    return parser.parse_args()


def load_checkpoint(path: Path) -> dict:
    if not path.exists():
        return {}
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError:
        return {}


def save_checkpoint(path: Path, payload: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temp_path = path.with_suffix(path.suffix + ".tmp")
    temp_path.write_text(json.dumps(payload, ensure_ascii=True, indent=2), encoding="utf-8")
    os.replace(temp_path, path)


def run_step(step: int, work_units: int) -> str:
    digest = f"seed-{step}".encode("utf-8")
    for unit in range(work_units):
        digest = hashlib.sha256(digest + f":{unit}".encode("utf-8")).digest()
    return digest.hex()


def main() -> int:
    args = parse_args()
    checkpoint_path = Path(args.checkpoint).expanduser()

    state = load_checkpoint(checkpoint_path) if args.resume else {}
    completed_steps = int(state.get("completed_steps", 0))
    if completed_steps < 0:
        completed_steps = 0
    if completed_steps > args.steps:
        completed_steps = args.steps

    print(
        f"starting demo task: steps={args.steps}, sleep_s={args.sleep_s}, "
        f"work_units={args.work_units}, resume={args.resume}, completed={completed_steps}",
        flush=True,
    )

    last_digest = str(state.get("last_digest", ""))
    started_at = state.get("started_at") or time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())

    for step in range(completed_steps, args.steps):
        current_step = step + 1
        last_digest = run_step(current_step, args.work_units)
        progress_pct = round(current_step / args.steps * 100.0, 1) if args.steps else 100.0
        payload = {
            "started_at": started_at,
            "updated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            "completed_steps": current_step,
            "total_steps": args.steps,
            "progress_pct": progress_pct,
            "last_digest": last_digest,
            "status": "running" if current_step < args.steps else "completed",
        }
        save_checkpoint(checkpoint_path, payload)
        print(
            f"progress step={current_step}/{args.steps} "
            f"pct={progress_pct:.1f}% digest={last_digest[:16]}",
            flush=True,
        )
        if current_step < args.steps:
            time.sleep(args.sleep_s)

    print("demo task completed successfully", flush=True)
    return 0


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