dibang-xiaoshou/auto-create-plan/auto_create_monthly_plan.py

from __future__ import annotations

"""
按月自动生成采购计划单。

运行方式（示例）：
    cd auto-create-plan
    python auto_create_monthly_plan.py

建议在服务器上用守护进程（如 systemd、supervisor）长期运行本脚本，
不依赖 crontab，由脚本内部每小时自行判断是否需要执行当月生成逻辑。
"""

import configparser
import logging
import sys
import time
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Optional

from sqlalchemy import and_, or_, select
from sqlalchemy.orm import Session
from sqlmodel import Field, SQLModel, create_engine

ROOT_DIR = Path(__file__).resolve().parents[1]

class PurchasePlan(SQLModel, table=True):
    """计划表 - 对应 purchase_plans。"""

    __tablename__ = "purchase_plans"

    id: int | None = Field(default=None, primary_key=True)
    plan_no: str | None = Field(default=None, max_length=255, index=True)
    plan_date: str = Field(max_length=255, index=True)
    start_date: str | None = Field(default=None, max_length=255, index=True)
    end_date: str | None = Field(default=None, max_length=255, index=True)
    business_type: str = Field(max_length=255, index=True)
    product: str = Field(max_length=255, index=True)
    specification: str | None = Field(default=None, max_length=255)
    planned_quantity: float = Field()
    vehicle_no: str | None = Field(default="", max_length=255)
    customer_name: str | None = Field(default=None, max_length=255)
    remarks: str | None = Field(default=None)
    is_deleted: int = Field(default=0, index=True)
    execution_status: str | None = Field(default=None, max_length=50, index=True)
    executed_at: str | None = Field(default=None, max_length=255)
    last_weighing_slip_id: int | None = Field(default=None)
    executed_quantity: float | None = Field(default=0)
    created_at: str = Field(max_length=255)
    auto_generate_enabled: int = Field(default=0, index=True)
    auto_generate_end_month: str | None = Field(default=None, max_length=7, index=True)
    auto_generate_last_month: str | None = Field(default=None, max_length=7)


@dataclass
class JobConfig:
    db_url: str
    loop_interval_seconds: int
    tz: timezone
    log_file: Optional[Path]


def load_config(conf_path: Path) -> JobConfig:
    cfg = configparser.ConfigParser()
    if not conf_path.exists():
        raise FileNotFoundError(f"配置文件不存在: {conf_path}")

    cfg.read(conf_path, encoding="utf-8")

    db_url = cfg.get("database", "url", fallback="").strip()
    if not db_url:
        raise ValueError("配置 [database] url 不能为空")

    loop_interval_seconds = cfg.getint("job", "loop_interval_seconds", fallback=3600)

    tz_name = cfg.get("job", "timezone", fallback="Asia/Shanghai").strip()
    # 简单处理：目前只考虑 Asia/Shanghai，其他时区可以按需扩展
    if tz_name == "Asia/Shanghai":
        tz = timezone(timedelta(hours=8))
    else:
        # 回退为 UTC
        tz = timezone.utc

    log_file_raw = cfg.get("job", "log_file", fallback="").strip()
    log_file = ROOT_DIR / "auto-create-plan" / log_file_raw if log_file_raw else None

    return JobConfig(
        db_url=db_url,
        loop_interval_seconds=loop_interval_seconds,
        tz=tz,
        log_file=log_file,
    )


def setup_logging(log_file: Optional[Path]) -> None:
    handlers: list[logging.Handler] = [logging.StreamHandler(sys.stdout)]
    if log_file:
        log_file.parent.mkdir(parents=True, exist_ok=True)
        handlers.append(logging.FileHandler(log_file, encoding="utf-8"))

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
        handlers=handlers,
    )


def get_now(cfg: JobConfig) -> datetime:
    return datetime.now(tz=cfg.tz)


def month_range(dt: datetime) -> tuple[str, str]:
    """返回 dt 所在月份的首日、末日（字符串 YYYY-MM-DD）。"""
    first = dt.replace(day=1)
    # 下月一号减一天
    if first.month == 12:
        next_month = first.replace(year=first.year + 1, month=1, day=1)
    else:
        next_month = first.replace(month=first.month + 1, day=1)
    last = next_month - timedelta(days=1)
    return first.strftime("%Y-%m-%d"), last.strftime("%Y-%m-%d")


