# -*- coding: utf-8 -*-
"""
留学生公寓电表数据采集脚本（按终端分组并发）

并发策略:
    不同终端并行处理，同一终端同时只下发一个任务。
    每个终端分组由一个独立线程串行遍历其下所有测量点。

用法:
    python fetch_energy.py                       # 采集：每终端一个线程（默认全部终端并行）
    python fetch_energy.py <目录>                # 指定目录
    python fetch_energy.py -j 2                  # 最多 2 个终端同时采集（默认 0 = 终端数）
    python fetch_energy.py -list                 # 仅解析 xls 列出设备清单，不发起任何请求（用于核对）
    python fetch_energy.py -list -j 2            # 参数可组合

过滤: 只采集电表号（通信地址）以 260498 为前缀的表，其余跳过（-list 也会标注）

输出:
    result.jsonl  每行一条原始返回（source/line/宿舍名/电表号/txdz/cldh/response 或 error），按采集完成顺序写入
    result.xls    汇总表（按 xls 原始行序）：文件名 | 宿舍名 | 电表号 | 终端号 | 端点号 | <按 sjxmc 动态生成的数据列> | 备注

依赖: pip install xlrd xlwt
"""
import argparse
import glob
import json
import os
import threading
import time
from datetime import datetime
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor

import xlrd

API_URL = "https://energy.scnu.edu.cn/api-collect/sb/sjzc/sjzc"
TOKEN = "f1736b22-e741-4c8b-84cd-ce408ed617ed"
CJSJX = "02010100,02020100,00010000"
MAX_RETRY = 3
METER_PREFIX = "260498"  # 只采集电表号以该前缀开头的表

COL_TXDZ = "设备所属终端资产编号"      # 终端号
COL_CLDH = "电表测量点号"              # 端点号
COL_DORM = "宿舍名称（房间名称）"      # 宿舍名
COL_METER = "通信地址"                 # 电表号
HEADER_SCAN = 10  # 在前 N 行内扫描表头

XLS_FIXED_COLS = ["文件名", "宿舍名", "电表号", "终端号", "端点号"]
XLS_REMARK_COL = "备注"
COL_WIDTH = {"文件名": 44, "宿舍名": 14, "电表号": 18, "终端号": 16, "端点号": 10, XLS_REMARK_COL: 24}
DYN_COL_WIDTH = 16  # 动态数据列宽
XLS_FLUSH_EVERY = 50  # 每采集 N 条保存一次 xls，防止中途中断丢失

HEADERS = {
    "accept": "application/json, text/javascript, */*; q=0.01",
    "accept-language": "zh-CN,zh;q=0.9",
    "authorization": "Bearer " + TOKEN,
    "content-type": "application/json",
    "dnt": "1",
    "origin": "https://energy.scnu.edu.cn",
    "priority": "u=1, i",
    "referer": "https://energy.scnu.edu.cn/",
    "sec-ch-ua": '"Chromium";v="152", "Not?A_Brand";v="24", "Google Chrome";v="152"',
    "sec-ch-ua-mobile": "?0",
    "sec-ch-ua-platform": '"Windows"',
    "sec-fetch-dest": "empty",
    "sec-fetch-mode": "cors",
    "sec-fetch-site": "same-origin",
    "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/152.0.0.0 Safari/537.36",
    "x-requested-with": "XMLHttpRequest",
}


def format_id(v):
    """清洗单元格值：数字型（2.0 / 170.0）格式化为整数字符串，字符串原样去空白"""
    if isinstance(v, float):
        return str(int(v)) if v == int(v) else repr(v)
    if isinstance(v, int):
        return str(v)
    s = str(v).strip()
    return s


def read_xls(path):
    """读取 xls 中所有 sheet，自动定位含目标表头的 sheet（如"电表"），
    跳过端口对照等无关 sheet（如"Sheet1"）。
    返回 [dict(line, dorm, meter, txdz, cldh), ...]（行号 1 起始）"""
    wb = xlrd.open_workbook(path)
    rows = []
    for sh in wb.sheets():
        cols = {}
        header_row = -1
        # 前 HEADER_SCAN 行内扫描表头
        for r in range(min(HEADER_SCAN, sh.nrows)):
            for c in range(sh.ncols):
                v = str(sh.cell_value(r, c)).strip()
                if v in (COL_TXDZ, COL_CLDH, COL_DORM, COL_METER):
                    cols[v] = c
            if COL_TXDZ in cols and COL_CLDH in cols:
                header_row = r
                break
            cols = {}
        if header_row < 0:
            continue  # 本 sheet 无目标表头，跳过

        for r in range(header_row + 1, sh.nrows):
            txdz = format_id(sh.cell_value(r, cols[COL_TXDZ]))
            cldh = format_id(sh.cell_value(r, cols[COL_CLDH]))
            if not txdz or not cldh:
                continue
            if txdz == COL_TXDZ or cldh == COL_CLDH:
                continue  # 跳过重复出现的表头行
            dorm = format_id(sh.cell_value(r, cols[COL_DORM])) if COL_DORM in cols else ""
            meter = format_id(sh.cell_value(r, cols[COL_METER])) if COL_METER in cols else ""
            rows.append({"line": r + 1, "dorm": dorm, "meter": meter, "txdz": txdz, "cldh": cldh})
    return rows


