Blame for src/loader.py

from __future__ import annotations

"""
数据加载：支持 Google Sheets（优先）和本地 Excel 两种数据源。
处理编码异常列名、过滤汇总行、从 Date 派生 year/week/month。
"""
import logging
import os
from pathlib import Path
from typing import Any

import pandas as pd
import yaml

log = logging.getLogger("ad-agent")


def _project_root() -> Path:
    return Path(__file__).resolve().parent.parent


def load_config(config_path: Path | None = None) -> dict[str, Any]:
    path = config_path or (_project_root() / "config.yaml")
    with open(path, "r", encoding="utf-8") as f:
        return yaml.safe_load(f)


# ─────────────────────────────────────────────
# Google Sheets 数据源
# ─────────────────────────────────────────────

def _build_google_creds(config: dict[str, Any]):
    """
    构建 Google API 凭证，支持三种认证方式（按优先级）：
      1. GOOGLE_CREDENTIALS 环境变量（服务账号 JSON 字符串，云端部署推荐）
      2. GOOGLE_OAUTH_REFRESH_TOKEN 环境变量（OAuth refresh token，GitHub Actions 等场景）
      3. credentials.json 本地文件（本地开发）
    """
    from google.oauth2.service_account import Credentials as SACredentials
    from google.oauth2.credentials import Credentials as OAuthCredentials
    import json

    scopes = [
        "https://www.googleapis.com/auth/spreadsheets",
        "https://www.googleapis.com/auth/drive",
    ]

    creds_json = os.environ.get("GOOGLE_CREDENTIALS", "").strip()
    if creds_json:
        info = json.loads(creds_json)
        creds = SACredentials.from_service_account_info(info, scopes=scopes)
        log.info("使用环境变量 GOOGLE_CREDENTIALS（服务账号）认证")
        return creds

    refresh_token = os.environ.get("GOOGLE_OAUTH_REFRESH_TOKEN", "").strip()
    client_id = os.environ.get("GOOGLE_OAUTH_CLIENT_ID", "").strip()
    client_secret = os.environ.get("GOOGLE_OAUTH_CLIENT_SECRET", "").strip()
    if refresh_token and client_id and client_secret:
        creds = OAuthCredentials(
            token=None,
            refresh_token=refresh_token,
            client_id=client_id,
            client_secret=client_secret,
            token_uri="https://oauth2.googleapis.com/token",
        )
        log.info("使用 OAuth refresh token 认证")
        return creds

    gs_cfg = config.get("google_sheets", {})
    creds_file = gs_cfg.get("credentials_file", "credentials.json")
    creds_path = Path(creds_file)
    if not creds_path.is_absolute():
        creds_path = _project_root() / creds_path
    if creds_path.exists():
        creds = SACredentials.from_service_account_file(str(creds_path), scopes=scopes)
        log.info(f"使用凭证文件 {creds_path} 认证")
        return creds

    return None


def _load_from_google_sheets(config: dict[str, Any]) -> pd.DataFrame | None:
    """
    从 Google Sheets 分平台表（Meta/Google/Criteo/Bing）读取 API 同步数据。
    单一数据源，不读 Export 表，避免多源合并导致的重复。
    """
    gs_cfg = config.get("google_sheets", {})
    if not gs_cfg.get("enabled"):
        return None

    url = gs_cfg.get("spreadsheet_url", "").strip()
    if not url:
        log.warning("Google Sheets 已启用但未配置 spreadsheet_url，回退到本地 Excel")
        return None

    try:
        import gspread

        creds = _build_google_creds(config)
        if creds is None:
            log.warning("无可用 Google 认证凭证，回退到本地 Excel")
            return None

        gc = gspread.authorize(creds)
        sh = gc.open_by_url(url)

        platform_sheets = gs_cfg.get("platform_sheets", ["Meta", "Google", "Criteo", "Bing"])

        col_cfg = config.get("columns", {})
        idx_rename = {
            col_cfg.get("country_idx", 0): "country",
            col_cfg.get("platform_idx", 1): "platform",
            col_cfg.get("campaign_type_idx", 2): "campaign_type",
            col_cfg.get("campaign_name_idx", 3): "campaign_name",
            col_cfg.get("date_idx", 4): "date",
        }

        def _read_one_sheet(ws_name: str) -> pd.DataFrame | None:
            try:
                ws = sh.worksheet(ws_name)
                records = ws.get_all_values()
                if len(records) < 2:
                    return None
                header = records[0]
                seen: dict[str, int] = {}
                unique_header = []
                for col in header:
                    if col in seen:
                        seen[col] += 1
                        unique_header.append(f"{col}_{seen[col]}")
                    else:
                        seen[col] = 0
                        unique_header.append(col)
                sdf = pd.DataFrame(records[1:], columns=unique_header)
                rename = {}
                for idx, name in idx_rename.items():
                    if idx < len(sdf.columns):
                        rename[sdf.columns[idx]] = name
                sdf = sdf.rename(columns=rename)
                log.info(f"从 Sheet [{ws_name}] 读取 {len(sdf)} 行")
                return sdf
            except Exception:
                return None

        all_dfs: list[pd.DataFrame] = []
        for ps in platform_sheets:
            ps_df = _read_one_sheet(ps)
            if ps_df is not None and not ps_df.empty:
                all_dfs.append(ps_df)

        if not all_dfs:
            log.warning("Google Sheets 所有分平台工作表均无数据")
            return None

        df = pd.concat(all_dfs, ignore_index=True)
        log.info(f"从 Google Sheets 分平台表共读取 {len(df)} 行数据（{', '.join(platform_sheets)}）")
        return df

    except Exception as e:
        log.error(f"Google Sheets 读取失败: {e}，回退到本地 Excel")
        return None


