import os

import pandas as pd


def fill_time_gaps(
    data: pd.DataFrame,
    time_col: str = "time",
    freq: str = "1min",
    short_gap_threshold: int = 10,
) -> pd.DataFrame:
    """
    补齐缺失时间戳并填补数据缺口。

    Args:
        data: 包含时间列的 DataFrame
        time_col: 时间列名（默认 'time'）
        freq: 重采样频率（默认 '1min'）
        short_gap_threshold: 短缺口阈值（分钟），<=此值用线性插值，>此值用前向填充

    Returns:
        补齐时间后的 DataFrame（保留原时间列格式）
    """
    if time_col not in data.columns:
        raise ValueError(f"时间列 '{time_col}' 不存在于数据中")

    # 解析时间列并设为索引
    data = data.copy()
    data[time_col] = pd.to_datetime(data[time_col], utc=True)
    data_indexed = data.set_index(time_col)

    # 生成完整时间范围
    full_range = pd.date_range(
        start=data_indexed.index.min(), end=data_indexed.index.max(), freq=freq
    )

    # 重索引以补齐缺失时间点，同时保留原始时间戳
    combined_index = data_indexed.index.union(full_range).sort_values().unique()
    data_reindexed = data_indexed.reindex(combined_index)

    # 按列处理缺口
    for col in data_reindexed.columns:
        # 识别缺失值位置
        is_missing = data_reindexed[col].isna()

        # 计算连续缺失的长度
        missing_groups = (is_missing != is_missing.shift()).cumsum()
        gap_lengths = is_missing.groupby(missing_groups).transform("sum")

        # 短缺口：时间插值
        short_gap_mask = is_missing & (gap_lengths <= short_gap_threshold)
        if short_gap_mask.any():
            data_reindexed.loc[short_gap_mask, col] = (
                data_reindexed[col]
                .interpolate(method="time", limit_area="inside")
                .loc[short_gap_mask]
            )

        # 长缺口：前向填充
        long_gap_mask = is_missing & (gap_lengths > short_gap_threshold)
        if long_gap_mask.any():
            data_reindexed.loc[long_gap_mask, col] = (
                data_reindexed[col].ffill().loc[long_gap_mask]
            )

    # 重置索引并恢复时间列（保留原格式）
    data_result = data_reindexed.reset_index()
    data_result.rename(columns={"index": time_col}, inplace=True)

    # 保留时区信息
    data_result[time_col] = data_result[time_col].dt.strftime("%Y-%m-%dT%H:%M:%S%z")
    # 修正时区格式（Python的%z输出为+0000，需转为+00:00）
    data_result[time_col] = data_result[time_col].str.replace(
        r"(\+\d{2})(\d{2})$", r"\1:\2", regex=True
    )

    return data_result


def _cleanup_temp_files(prefix: str) -> None:
    """清理 EPANET 仿真产生的临时文件。"""
    for ext in [".inp", ".rpt", ".bin", ".out"]:
        temp_file = prefix + ext
        if os.path.exists(temp_file):
            try:
                os.remove(temp_file)
            except OSError:
                pass