"""第 24 课第 3 步：明确输出格式、文件保存与错误返回值。"""
import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(ROOT))
from poetry_gpt.common import DEFAULT_RUN, write_json
from poetry_gpt.generate import Writer, parse_keywords


def main(argv=None):
    parser = argparse.ArgumentParser(description="我的可交付写诗命令")
    parser.add_argument("--start", default="")
    parser.add_argument("--keywords", default="")
    parser.add_argument("--form", choices=["five", "seven"], default="five")
    parser.add_argument("--device", choices=["auto", "cpu", "mps", "cuda"], default="auto")
    parser.add_argument("--checkpoint", type=Path, default=DEFAULT_RUN / "best.pt")
    parser.add_argument("--json", action="store_true")
    parser.add_argument("--output", type=Path)
    parser.add_argument("--count", type=int, default=1)
    args = parser.parse_args(argv)
    try:
        writer = Writer(args.checkpoint, device=args.device)
        poems = writer.write(
            start=args.start, keywords=parse_keywords(args.keywords),
            form=5 if args.form == "five" else 7, count=args.count,
        )
        if args.output:
            write_json(args.output, poems)
        if args.json:
            print(json.dumps(poems, ensure_ascii=False, indent=2))
        else:
            for poem in poems:
                print("\n".join(poem["lines"]))
    except (ValueError, FileNotFoundError) as error:
        print(f"提示：{error}", file=sys.stderr)
        return 2
    return 0


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