# ─────────────────────────────────────────────
# 本地 Excel 数据源
# ─────────────────────────────────────────────

def _load_from_excel(config: dict[str, Any]) -> pd.DataFrame:
    """从 data 目录读取最新的 .xlsx 文件。"""
    root = _project_root()
    data_dir = Path(config.get("data_dir", "data"))
    if not data_dir.is_absolute():
        data_dir = root / data_dir

    files = sorted(
        (f for f in data_dir.glob("*.xlsx") if not f.name.startswith("~$")),
        key=lambda p: p.stat().st_mtime, reverse=True,
    )
    if not files:
        raise FileNotFoundError(f"在 {data_dir} 中未找到 .xlsx 文件")

    sheet = config.get("excel", {}).get("sheet", 0)
    df = pd.read_excel(files[0], sheet_name=sheet, engine="openpyxl")
    log.info(f"从本地 Excel 读取: {files[0].name}（{len(df)} 行）")
    return df


# ─────────────────────────────────────────────
# 写回 Google Sheets
# ─────────────────────────────────────────────

def _get_spreadsheet(config: dict[str, Any]):
    """获取 gspread spreadsheet 对象（读写权限）。"""
    import gspread

    gs_cfg = config.get("google_sheets", {})
    url = gs_cfg.get("spreadsheet_url", "").strip()
    if not url:
        raise RuntimeError("未配置 spreadsheet_url")

    creds = _build_google_creds(config)
    if creds is None:
        raise RuntimeError("无可用 Google 认证凭证（需要 GOOGLE_CREDENTIALS、GOOGLE_OAUTH_REFRESH_TOKEN 或 credentials.json）")

    gc = gspread.authorize(creds)
    return gc.open_by_url(url)


def _get_or_create_worksheet(sh, sheet_name: str, header: list[str]):
    """获取或创建指定名称的 worksheet，确保有表头。"""
    import gspread
    try:
        ws = sh.worksheet(sheet_name)
    except gspread.WorksheetNotFound:
        ws = sh.add_worksheet(title=sheet_name, rows=1000, cols=len(header))
        ws.append_row(header, value_input_option="RAW")
        log.info(f"创建新 Sheet: {sheet_name}")
    return ws


def _build_header(config: dict[str, Any]) -> list[str]:
    """根据 config 构建标准表头。"""
    col_cfg = config.get("columns", {})
    dim_count = 5
    if col_cfg.get("ad_id_idx") is not None:
        dim_count = max(dim_count, col_cfg["ad_id_idx"] + 1)
    dims = [""] * dim_count
    dims[col_cfg.get("country_idx", 0)] = "Country"
    dims[col_cfg.get("platform_idx", 1)] = "Platform"
    dims[col_cfg.get("campaign_type_idx", 2)] = "Campaign Type"
    dims[col_cfg.get("campaign_name_idx", 3)] = "Campaign Name"
    dims[col_cfg.get("date_idx", 4)] = "Date"
    if col_cfg.get("adset_name_idx") is not None:
        dims[col_cfg["adset_name_idx"]] = "Ad Set Name"
    if col_cfg.get("ad_name_idx") is not None:
        dims[col_cfg["ad_name_idx"]] = "Ad Name"
    if col_cfg.get("ad_id_idx") is not None:
        dims[col_cfg["ad_id_idx"]] = "Ad ID"
    return dims + config.get("metrics", [])


