shanxia2.0/backend/scripts/import_vehicle_fuel_register.py

#!/usr/bin/env python3
"""导入客户「加油汇总表」到 pl_vehicle_fuel_register。

用法（在 backend 目录）:
  # 测试环境 dry-run
  python scripts/import_vehicle_fuel_register.py --env-file .env --dry-run

  # 测试环境正式导入
  python scripts/import_vehicle_fuel_register.py --env-file .env

  # 生产环境
  python scripts/import_vehicle_fuel_register.py --env-file .env.production
"""
from __future__ import annotations

import argparse
import os
import re
import sys
from dataclasses import dataclass, field
from datetime import date, datetime
from decimal import Decimal, InvalidOperation
from pathlib import Path
from typing import Optional

import openpyxl

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

IMPORT_TAG = "import:fuel-summary-20260916"
DEFAULT_DATA_DIR = ROOT.parent / "加油汇总表(202609161514)"
CATEGORY_MAP = {
    "百链": "bailian",
    "加油站": "gas_station",
    "山峡厂": "shanxia_factory",  # 历史 Excel 文案
    "山峡": "shanxia_factory",
}
PLATE_MAP = {
    "皖H3190": "皖HA3190",
    "叉车": "皖H21522",
    "吨桶": "皖H21522",
}
REGISTRANT_MAP = {
    "皖H13532": "王晓阳",
}
FOLDER = "pl/vehicle-fuel/register"


@dataclass
class ImportRow:
    source: str
    seq: str
    fuel_category: str
    register_date: date
    registrant_name: str
    license_plate: str
    odometer_km: Optional[Decimal]
    fuel_liters: Decimal
    remark: Optional[str] = None
    attachment_files: list[Path] = field(default_factory=list)

    @property
    def import_key(self) -> str:
        return f"[{IMPORT_TAG}|{self.source}:{self.seq}]"


def _load_env_file(path: Path) -> None:
    if not path.exists():
        raise SystemExit(f"环境文件不存在: {path}")
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        key = key.strip()
        value = value.strip().strip('"').strip("'")
        os.environ[key] = value


def _to_decimal(value, *, field_name: str) -> Decimal:
    if value is None or str(value).strip() == "":
        raise ValueError(f"{field_name} 为空")
    try:
        return Decimal(str(value).strip())
    except (InvalidOperation, ValueError) as exc:
        raise ValueError(f"{field_name} 无效: {value!r}") from exc


def _to_date(value) -> date:
    if isinstance(value, datetime):
        return value.date()
    if isinstance(value, date):
        return value
    text = str(value).strip()
    return datetime.strptime(text[:10], "%Y-%m-%d").date()


def _norm_seq(raw) -> str:
    text = str(raw).strip()
    if text.isdigit():
        return f"{int(text):02d}" if len(text) <= 2 else text
    return text


def _load_attachment_index(att_dir: Path) -> dict[str, list[Path]]:
    index: dict[str, list[Path]] = {}
    if not att_dir.is_dir():
        return index
    for path in sorted(att_dir.iterdir()):
        if not path.is_file() or path.name.startswith("."):
            continue
        m = re.match(r"^(\d+)_", path.name)
        if not m:
            print(f"  警告: 附件名无法解析序号: {path.name}")
            continue
        seq = f"{int(m.group(1)):02d}"
        index.setdefault(seq, []).append(path)
    return index