def query_once(txdz, cldh):
    """按调用约定发起单次请求，带重试（线程内执行，线程安全）"""
    body = json.dumps(
        {"sjzcdx": "0", "txdz": txdz, "cldh": cldh, "cjsjx": CJSJX},
        ensure_ascii=False,
    ).encode("utf-8")
    last_err = None
    for attempt in range(1, MAX_RETRY + 1):
        req = urllib.request.Request(API_URL, data=body, headers=HEADERS, method="POST")
        try:
            with urllib.request.urlopen(req, timeout=60) as resp:
                return resp.read()
        except urllib.error.HTTPError as e:
            detail = e.read().decode("utf-8", "replace")
            last_err = "HTTP %d: %s" % (e.code, detail)
            if e.code in (401, 403):  # 鉴权失败不重试
                break
        except Exception as e:  # 网络超时等
            last_err = str(e)
        time.sleep(attempt)  # 1s/2s/3s 递增间隔
    raise RuntimeError("重试 %d 次仍失败: %s" % (MAX_RETRY, last_err))


def extract_values(response):
    """从接口返回中提取数据：返回 ({sjxmc: value, ...}, 备注信息)"""
    if not isinstance(response, dict):
        return {}, "返回非JSON结构"
    msg = str(response.get("msg") or "")
    data = response.get("data")
    if not isinstance(data, list):
        return {}, msg or "返回无data字段"
    values = {}
    for item in data:
        if isinstance(item, dict) and item.get("sjxmc") is not None:
            values[str(item["sjxmc"])] = "" if item.get("value") is None else str(item["value"])
    if not values:
        return {}, msg or "返回data为空"
    return values, msg


def save_xls(records, out_path):
    """把采集记录导出为 xls：固定列 + 按 sjxmc 动态数据列 + 备注"""
    import xlwt  # 延迟导入，-list 模式无需安装

    # 收集动态列（按首次出现顺序）
    dyn_cols = []
    for rec in records:
        for k in rec["values"]:
            if k not in dyn_cols:
                dyn_cols.append(k)

    headers = XLS_FIXED_COLS + dyn_cols + [XLS_REMARK_COL]
    wb = xlwt.Workbook(encoding="utf-8")
    ws = wb.add_sheet("采集结果")
    for i, h in enumerate(headers):
        ws.write(0, i, h)
        ws.col(i).width = 256 * COL_WIDTH.get(h, DYN_COL_WIDTH)
    for ri, rec in enumerate(records, 1):
        row = [rec["source"], rec["dorm"], rec["meter"], rec["txdz"], rec["cldh"]]
        row += [rec["values"].get(k, "") for k in dyn_cols]
        row.append(rec["remark"])
        for ci, v in enumerate(row):
            ws.write(ri, ci, v)
    wb.save(out_path)


def collect_one(task, out, lock, stats, records, xls_path):
    """单个采集任务：请求 + 写 jsonl + 更新统计（lock 保护共享输出）"""
    with lock:
        stats["done"] += 1
        total = stats["done"]
    idx, f, r = task
    rec = {
        "source": os.path.basename(f),
        "line": r["line"],
        "dorm": r["dorm"],
        "meter": r["meter"],
        "txdz": r["txdz"],
        "cldh": r["cldh"],
        "values": {},
        "remark": "",
    }
    jsonl_rec = {k: rec[k] for k in ("source", "line", "dorm", "meter", "txdz", "cldh")}
    try:
        resp_body = query_once(r["txdz"], r["cldh"])
        try:
            resp = json.loads(resp_body)
        except (ValueError, UnicodeDecodeError):
            resp = resp_body.decode("utf-8", "replace")
        jsonl_rec["response"] = resp
        rec["values"], rec["remark"] = extract_values(resp)
        if not rec["remark"]:
            rec["remark"] = "成功"
        with lock:
            stats["ok"] += 1
        print("[%d/%d] %s %s/%s 完成 data: %s"
              % (total, stats["total"],rec["dorm"], r["txdz"], r["cldh"],
                 resp_body.decode("utf-8", "replace")))
    except Exception as e:
        jsonl_rec["error"] = str(e)
        rec["remark"] = "失败: %s" % e
        with lock:
            stats["fail"] += 1
        print("[%d/%d] %s/%s 失败: %s" % (total, stats["total"], r["txdz"], r["cldh"], e))

    with lock:
        records[idx] = rec  # 按原始行序放入，xls 输出保持源表顺序
        out.write(json.dumps(jsonl_rec, ensure_ascii=False) + "\n")
        out.flush()
        # 定期保存 xls（按完成数计），防止中途中断丢失已完成部分
        if stats["done"] % XLS_FLUSH_EVERY == 0:
            try:
                save_xls([x for x in records if x], xls_path)
            except Exception as e:
                print("警告: 中途保存 xls 失败: %s" % e, file=sys.stderr)


