# -*- coding: utf-8 -*-
"""抓取中国食物成分表(nlc.chinanutri.cn/fq)全部食物, 生成 MySQL INSERT 脚本.

    python tools/fetch_food_nutrition.py            # 联网抓取(带缓存,可续跑)
    python tools/fetch_food_nutrition.py --offline  # 只用缓存重新生成 SQL

产物: tools/_cache/  接口原始响应缓存
      src/main/resources/db/data_food_nutrition.sql
前置: 先执行 db/migration_add_food_nutrition.sql 建表.
"""
from __future__ import annotations

import argparse
import json
import pathlib
import re
import ssl
import time
import urllib.parse
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed

BASE = "https://nlc.chinanutri.cn/fq/"
QUERY_URL = BASE + "FoodInfoQueryAction!queryFoodInfoList.do"
ROOT = pathlib.Path(__file__).resolve().parent.parent
CACHE = ROOT / "tools" / "_cache"
OUT_SQL = ROOT / "src" / "main" / "resources" / "db" / "data_food_nutrition.sql"
CTX = ssl.create_default_context()
CTX.check_hostname = False
CTX.verify_mode = ssl.CERT_NONE
WORKERS, RETRY, BATCH = 4, 4, 200

# 数组下标 -> 列名 (下标 = 列表页表头序号 + 2)
NUTRIENTS = [
    (5, "edible_percent"), (6, "water"), (7, "energy_kj"), (8, "protein"),
    (9, "fat"), (10, "cholesterol"), (11, "ash"), (12, "carbohydrate"),
    (13, "dietary_fiber"), (14, "carotene"), (15, "retinol"), (16, "vitamin_e"),
    (17, "thiamin"), (18, "riboflavin"), (19, "niacin"), (20, "vitamin_c"),
    (21, "calcium"), (22, "phosphorus"), (23, "potassium"), (24, "sodium"),
    (25, "magnesium"), (26, "iron"), (27, "zinc"), (28, "selenium"),
    (29, "copper"), (30, "manganese"), (31, "iodine"), (32, "sfa_percent"),
    (33, "mufa_percent"), (34, "pufa_percent"), (35, "fa_total_percent"),
]
COLUMNS = ["food_id", "name", "alias", "english_name", "image_path",
           "category_id", "category_name", "sub_category_id", "sub_category_name"] \
    + [c for _, c in NUTRIENTS] + ["raw_values"]

NUM_RE = re.compile(r"^([+-]?\d+(?:\.\d+)?)\s*(?:g|mg|μg|µg|ug|kJ|kj|%)?$")
NULL_TOKENS = {"", "—", "–", "-", "…", "tr", "微量", "未检出", "未测"}


def to_num(raw) -> str | None:
    s = "" if raw is None else str(raw).strip()
    if s.lower() in NULL_TOKENS:
        return None
    m = NUM_RE.match(s)
    return m.group(1) if m else None


def sql_str(v) -> str:
    if v is None or str(v) == "":
        return "NULL"
    return "'" + str(v).replace("\\", "\\\\").replace("'", "''") + "'"


def post_form(pairs: dict) -> str:
    req = urllib.request.Request(
        QUERY_URL, data=urllib.parse.urlencode(pairs).encode(),
        headers={"Content-Type": "application/x-www-form-urlencoded; charset=UTF-8",
                 "X-Requested-With": "XMLHttpRequest",
                 "Referer": BASE + "foodlist_0_0_0_0_0_1.htm",
                 "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"})
    with urllib.request.urlopen(req, timeout=30, context=CTX) as r:
        return r.read().decode("utf-8", errors="replace")


def page_json(c1: int, c2: int, page: int, offline: bool) -> dict:
    path = CACHE / f"cat_{c1}_{c2}_p{page}.json"
    if path.exists():
        return json.loads(path.read_text(encoding="utf-8"))
    if offline:
        raise RuntimeError(f"离线模式缺缓存 {path.name}")
    err = None
    for i in range(1, RETRY + 1):
        try:
            data = json.loads(post_form({"categoryOne": str(c1), "categoryTwo": str(c2),
                                         "foodName": "", "pageNum": str(page),
                                         "field": "", "flag": ""}))
            path.write_text(json.dumps(data, ensure_ascii=False), encoding="utf-8")
            return data
        except Exception as e:  # noqa: BLE001
            err = e
            time.sleep(1.2 * i)
    raise RuntimeError(f"抓取 {c1}/{c2} p{page} 失败: {err}")


def fetch_group(c1: int, c2: int, offline: bool, label: str = "") -> list:
    first = page_json(c1, c2, 1, offline)
    total = int(first.get("totalPages") or 1)
    pages = [first] + [None] * (total - 1)
    if total > 1:
        with ThreadPoolExecutor(max_workers=WORKERS) as ex:
            futs = {ex.submit(page_json, c1, c2, p, offline): p for p in range(2, total + 1)}
            for n, fut in enumerate(as_completed(futs), 1):
                pages[futs[fut] - 1] = fut.result()
                if n % 20 == 0 or n == len(futs):
                    print(f"      {n}/{len(futs)} 页", flush=True)
    rows = [r for pg in pages for r in (pg.get("list") or [])]
    print(f"   {label or f'{c1}/{c2}'}: {total} 页 -> {len(rows)} 条")
    return rows


