jdf-energy-collector/main.py

import argparse
import csv
import json
import logging
import sys
import time
from datetime import datetime, timedelta
from logging.handlers import TimedRotatingFileHandler
from pathlib import Path

import requests
import yaml


LOGIN_RETRIES = 3
LOGIN_RETRY_INTERVAL_SECONDS = 60
QUERY_TIMEOUT_RETRIES = 3
QUERY_TIMEOUT_RETRY_INTERVAL_SECONDS = 10
DEFAULT_LOG_FILE = "logs/run.log"
DEFAULT_LOG_RETENTION_DAYS = 3
DEFAULT_ADDPOINTDATUM_TIMEOUT_SECONDS = 5
DEFAULT_CALCAGG_TIMEOUT_SECONDS = 60 * 50  # 50minute
DEFAULT_TOKEN_REFRESH_INTERVAL_HOURS = 6

LOGGER = logging.getLogger("jdf_energy_collector")


def load_config(config_path: str) -> dict:
    with open(config_path, "r", encoding="utf-8") as file:
        config = yaml.safe_load(file) or {}

    required_sections = ["energy_chaowang", "points", "services", "schedule"]
    missing_sections = [section for section in required_sections if section not in config]
    if missing_sections:
        raise ValueError(f"配置文件缺少配置段: {', '.join(missing_sections)}")

    return config


def setup_logging(config: dict) -> None:
    log_config = config.get("logging") or {}
    log_file = Path(log_config.get("file", DEFAULT_LOG_FILE))
    retention_days = int(log_config.get("retention_days", DEFAULT_LOG_RETENTION_DAYS))

    log_file.parent.mkdir(parents=True, exist_ok=True)
    formatter = logging.Formatter("%(asctime)s %(levelname)s %(message)s")

    file_handler = TimedRotatingFileHandler(
        log_file,
        when="midnight",
        interval=1,
        backupCount=retention_days,
        encoding="utf-8",
    )
    file_handler.setFormatter(formatter)

    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setFormatter(formatter)

    LOGGER.handlers.clear()
    LOGGER.setLevel(logging.INFO)
    LOGGER.addHandler(file_handler)
    LOGGER.addHandler(console_handler)
    LOGGER.propagate = False


def mask_sensitive(value):
    if isinstance(value, dict):
        return {
            key: "***" if str(key).lower() in {"authorization", "password", "token"} else mask_sensitive(item)
            for key, item in value.items()
        }
    if isinstance(value, list):
        return [mask_sensitive(item) for item in value]
    return value


def to_log_text(value) -> str:
    return json.dumps(mask_sensitive(value), ensure_ascii=False, default=str)


def post_json(name: str, url: str, payload: dict, headers: dict | None = None, timeout: int | float | None = None):
    LOGGER.info(
        "HTTP请求开始: name=%s url=%s timeout=%s headers=%s payload=%s",
        name,
        url,
        timeout,
        to_log_text(headers or {}),
        to_log_text(payload),
    )
    started_at = time.monotonic()

    try:
        response = requests.post(url, headers=headers, json=payload, timeout=timeout)
        elapsed_ms = int((time.monotonic() - started_at) * 1000)
        try:
            response_body = response.json()
        except ValueError:
            response_body = response.text[:2000]

        LOGGER.info(
            "HTTP请求结束: name=%s status_code=%s elapsed_ms=%s response=%s",
            name,
            response.status_code,
            elapsed_ms,
            to_log_text(response_body),
        )
        return response
    except Exception as exc:
        elapsed_ms = int((time.monotonic() - started_at) * 1000)
        LOGGER.exception("HTTP请求异常: name=%s elapsed_ms=%s error=%s", name, elapsed_ms, exc)
        raise


def read_points(points_file: str) -> tuple[list[str], dict[str, str]]:
    with open(points_file, encoding="utf-8-sig", newline="") as file:
        reader = csv.DictReader(file, delimiter=",")
        if reader.fieldnames is None:
            raise ValueError(f"点位文件为空: {points_file}")

        reader.fieldnames = [name.strip() if name else name for name in reader.fieldnames]
        if "id" not in reader.fieldnames or "point_id" not in reader.fieldnames:
            raise ValueError(f"{points_file} 必须包含 id 和 point_id 两列")

        tag_ids = []
        id_to_point_id = {}
        seen_tag_ids = set()
        seen_point_ids = set()
        duplicated_tag_ids = set()
        duplicated_point_ids = set()

        for row in reader:
            tag_id = (row.get("id") or "").strip()
            point_id = (row.get("point_id") or "").strip()
            if not tag_id or not point_id:
                continue

            if tag_id in seen_tag_ids:
                duplicated_tag_ids.add(tag_id)
            if point_id in seen_point_ids:
                duplicated_point_ids.add(point_id)

            seen_tag_ids.add(tag_id)
            seen_point_ids.add(point_id)
            tag_ids.append(tag_id)
            id_to_point_id[tag_id] = point_id

    if duplicated_tag_ids:
        raise ValueError(f"{points_file} 的 id 列存在重复值: {sorted(duplicated_tag_ids)}")
    if duplicated_point_ids:
        raise ValueError(f"{points_file} 的 point_id 列存在重复值: {sorted(duplicated_point_ids)}")
    if not tag_ids:
        raise ValueError(f"点位文件没有有效点位: {points_file}")

    LOGGER.info("点位文件读取完成: file=%s count=%s", points_file, len(tag_ids))
    return tag_ids, id_to_point_id