def run_group(group_tasks, out, lock, stats, records, xls_path):
    """一个终端分组的采集线程：串行遍历该终端下所有测量点，
    保证同一终端同时只有一个在途请求"""
    for t in group_tasks:
        collect_one(t, out, lock, stats, records, xls_path)


def main():
    ap = argparse.ArgumentParser(description="留学生公寓电表数据采集")
    ap.add_argument("dir", nargs="?", default=".", help="xls 所在目录（默认当前目录）")
    ap.add_argument("-list", action="store_true", help="仅解析 xls 列出设备清单，不发起请求")
    ap.add_argument("-j", "--concurrency", type=int, default=0,
                    help="最多同时采集的终端数（默认 0 = 所有终端并行，每终端单线程串行）")
    args = ap.parse_args()
    target_dir = args.dir

    # Windows 下 glob 大小写不敏感，用集合去重，避免同一文件被处理两遍
    xls_files = sorted(set(glob.glob(os.path.join(target_dir, "*.xls"))))
    if not xls_files:
        print("目录下未找到 xls 文件:", target_dir, file=sys.stderr)
        raise SystemExit(1)

    # 解析 + 按电表号前缀过滤
    tasks = []  # (全局序号, 文件, 行数据)
    skipped = 0
    matched = 0
    for f in xls_files:
        rows = read_xls(f)
        kept = [r for r in rows if r["meter"].startswith(METER_PREFIX)]
        skipped += len(rows) - len(kept)
        matched += len(kept)
        print("=== %s === 解析到 %d 台电表，前缀 %s 匹配 %d 台，跳过 %d 台"
              % (os.path.basename(f), len(rows), METER_PREFIX, len(kept), len(rows) - len(kept)))
        if args.list:
            for r in rows:
                mark = "" if r["meter"].startswith(METER_PREFIX) else "  [跳过: 前缀不匹配]"
                print("    [行%d] 宿舍=%s 电表号=%s 终端号=%s 端点号=%s%s"
                      % (r["line"], r["dorm"], r["meter"], r["txdz"], r["cldh"], mark))
            continue
        for r in kept:
            tasks.append((len(tasks), f, r))
    if args.list:
        print("\n清单汇总: 匹配 %d 台，跳过 %d 台（前缀 %s）" % (matched, skipped, METER_PREFIX))
        return
    if not tasks:
        print("过滤后没有可采集的电表（前缀 %s）" % METER_PREFIX, file=sys.stderr)
        raise SystemExit(1)

    # 按终端号分组（保持源表出现顺序）；同一终端的任务由同一线程串行执行
    groups = {}  # txdz -> [(全局序号, 文件, 行数据), ...]
    for idx, f, r in tasks:
        groups.setdefault(r["txdz"], []).append((idx, f, r))
    group_list = list(groups.values())
    workers = len(group_list) if args.concurrency <= 0 else min(args.concurrency, len(group_list))

    print("待采集 %d 台，分为 %d 个终端组: %s"
          % (len(tasks), len(groups),
             ", ".join("%s(%d台)" % (t, len(g)) for t, g in groups.items())))
    print("最多 %d 个终端同时采集，同一终端内串行\n" % workers)

    now = datetime.now()
    time_str = now.strftime("%Y-%m-%d_%H")   # 例：2026-09-08_15

    base_name = "完整召测数据"
    jsonl_path = os.path.join(target_dir, f"{base_name}_{time_str}.jsonl")
    xls_path = os.path.join(target_dir, f"{base_name}_{time_str}.xls")
    #jsonl_path = os.path.join(target_dir, "result.jsonl")
    #xls_path = os.path.join(target_dir, "result.xls")
    records = [None] * len(tasks)
    stats = {"total": len(tasks), "done": 0, "ok": 0, "fail": 0}
    lock = threading.Lock()
    t0 = time.time()

    with open(jsonl_path, "w", encoding="utf-8") as out:
        with ThreadPoolExecutor(max_workers=workers) as pool:
            futures = [pool.submit(run_group, g, out, lock, stats, records, xls_path)
                       for g in group_list]
            for fu in futures:
                fu.result()  # 等待全部结束并传播异常

    save_xls([x for x in records if x], xls_path)
    dt = time.time() - t0
    print("\n完成: 总计 %d, 成功 %d, 失败 %d, 耗时 %.1fs (%d 个终端组并行，组内串行)"
          % (stats["total"], stats["ok"], stats["fail"], dt, workers))
    print("原始返回已写入:", jsonl_path)
    print("汇总表已写入:", xls_path)


if __name__ == "__main__":
    main()