def load_categories(offline: bool) -> list:
    f = CACHE / "index.htm"
    if not f.exists():
        if offline:
            raise RuntimeError("离线模式缺 index.htm")
        req = urllib.request.Request(BASE, headers={"User-Agent": "Mozilla/5.0"})
        with urllib.request.urlopen(req, timeout=30, context=CTX) as r:
            f.write_bytes(r.read())
    html = f.read_bytes().decode("utf-8", errors="replace")
    out = []
    for m in re.finditer(r'<a href="foodlist_0_\d+_\d+_0_0_1\.htm"([^>]*)>([^<]*)</a>', html):
        attrs, name = m.group(1), m.group(2).strip()
        sid = re.search(r'data_id="(\d+)"', attrs)
        pid = re.search(r'data_pid="(\d+)"', attrs)
        if sid and pid:
            out.append((int(pid.group(1)), int(sid.group(1)), name))
    return out


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--offline", action="store_true")
    args = ap.parse_args()
    CACHE.mkdir(parents=True, exist_ok=True)

    cats = load_categories(args.offline)
    tops = {i: n for p, i, n in cats if p == 0}
    subs = [(p, i, n) for p, i, n in cats if p != 0]
    print(f"分类: {len(tops)} 个大类, {len(subs)} 个子类")

    print("[1/3] 抓取全量食物 ...")
    master = {}
    for row in fetch_group(0, 0, args.offline, "全量"):
        master[int(row[0])] = row
    print(f"   全量去重后 {len(master)} 条")

    print("[2/3] 抓取分类归属 ...")
    cat_of: dict[int, tuple] = {}
    for pid, sid, name in subs:
        for row in fetch_group(pid, sid, args.offline, f"{tops.get(pid,'')} > {name}"):
            fid = int(row[0])
            master.setdefault(fid, row)
            cat_of[fid] = (pid, tops.get(pid, ""), sid, name)
    for tid, tname in tops.items():
        for row in fetch_group(tid, 0, args.offline, f"{tname}(大类)"):
            fid = int(row[0])
            master.setdefault(fid, row)
            cat_of.setdefault(fid, (tid, tname, None, None))

    print("[3/3] 生成 SQL ...")
    records = []
    no_cat, bad = [], 0
    for fid in sorted(master):
        row = master[fid]
        name = (row[2] or "").strip()
        if not name:
            bad += 1
            continue
        c1, c1n, c2, c2n = cat_of.get(fid, (None, None, None, None))
        if c1 is None:
            no_cat.append(fid)
        vals = [str(fid), sql_str(name), sql_str(row[3] if len(row) > 3 else ""),
                sql_str(row[4] if len(row) > 4 else ""), sql_str(row[1] if len(row) > 1 else "")]
        vals += [str(c1) if c1 else "NULL", sql_str(c1n),
                 str(c2) if c2 else "NULL", sql_str(c2n)]
        vals += ["NULL" if (n := to_num(row[i] if len(row) > i else None)) is None else n
                 for i, _ in NUTRIENTS]
        vals.append(sql_str(json.dumps(row, ensure_ascii=False)))
        records.append("(" + ",".join(vals) + ")")

    cols = ",".join(f"`{c}`" for c in COLUMNS)
    lines = [
        "-- 食物营养成分数据 (自动生成, 请勿手工编辑)",
        f"-- 来源: {BASE}  | 接口: FoodInfoQueryAction!queryFoodInfoList.do",
        f"-- 生成时间: {time.strftime('%Y-%m-%d %H:%M:%S')}",
        f"-- 共 {len(records)} 条; 无分类归属 {len(no_cat)} 条; 跳过无名 {bad} 条",
        "-- 前置: 先执行 db/migration_add_food_nutrition.sql 建表",
        "-- 注意: 本脚本会先清空 food_nutrition 再全量写入",
        "",
        "SET NAMES utf8mb4;",
        "DELETE FROM `food_nutrition`;",
        "",
    ]
    for i in range(0, len(records), BATCH):
        lines.append(f"INSERT INTO `food_nutrition` ({cols}) VALUES")
        lines.append(",\n".join(records[i:i + BATCH]) + ";")
        lines.append("")
    OUT_SQL.write_text("\n".join(lines), encoding="utf-8")

    size_mb = OUT_SQL.stat().st_size / 1024 / 1024
    print(f"\n完成: {OUT_SQL}")
    print(f"  {len(records)} 条记录, 文件 {size_mb:.2f} MB")
    if no_cat:
        print(f"  未归属分类的 food_id (前20): {no_cat[:20]}")


if __name__ == "__main__":
    main()