def load_points_from_config(config: dict) -> tuple[list[str], dict[str, str]]:
    return read_points(str(Path(config["points"]["file"])))


def calculate_hour_range(now: datetime | None = None) -> tuple[datetime, datetime, datetime]:
    current = now or datetime.now()
    current_hour = current.replace(minute=0, second=0, microsecond=0)
    start_time = current_hour - timedelta(hours=1)
    end_time = current_hour - timedelta(seconds=1)
    return start_time, end_time, current_hour


def format_api_time(value: datetime) -> str:
    return value.strftime("%Y-%m-%d %H:%M:%S")


def local_timestamp(value: datetime) -> int:
    return int(value.timestamp())


def parse_full_time(full_time: str) -> datetime:
    return datetime.strptime(full_time, "%Y-%m-%d %H:%M:%S")


def login(energy_config: dict) -> str:
    login_url = f"{energy_config['base_url'].rstrip('/')}/api/thingshome-admin/sys/service/getToken"
    payload = {
        "username": energy_config["username"],
        "password": energy_config["password"],
    }
    timeout = energy_config.get("login_timeout_seconds", 15)

    for attempt in range(1, LOGIN_RETRIES + 1):
        try:
            response = post_json("energy_chaowang_login", login_url, payload, timeout=timeout)
            response.raise_for_status()
            result = response.json()
            if result.get("code") == 0 and result.get("token"):
                return result["token"]

            LOGGER.warning("登录失败，第 %s 次: %s", attempt, to_log_text(result))
        except Exception as exc:
            LOGGER.warning("登录异常，第 %s 次: %s", attempt, exc)

        if attempt < LOGIN_RETRIES:
            time.sleep(LOGIN_RETRY_INTERVAL_SECONDS)

    raise RuntimeError("登录失败，已达到最大重试次数")


class TokenManager:
    def __init__(self, energy_config: dict):
        self.energy_config = energy_config
        self.token = None
        self.last_login_at = None
        self.refresh_interval = timedelta(
            hours=float(energy_config.get("token_refresh_interval_hours", DEFAULT_TOKEN_REFRESH_INTERVAL_HOURS))
        )

    def get_token(self) -> str:
        now = datetime.now()
        if self.token and self.last_login_at and now < self.last_login_at + self.refresh_interval:
            LOGGER.info("复用登录 Token，下次登录时间: %s", format_api_time(self.last_login_at + self.refresh_interval))
            return self.token

        return self.refresh_token()

    def refresh_token(self) -> str:
        LOGGER.info("开始获取登录 Token")
        self.token = login(self.energy_config)
        self.last_login_at = datetime.now()
        LOGGER.info("登录 Token 获取成功，下次登录时间: %s", format_api_time(self.last_login_at + self.refresh_interval))
        return self.token


def query_hour_data(energy_config: dict, token: str, tag_ids: list[str], start_time: datetime, end_time: datetime) -> dict:
    query_url = f"{energy_config['base_url'].rstrip('/')}/api/thingshome-ems/realTime/selectByCustom"
    headers = {
        "Content-Type": "application/json",
        "token": token,
    }
    payload = {
        "type": "1",
        "tagNameList": tag_ids,
        "energyTypeList": ["electric power"],
        "tb": "0",
        "startTime": format_api_time(start_time),
        "endTime": format_api_time(end_time),
        "dateCode": "h",
    }
    timeout = energy_config.get("query_timeout_seconds", 30)

    for attempt in range(1, QUERY_TIMEOUT_RETRIES + 1):
        try:
            response = post_json("energy_chaowang_query_hour", query_url, payload, headers=headers, timeout=timeout)
            response.raise_for_status()
            result = response.json()
            if result.get("code") != 0:
                raise RuntimeError(f"查询接口返回失败: {result}")
            return result
        except requests.Timeout as exc:
            LOGGER.warning("查询接口 timeout，第 %s 次: %s", attempt, exc)
            if attempt < QUERY_TIMEOUT_RETRIES:
                time.sleep(QUERY_TIMEOUT_RETRY_INTERVAL_SECONDS)

    raise RuntimeError("查询接口 timeout，已达到最大重试次数")


def query_hour_data_with_token_refresh(
    energy_config: dict,
    token_manager: TokenManager,
    tag_ids: list[str],
    start_time: datetime,
    end_time: datetime,
) -> dict:
    token = token_manager.get_token()
    try:
        return query_hour_data(energy_config, token, tag_ids, start_time, end_time)
    except requests.HTTPError as exc:
        if exc.response is None or exc.response.status_code != 401:
            raise

    LOGGER.warning("查询接口返回 401，重新登录获取 Token 后重试")
    token = token_manager.refresh_token()
    return query_hour_data(energy_config, token, tag_ids, start_time, end_time)