def write_api_data_to_sheets(
    df: pd.DataFrame,
    config_path: Path | None = None,
) -> dict[str, int]:
    """
    将 API 拉取的数据按平台写入不同的 Google Sheets worksheet。
    每个平台一个 sheet（如 Meta、Google、Criteo），自动创建不存在的 sheet。
    按 (country, date, campaign_name, ad_name) 去重，只追加新行。
    """
    config = load_config(config_path)
    sh = _get_spreadsheet(config)
    col_cfg = config.get("columns", {})
    country_idx = col_cfg.get("country_idx", 0)
    date_idx = col_cfg.get("date_idx", 4)
    campaign_name_idx = col_cfg.get("campaign_name_idx", 3)
    metrics_list = config.get("metrics", [])
    header = _build_header(config)

    key_indices = [country_idx, date_idx, campaign_name_idx]

    platforms = df["platform"].unique() if "platform" in df.columns else []
    total_added = 0

    for platform in platforms:
        p_df = df[df["platform"] == platform]
        ws = _get_or_create_worksheet(sh, platform, header)

        all_values = ws.get_all_values()
        col_order = all_values[0] if all_values else header
        existing_rows = all_values[1:] if len(all_values) > 1 else []

        existing_keys = set()
        for row in existing_rows:
            if len(row) > max(key_indices):
                key = tuple(row[i].strip() for i in key_indices)
                existing_keys.add(key)

        new_rows = []
        for _, row_data in p_df.iterrows():
            dt = row_data.get("date")
            date_str = dt.strftime("%Y-%m-%d") if hasattr(dt, "strftime") else str(dt).split(" ")[0]
            campaign_name = str(row_data.get("campaign_name", ""))
            ad_name = str(row_data.get("ad_name", ""))
            country = str(row_data.get("country", ""))
            key = (country, date_str, campaign_name, ad_name)

            if key in existing_keys:
                continue

            out_row = [""] * len(col_order)
            dim_map = {
                col_cfg.get("country_idx", 0): str(row_data.get("country", "")),
                col_cfg.get("platform_idx", 1): platform,
                col_cfg.get("campaign_type_idx", 2): str(row_data.get("campaign_type", "")),
                campaign_name_idx: campaign_name,
                date_idx: date_str,
            }
            adset_name_idx = col_cfg.get("adset_name_idx")
            ad_name_idx = col_cfg.get("ad_name_idx")
            ad_id_idx = col_cfg.get("ad_id_idx")
            if adset_name_idx is not None:
                dim_map[adset_name_idx] = str(row_data.get("adset_name", ""))
            if ad_name_idx is not None:
                dim_map[ad_name_idx] = ad_name
            if ad_id_idx is not None:
                dim_map[ad_id_idx] = str(row_data.get("ad_id", ""))
            for idx, val in dim_map.items():
                if idx < len(out_row):
                    out_row[idx] = val

            for col_name in col_order:
                if col_name in metrics_list and col_name in row_data.index:
                    val = row_data[col_name]
                    if val is None or (isinstance(val, float) and val != val):
                        continue
                    col_idx = col_order.index(col_name)
                    v = float(val)
                    out_row[col_idx] = str(int(v)) if v == int(v) else str(v)

            new_rows.append(out_row)
            existing_keys.add(key)

        if new_rows:
            ws.append_rows(new_rows, value_input_option="USER_ENTERED")
            log.info(f"写入 Sheet [{platform}]: {len(new_rows)} 行新数据")
            total_added += len(new_rows)

    return {"added": total_added}


# ─────────────────────────────────────────────
# 统一加载入口
# ─────────────────────────────────────────────

_data_source: str = "unknown"


def load_raw(config_path: Path | None = None) -> pd.DataFrame:
    """优先 Google Sheets 分平台表，回退本地 Excel。"""
    global _data_source
    config = load_config(config_path)
    df = _load_from_google_sheets(config)
    if df is not None:
        _data_source = "sheets"
    else:
        df = _load_from_excel(config)
        _data_source = "excel"
    return df


def normalize_columns(df: pd.DataFrame, config: dict[str, Any]) -> pd.DataFrame:
    """用列索引将编码异常的列重命名为可读字段名。"""
    col_cfg = config.get("columns", {})
    idx_map = {
        col_cfg.get("country_idx", 0): "country",
        col_cfg.get("platform_idx", 1): "platform",
        col_cfg.get("campaign_type_idx", 2): "campaign_type",
        col_cfg.get("campaign_name_idx", 3): "campaign_name",
        col_cfg.get("date_idx", 4): "date",
    }
    rename = {}
    for idx, name in idx_map.items():
        if idx < len(df.columns):
            rename[df.columns[idx]] = name
    return df.rename(columns=rename)