def _parse_excel(
    path: Path,
    *,
    source: str,
    with_attachments: bool,
    att_index: dict[str, list[Path]],
) -> list[ImportRow]:
    wb = openpyxl.load_workbook(path, read_only=True, data_only=True)
    rows = list(wb.active.iter_rows(values_only=True))
    wb.close()
    if not rows:
        return []
    headers = [str(h).strip() if h is not None else "" for h in rows[0]]
    out: list[ImportRow] = []
    for raw in rows[1:]:
        if not any(c is not None and str(c).strip() for c in raw):
            continue
        item = dict(zip(headers, raw))
        seq = _norm_seq(item.get("序号"))
        cat_label = str(item.get("加油类别") or "").strip()
        if cat_label not in CATEGORY_MAP:
            raise ValueError(f"{source} 序号{seq}: 未知加油类别 {cat_label!r}")
        name = str(item.get("登记人") or "").strip()
        name = REGISTRANT_MAP.get(name, name)
        plate = str(item.get("车号") or "").strip()
        plate = PLATE_MAP.get(plate, plate)
        odo_raw = item.get("期初公里数/Km")
        odometer = None
        if odo_raw is not None and str(odo_raw).strip() != "":
            odometer = _to_decimal(odo_raw, field_name="期初公里数")
        liters = _to_decimal(item.get("加油升数/L"), field_name="加油升数")
        files: list[Path] = []
        if with_attachments:
            files = list(att_index.get(seq, []))
            if not files:
                raise ValueError(f"{source} 序号{seq}: 声明有附件但未找到文件")
        out.append(
            ImportRow(
                source=source,
                seq=seq,
                fuel_category=CATEGORY_MAP[cat_label],
                register_date=_to_date(item.get("日期")),
                registrant_name=name,
                license_plate=plate,
                odometer_km=odometer,
                fuel_liters=liters,
                attachment_files=files,
            )
        )
    return out


def load_rows(data_dir: Path) -> list[ImportRow]:
    att_index = _load_attachment_index(data_dir / "附件")
    with_att = _parse_excel(
        data_dir / "(总表)表格视图-有附件.xlsx",
        source="with",
        with_attachments=True,
        att_index=att_index,
    )
    no_att = _parse_excel(
        data_dir / "(总表)表格视图——没有附件.xlsx",
        source="no",
        with_attachments=False,
        att_index=att_index,
    )
    return with_att + no_att