def get_last_month_range(now: datetime) -> tuple[str, str]:
    """返回上个月的首日、末日（字符串 YYYY-MM-DD）。"""
    first_this, _ = month_range(now)
    first_this_dt = datetime.strptime(first_this, "%Y-%m-%d").replace(tzinfo=now.tzinfo)
    last_prev = first_this_dt - timedelta(days=1)
    return month_range(last_prev)


def generate_plan_no() -> str:
    """生成计划号：格式与后端 _generate_plan_no 一致：SX-YYYYMMDDHHMMSS001。"""
    now = datetime.now()
    return f"SX-{now.strftime('%Y%m%d%H%M%S')}001"


def default_execution_status(business_type: str) -> str:
    return "待入库" if business_type == "入库" else "待出库"


def run_job_once(session: Session, cfg: JobConfig) -> None:
    now = get_now(cfg)

    # 仅在每月 1 号执行
    if now.day != 1:
        logging.info("当前不是每月 1 日，跳过本次执行")
        return

    current_month = now.strftime("%Y-%m")
    this_month_first, this_month_last = month_range(now)
    last_month_first, last_month_last = get_last_month_range(now)

    logging.info(
        "开始执行自动生成计划任务，current_month=%s, last_month=[%s, %s], this_month=[%s, %s]",
        current_month,
        last_month_first,
        last_month_last,
        this_month_first,
        this_month_last,
    )

    # 选出“上个月的计划”，且开启了自动生成，且本月尚未生成过，且未超过结束月份
    stmt = (
        select(PurchasePlan)
        .where(
            and_(
                PurchasePlan.is_deleted == 0,
                PurchasePlan.auto_generate_enabled == 1,
                or_(
                    PurchasePlan.auto_generate_last_month.is_(None),
                    PurchasePlan.auto_generate_last_month != current_month,
                ),
                or_(
                    PurchasePlan.auto_generate_end_month.is_(None),
                    PurchasePlan.auto_generate_end_month >= current_month,
                ),
                PurchasePlan.start_date == last_month_first,
                PurchasePlan.end_date == last_month_last,
            )
        )
    )

    plans = list(session.execute(stmt).scalars())
    if not plans:
        logging.info("未找到需要自动生成下月计划的记录")
        return

    logging.info("找到 %d 条可自动生成的上月计划", len(plans))

    created_count = 0
    for p in plans:
        new_plan = PurchasePlan(
            plan_no=generate_plan_no(),
            plan_date=this_month_first,
            start_date=this_month_first,
            end_date=this_month_last,
            business_type=p.business_type,
            product=p.product,
            specification=p.specification,
            planned_quantity=p.planned_quantity,
            vehicle_no=p.vehicle_no or "",
            customer_name=p.customer_name,
            remarks=p.remarks,
            is_deleted=0,
            execution_status=default_execution_status(p.business_type),
            executed_at=None,
            last_weighing_slip_id=None,
            executed_quantity=0,
            created_at=now.strftime("%Y-%m-%d %H:%M:%S"),
            auto_generate_enabled=p.auto_generate_enabled,
            auto_generate_end_month=p.auto_generate_end_month,
            auto_generate_last_month=None,
        )
        session.add(new_plan)

        # 标记老计划这个月已经处理过
        p.auto_generate_last_month = current_month
        created_count += 1

    session.commit()
    logging.info("本次自动生成完成，共生成 %d 条新计划", created_count)


def main() -> None:
    conf_path = ROOT_DIR / "auto-create-plan" / "setting.conf"
    cfg = load_config(conf_path)
    setup_logging(cfg.log_file)

    logging.info("启动自动生成月度计划任务，使用数据库：%s", cfg.db_url)

    engine = create_engine(cfg.db_url, echo=False, future=True)

    # 确保元数据已加载（只针对 PurchasePlan 即可）
    SQLModel.metadata.create_all(engine, tables=[PurchasePlan.__table__])

    while True:
        try:
            with Session(engine) as session:
                run_job_once(session, cfg)
        except Exception:
            logging.exception("执行自动生成计划任务时发生异常")
        time.sleep(cfg.loop_interval_seconds)


if __name__ == "__main__":
    main()