def _clean_numeric(s: pd.Series) -> pd.Series:
    """
    清洗 Google Sheets 格式化字符串。
    带 % 的值先标记，去掉符号后除以 100 还原为小数。
    """
    raw = s.astype(str)
    is_pct = raw.str.contains('%', na=False)
    cleaned = raw.str.replace(r'[$,%\s]', '', regex=True)
    values = pd.to_numeric(cleaned, errors="coerce")
    values[is_pct] = values[is_pct] / 100.0
    return values


def load_daily_data(config_path: Path | None = None) -> pd.DataFrame:
    """
    加载并清洗数据，只保留日度明细行。
    自动从 date 列派生 year、week（ISO week）、month。
    兼容本地 Excel（数值）和 Google Sheets（格式化字符串）。
    """
    config = load_config(config_path)
    df = load_raw(config_path)
    df = normalize_columns(df, config)

    # Excel 层级结构中，子行的维度列为空字符串，需要前向填充
    # API 写入的 Sheets 分平台表每行都有完整字段，不需要 ffill
    if _data_source == "excel":
        for dim in ["country", "platform", "campaign_type", "campaign_name"]:
            if dim in df.columns:
                df[dim] = df[dim].replace("", pd.NA).ffill()

    df["date"] = pd.to_datetime(df["date"], format="mixed", errors="coerce")
    df = df[df["date"].notna()].copy()

    df = df[df["country"].notna() & (~df["country"].astype(str).str.contains("Total|Applied filters", na=False))].copy()
    df = df[df["platform"].notna() & (df["platform"].astype(str).str.strip() != "Total")].copy()
    df = df[df["campaign_type"].notna() & (df["campaign_type"].astype(str).str.strip() != "Total")].copy()
    df = df[df["campaign_name"].notna() & (df["campaign_name"].astype(str).str.strip() != "Total")].copy()

    df["year"] = df["date"].dt.isocalendar().year.astype(int)
    df["week"] = df["date"].dt.isocalendar().week.astype(int)
    df["month"] = df["date"].dt.month

    metrics = config.get("metrics", [])
    for col in metrics:
        if col in df.columns:
            df[col] = _clean_numeric(df[col])

    df = df.reset_index(drop=True)

    df = _fix_campaign_type_from_name(df)

    dedup_cols = ["date", "platform", "campaign_name", "country"]
    if all(c in df.columns for c in dedup_cols):
        before = len(df)
        df = df.drop_duplicates(subset=dedup_cols, keep="last")
        dropped = before - len(df)
        if dropped:
            log.info(f"去重删除 {dropped} 行重复数据")
        df = df.reset_index(drop=True)

    return df


_VALID_FUNNELS = {"Awareness", "Consideration", "Conversion"}


def _fix_campaign_type_from_name(df: pd.DataFrame) -> pd.DataFrame:
    """根据 campaign_name 中的第二段修正 campaign_type（如名称含明确漏斗标识）。"""
    if "campaign_name" not in df.columns or "campaign_type" not in df.columns:
        return df
    parts = df["campaign_name"].astype(str).str.split("_", n=2)
    name_funnel = parts.str[1]
    mask = name_funnel.isin(_VALID_FUNNELS) & (name_funnel != df["campaign_type"])
    if mask.any():
        log.info(f"根据 campaign_name 修正 {mask.sum()} 行 campaign_type")
        df.loc[mask, "campaign_type"] = name_funnel[mask]
    return df


# 兼容旧接口
load_weekly_data = load_daily_data


def get_available_weeks(df: pd.DataFrame) -> list[tuple[int, int]]:
    """返回可用的 (year, week) 列表，按时间排序。"""
    pairs = df[["year", "week"]].drop_duplicates().sort_values(["year", "week"])
    return list(pairs.itertuples(index=False, name=None))


def get_latest_week(df: pd.DataFrame) -> tuple[int, int]:
    """返回数据中最新一周的 (year, week)。"""
    weeks = get_available_weeks(df)
    if not weeks:
        raise ValueError("数据中无可用周数据")
    return weeks[-1]


def get_available_months(df: pd.DataFrame) -> list[tuple[int, int]]:
    """返回可用的 (year, month) 列表，按时间排序。"""
    cal_year = df["date"].dt.year
    pairs = pd.DataFrame({"year": cal_year, "month": df["month"]}).drop_duplicates().sort_values(["year", "month"])
    return list(pairs.itertuples(index=False, name=None))