def main() -> None:
    parser = argparse.ArgumentParser(description="导入加油汇总表")
    parser.add_argument("--env-file", default=".env", help="环境文件路径（相对 backend）")
    parser.add_argument(
        "--data-dir",
        default=str(DEFAULT_DATA_DIR),
        help="客户数据目录",
    )
    parser.add_argument("--dry-run", action="store_true", help="只校验不写库")
    parser.add_argument("--tenant-no", default="shanxia2.0")
    args = parser.parse_args()

    env_path = Path(args.env_file)
    if not env_path.is_absolute():
        env_path = ROOT / env_path
    _load_env_file(env_path)

    # 延迟导入，确保 env 已注入
    from sqlalchemy.orm import Session

    from app.core.database import SessionLocal
    from app.core.storage import storage
    from app.core.utils import generate_uuid
    from app.models.models import SysUser
    from app.models.pl_models import PlVehicle, PlVehicleFuelRegister
    from app.schemas.pl.vehicle_fuel import calc_fuel_tons

    data_dir = Path(args.data_dir)
    rows = load_rows(data_dir)
    print(f"环境文件: {env_path}")
    print(f"数据目录: {data_dir}")
    print(f"解析行数: {len(rows)} (附件文件合计 {sum(len(r.attachment_files) for r in rows)})")
    print(f"模式: {'DRY-RUN' if args.dry_run else 'WRITE'}")

    db: Session = SessionLocal()
    try:
        users = (
            db.query(SysUser)
            .filter(SysUser.del_flag == 0, SysUser.status == "1")
            .all()
        )
        by_nick = {}
        for u in users:
            nick = (u.nick_name or "").strip()
            if nick:
                by_nick.setdefault(nick, []).append(u)

        vehicles = {
            (v.license_plate or "").strip(): v
            for v in db.query(PlVehicle)
            .filter(PlVehicle.tenant_no == args.tenant_no, PlVehicle.del_flag == 0)
            .all()
        }

        existing_remarks = [
            (r.remark or "")
            for r in db.query(PlVehicleFuelRegister)
            .filter(
                PlVehicleFuelRegister.tenant_no == args.tenant_no,
                PlVehicleFuelRegister.del_flag == 0,
                PlVehicleFuelRegister.remark.like(f"%{IMPORT_TAG}%"),
            )
            .all()
        ]

        errors: list[str] = []
        prepared: list[tuple[ImportRow, SysUser, Optional[str]]] = []
        for row in rows:
            hits = by_nick.get(row.registrant_name, [])
            if len(hits) != 1:
                errors.append(
                    f"{row.source}:{row.seq} 登记人无法唯一匹配: {row.registrant_name!r} hits={len(hits)}"
                )
                continue
            if row.license_plate not in vehicles:
                errors.append(
                    f"{row.source}:{row.seq} 车牌不存在: {row.license_plate}"
                )
                continue
            key = row.import_key
            already = any(key in rm for rm in existing_remarks)
            if already:
                print(f"  skip 已导入 {key}")
                continue
            prepared.append((row, hits[0], None))

        if errors:
            print("校验失败:")
            for e in errors:
                print(" ", e)
            raise SystemExit(1)

        print(f"待写入: {len(prepared)} 行")
        if args.dry_run:
            for row, user, _ in prepared[:5]:
                print(
                    f"  sample {row.source}:{row.seq} {row.register_date} "
                    f"{row.license_plate} {row.fuel_category} {row.fuel_liters}L "
                    f"-> {calc_fuel_tons(row.fuel_liters)}t "
                    f"registrant={user.nick_name} atts={len(row.attachment_files)}"
                )
            if len(prepared) > 5:
                print(f"  ... 另有 {len(prepared) - 5} 行")
            print("DRY-RUN 完成，未写库")
            return

        def next_register_no(on_date: date) -> str:
            prefix = f"SXJY{on_date:%Y%m%d}"
            last = (
                db.query(PlVehicleFuelRegister)
                .filter(PlVehicleFuelRegister.register_no.like(f"{prefix}%"))
                .order_by(PlVehicleFuelRegister.register_no.desc())
                .first()
            )
            seq = int(last.register_no[-6:]) + 1 if last else 1
            return f"{prefix}{seq:06d}"

        created = 0
        for row, user, _ in prepared:
            urls: list[str] = []
            for fp in row.attachment_files:
                content = fp.read_bytes()
                full_url = storage.upload_file(content, fp.name, FOLDER)
                # 与管理端一致：有 BASE_URL 时存相对路径
                from app.core.config import settings

                base_url = settings.MINIO_BASE_URL.rstrip("/") if settings.MINIO_BASE_URL else ""
                if base_url and full_url.startswith(base_url + "/"):
                    urls.append(full_url[len(base_url) + 1 :])
                else:
                    # 存 object key（去掉 endpoint/bucket 前缀）更稳
                    from app.core.storage import storage as st

                    if hasattr(st, "_extract_object_key"):
                        key = st._extract_object_key(full_url)
                        urls.append(key or full_url)
                    else:
                        urls.append(full_url)

            now = datetime.now()
            remark = row.import_key
            entity = PlVehicleFuelRegister(
                id=generate_uuid(),
                register_no=next_register_no(row.register_date),
                fuel_category=row.fuel_category,
                register_date=row.register_date,
                registrant_id=user.id,
                license_plate=row.license_plate,
                odometer_km=row.odometer_km,
                fuel_liters=row.fuel_liters,
                fuel_tons=calc_fuel_tons(row.fuel_liters),
                attachment_urls=",".join(urls) if urls else None,
                status="normal",
                tenant_no=args.tenant_no,
                remark=remark,
                create_by=user.id,
                create_dept_id=user.dept_id or "",
                create_time=now,
                update_by=user.id,
                update_time=now,
                del_flag=0,
            )
            db.add(entity)
            db.commit()
            created += 1
            if created % 10 == 0 or created == len(prepared):
                print(f"  已写入 {created}/{len(prepared)}")

        print(f"完成: 新建 {created} 条")
        total = (
            db.query(PlVehicleFuelRegister)
            .filter(
                PlVehicleFuelRegister.tenant_no == args.tenant_no,
                PlVehicleFuelRegister.del_flag == 0,
                PlVehicleFuelRegister.remark.like(f"%{IMPORT_TAG}%"),
            )
            .count()
        )
        print(f"库中本批次标记记录数: {total}")
    finally:
        db.close()


if __name__ == "__main__":
    main()