def extract_records(result: dict) -> list[dict]:
    records = []
    data_root = result.get("data") or {}

    for point_data in data_root.get("data") or []:
        records.extend(point_data.get("dataList") or [])

    if records:
        return records

    return data_root.get("trend") or []


def call_addpointdatum(basedataportal_base_url: str, records: list[dict], id_to_point_id: dict[str, str]) -> list[str]:
    point_ids = []
    point_id_set = set()
    data_by_point_id = {}

    for record in records:
        tag_id = str(record.get("id", "")).strip()
        point_id = id_to_point_id.get(tag_id)
        if not point_id:
            LOGGER.warning("跳过未在 CSV 中配置的点位: %s", tag_id)
            continue

        name = record.get("name")
        full_time = record.get("fullTime")
        this_value = record.get("thisValue")
        if full_time is None or this_value is None:
            LOGGER.warning("跳过字段不完整的点位: %s", to_log_text(record))
            continue

        #if this_value == 0:
            #LOGGER.info("跳过 thisValue 为 0 的点位: %s", to_log_text(record))
            #continue

        full_time_ts = local_timestamp(parse_full_time(full_time))
        LOGGER.info("点位数据: name=%s fullTime=%s thisValue=%s", name, full_time, this_value)

        data_by_point_id.setdefault(point_id, []).append({"ts": full_time_ts, "value": str(this_value)})
        if point_id not in point_id_set:
            point_ids.append(point_id)
            point_id_set.add(point_id)

    if not point_ids:
        return []

    url = f"{basedataportal_base_url.rstrip('/')}/ai/addpointdatum"
    payload = [{"point_id": point_id, "data": data_by_point_id[point_id]} for point_id in point_ids]
    response = post_json(
        "basedataportal_addpointdatum",
        url,
        payload,
        timeout=DEFAULT_ADDPOINTDATUM_TIMEOUT_SECONDS,
    )
    response.raise_for_status()
    result = response.json()
    if result.get("state") != 0:
        raise RuntimeError(f"addpointdatum 接口返回异常: {to_log_text(result)}")

    LOGGER.info("addpointdatum 写入完成: count=%s", len(point_ids))
    return point_ids


def call_calcagg(calcagg_base_url: str, point_ids: list[str], begin: int, end: int) -> None:
    url = f"{calcagg_base_url.rstrip('/')}/api/calcagg/calc_agg_points_range"
    payload = {
        "end": end,
        "begin": begin,
        "sync_run": True,
        "point_ids": point_ids,
        "operator_name": "jdf-energy-collector",
    }
    response = post_json("calcagg_calc_agg_points_range", url, payload, timeout=DEFAULT_CALCAGG_TIMEOUT_SECONDS)
    response.raise_for_status()
    result = response.json()
    if result.get("state") != 0:
        LOGGER.warning("计算聚合接口返回异常: %s", to_log_text(result))


def run_once(config: dict, token_manager: TokenManager, tag_ids: list[str], id_to_point_id: dict[str, str]) -> None:
    start_time, end_time, current_hour = calculate_hour_range()
    LOGGER.info(
        "开始执行小时数据采集: %s - %s, ts:%s-%s",
        format_api_time(start_time),
        format_api_time(end_time),
        start_time,
        end_time,
    )

    result = query_hour_data_with_token_refresh(
        config["energy_chaowang"], token_manager, tag_ids, start_time, end_time
    )
    records = extract_records(result)

    written_point_ids = call_addpointdatum(config["services"]["basedataportal"], records, id_to_point_id)
    if not written_point_ids:
        LOGGER.info("没有成功写入 addpointdatum 的点位，跳过计算聚合接口调用")
        return

    begin = local_timestamp(start_time)
    end = local_timestamp(end_time)
    call_calcagg(config["services"]["calcagg"], written_point_ids, begin, end)
    LOGGER.info("本轮采集完成，写入点位数量: %s", len(written_point_ids))


def wait_until_next_run(schedule_minute: int) -> None:
    now = datetime.now()
    next_run = now.replace(minute=schedule_minute, second=0, microsecond=0)
    if now > next_run:
        next_run += timedelta(hours=1)

    sleep_seconds = max(0, (next_run - now).total_seconds())
    LOGGER.info("下次执行时间: %s", format_api_time(next_run))
    time.sleep(sleep_seconds)


def main() -> None:
    parser = argparse.ArgumentParser(description="小时能源数据采集入库程序")
    parser.add_argument("--config", default="config.yaml", help="YAML 配置文件路径")
    args = parser.parse_args()

    config = load_config(args.config)
    setup_logging(config)
    tag_ids, id_to_point_id = load_points_from_config(config)
    token_manager = TokenManager(config["energy_chaowang"])

    schedule_minute = int(config["schedule"].get("minute", 30))
    if schedule_minute < 0 or schedule_minute > 59:
        raise ValueError("schedule.minute 必须在 0 到 59 之间")

    while True:
        wait_until_next_run(schedule_minute)
        try:
            run_once(config, token_manager, tag_ids, id_to_point_id)
        except Exception as exc:
            LOGGER.exception("本轮采集失败: %s", exc)


if __name__ == "__main__":
    main()