#!/usr/bin/env python3
"""
月度学而思/讯飞/希沃/京东方渠道订单汇总报表

功能：
  - 按指定月份，汇总这4个渠道每一天的：
      新增注册数、订单数（剔除退费）、GMV（有效订单金额加和）
  - 自动剔除测试账号
  - 输出 Excel 文件到 output/ 目录

渠道映射（sale_channel，仅对 key_from='app-active-h5-0-0' 的订单生效）：
  21 = 学而思
  22 = 讯飞
  26 = 希沃
  27 = 京东方

新增注册：download_channel LIKE '%学而思%' / '%讯飞%' / '%希沃%' / '%京东方%'
"""

import os
import sys
from datetime import datetime, timedelta
import calendar
import psycopg2
import pandas as pd

# ── 数据库连接 ──────────────────────────────────────────────
DB_HOST = "bj-postgres-16pob4sg.sql.tencentcdb.com"
DB_PORT = 28591
DB_USER = "ai_member"
DB_NAME = "vala_bi"
DB_PASSWORD = os.environ.get("PG_ONLINE_PASSWORD", "")

if not DB_PASSWORD:
    # 尝试从 secrets.env 读取
    secrets_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "secrets.env")
    if os.path.exists(secrets_path):
        with open(secrets_path) as f:
            for line in f:
                if line.startswith("PG_ONLINE_PASSWORD="):
                    DB_PASSWORD = line.strip().split("=", 1)[1].strip("'\"")
                    break

# ── 渠道配置 ────────────────────────────────────────────────
# (sale_channel值, 渠道中文名, CPS分成比例, None表示不用CPS分成)
CHANNELS = [
    ("21", "学而思", 0.45),
    ("22", "讯飞", None),
    ("26", "希沃", 0.40),
    ("27", "京东方", 0.30),
]

# ── 参数 ────────────────────────────────────────────────────
def parse_month(arg: str) -> str:
    """将用户输入解析为 YYYY-MM 格式，默认当前月"""
    arg = arg.strip()
    if len(arg) == 7 and arg[4] == "-":
        return arg
    if len(arg) == 6 and arg.isdigit():
        return f"{arg[:4]}-{arg[4:]}"
    # 默认当前月
    return datetime.now().strftime("%Y-%m")


def get_conn():
    return psycopg2.connect(
        host=DB_HOST,
        port=DB_PORT,
        user=DB_USER,
        password=DB_PASSWORD,
        dbname=DB_NAME,
        connect_timeout=10,
    )


# ── SQL 定义 ────────────────────────────────────────────────

# 新增注册：按天 + 渠道汇总（这4个渠道 download_channel 关键字包含匹配）
SQL_NEW_USERS = """
SELECT
    DATE(created_at) AS reg_date,
    CASE
        WHEN download_channel LIKE '%%学而思%%' THEN '学而思'
        WHEN download_channel LIKE '%%讯飞%%'   THEN '讯飞'
        WHEN download_channel LIKE '%%希沃%%'   THEN '希沃'
        WHEN download_channel LIKE '%%京东方%%' THEN '京东方'
    END AS channel_name,
    COUNT(*) AS new_users
FROM bi_vala_app_account
WHERE status = 1
  AND deleted_at IS NULL
  AND created_at >= %s
  AND created_at <  %s
  AND (
      download_channel LIKE '%%学而思%%'
   OR download_channel LIKE '%%讯飞%%'
   OR download_channel LIKE '%%希沃%%'
   OR download_channel LIKE '%%京东方%%'
  )
GROUP BY reg_date, channel_name
ORDER BY reg_date;
"""

# 有效订单（剔除全部退费订单）：端内 + sale_channel 在 21/22/26/27
# 剔除逻辑：排除 order_status=4 且 bi_refund_order.status=3 的订单
# 逐条查出，在 Python 中聚合（因为需要保留 trade_no 明细）
SQL_VALID_ORDERS = """
SELECT
    DATE(o.pay_success_date) AS order_date,
    CASE o.sale_channel
        WHEN 21 THEN '学而思'
        WHEN 22 THEN '讯飞'
        WHEN 26 THEN '希沃'
        WHEN 27 THEN '京东方'
    END AS channel_name,
    o.pay_amount_int,
    o.trade_no
FROM bi_vala_order o
INNER JOIN bi_vala_app_account a ON o.account_id = a.id AND a.status = 1
WHERE o.key_from = 'app-active-h5-0-0'
  AND o.sale_channel IN (21, 22, 26, 27)
  AND o.pay_success_date >= %s
  AND o.pay_success_date <  %s
  AND o.deleted_at IS NULL
  -- 剔除已退费订单
  AND NOT EXISTS (
      SELECT 1 FROM bi_refund_order r
      WHERE r.trade_no = o.trade_no
        AND r.status = 3
        AND r.deleted_at IS NULL
  )
  AND o.order_status != 4
ORDER BY order_date, channel_name;
"""


# ── 主逻辑 ──────────────────────────────────────────────────

def run(month_str: str):
    print(f"开始生成 {month_str} 学而思/讯飞/希沃/京东方渠道汇总报表...")

    # 解析月份起止
    year, month = int(month_str[:4]), int(month_str[5:])
    _, last_day = calendar.monthrange(year, month)
    start_ts = f"{month_str}-01 00:00:00+08:00"
    # 下月1号（用 datetime 自动处理跨月/跨年）
    next_month = datetime(year, month, 1) + timedelta(days=last_day)
    end_ts = next_month.strftime("%Y-%m-%d 00:00:00+08:00")

    conn = get_conn()
    try:
        # 1. 新增注册
        cur = conn.cursor()
        cur.execute(SQL_NEW_USERS, (start_ts, end_ts))
        rows_new = cur.fetchall()
        df_new = pd.DataFrame(rows_new, columns=[desc[0] for desc in cur.description]) if rows_new else pd.DataFrame(columns=["reg_date", "channel_name", "new_users"])
        cur.close()
        print(f"  新增注册记录: {len(df_new)} 行")

        # 2. 有效订单（逐条）
        cur = conn.cursor()
        cur.execute(SQL_VALID_ORDERS, (start_ts, end_ts))
        rows_ord = cur.fetchall()
        df_orders_raw = pd.DataFrame(rows_ord, columns=[desc[0] for desc in cur.description]) if rows_ord else pd.DataFrame(columns=["order_date", "channel_name", "pay_amount_int", "trade_no"])
        cur.close()
        print(f"  有效订单记录: {len(df_orders_raw)} 行")

        # 按 channel_name + order_date 聚合成每日汇总
        if not df_orders_raw.empty:
            df_orders = df_orders_raw.groupby(["channel_name", "order_date"], as_index=False).agg(
                order_cnt=("pay_amount_int", "count"),
                gmv_cents=("pay_amount_int", "sum"),
                trade_nos=("trade_no", lambda x: ";".join(str(t) for t in x)),
            )
        else:
            df_orders = pd.DataFrame(columns=["channel_name", "order_date", "order_cnt", "gmv_cents", "trade_nos"])
    finally:
        conn.close()

    # 3. 构建每个渠道的结果表（每一天一行）
    date_range = [
        f"{month_str}-{d:02d}" for d in range(1, last_day + 1)
    ]

    channel_names = [ch[1] for ch in CHANNELS]
    channel_dfs = {}
    channel_totals = {}

    for ch_val, ch_name, cps_rate in CHANNELS:
        if cps_rate is not None:
            # ── 学而思/希沃/京东方：只展示有订单的日期，无新增注册列，有CPS分成 ──
            rows = []
            for d in date_range:
                day_ord = df_orders[(df_orders["order_date"].astype(str) == d) & (df_orders["channel_name"] == ch_name)]
                if day_ord.empty:
                    continue
                order_cnt = int(day_ord["order_cnt"].sum())
                gmv = day_ord["gmv_cents"].sum() / 100.0
                trade_nos = ";".join(day_ord["trade_nos"].values)
                rows.append({
                    "日期": d,
                    "订单数": order_cnt,
                    "GMV（元）": round(gmv, 2),
                    "CPS分成（元）": round(gmv * cps_rate, 2),
                    "订单号": trade_nos,
                })

            if rows:
                df_sheet = pd.DataFrame(rows)
                total = {
                    "日期": "合计",
                    "订单数": df_sheet["订单数"].sum(),
                    "GMV（元）": round(df_sheet["GMV（元）"].sum(), 2),
                    "CPS分成（元）": round(df_sheet["CPS分成（元）"].sum(), 2),
                    "订单号": "",
                }
            else:
                df_sheet = pd.DataFrame(columns=["日期", "订单数", "GMV（元）", "CPS分成（元）", "订单号"])
                total = {"日期": "合计", "订单数": 0, "GMV（元）": 0.0, "CPS分成（元）": 0.0, "订单号": ""}
            df_sheet = pd.concat([df_sheet, pd.DataFrame([total])], ignore_index=True)
            channel_dfs[ch_name] = df_sheet
            channel_totals[ch_name] = total

        else:
            # ── 讯飞：每天一行，含新增注册 + CPA分成 + 订单号 ──
            rows = []
            for d in date_range:
                day_new = df_new[(df_new["reg_date"].astype(str) == d) & (df_new["channel_name"] == ch_name)]
                new_users = int(day_new["new_users"].sum()) if not day_new.empty else 0

                day_ord = df_orders[(df_orders["order_date"].astype(str) == d) & (df_orders["channel_name"] == ch_name)]
                if not day_ord.empty:
                    order_cnt = int(day_ord["order_cnt"].sum())
                    gmv = day_ord["gmv_cents"].sum() / 100.0
                    trade_nos = ";".join(day_ord["trade_nos"].values)
                else:
                    order_cnt = 0
                    gmv = 0.0
                    trade_nos = ""

                rows.append({
                    "日期": d,
                    "新增注册": new_users,
                    "订单数": order_cnt,
                    "GMV（元）": round(gmv, 2),
                    "CPA分成（元）": round(new_users * 7, 2),
                    "订单号": trade_nos,
                })

            df_sheet = pd.DataFrame(rows)
            total = {
                "日期": "合计",
                "新增注册": df_sheet["新增注册"].sum(),
                "订单数": df_sheet["订单数"].sum(),
                "GMV（元）": round(df_sheet["GMV（元）"].sum(), 2),
                "CPA分成（元）": round(df_sheet["CPA分成（元）"].sum(), 2),
                "订单号": "",
            }
            df_sheet = pd.concat([df_sheet, pd.DataFrame([total])], ignore_index=True)
            channel_dfs[ch_name] = df_sheet
            channel_totals[ch_name] = total

    # 4. 输出 Excel（4个Sheet，每个渠道一张）
    output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "output")
    os.makedirs(output_dir, exist_ok=True)
    filename = f"学而思讯飞希沃京东方_{month_str}_渠道汇总.xlsx"
    output_path = os.path.join(output_dir, filename)

    with pd.ExcelWriter(output_path, engine="openpyxl") as writer:
        for ch_name in channel_names:
            df_sheet = channel_dfs[ch_name]
            df_sheet.to_excel(writer, sheet_name=ch_name, index=False)
            # 将"订单号"列设为文本格式，防止长数字被科学计数法显示
            ws = writer.sheets[ch_name]
            # 找到"订单号"列的位置
            headers = list(df_sheet.columns)
            if "订单号" in headers:
                col_idx = headers.index("订单号")
                col_letter = chr(65 + col_idx)  # A=65
                for row in ws.iter_rows(min_row=2, max_row=ws.max_row, min_col=col_idx+1, max_col=col_idx+1):
                    for cell in row:
                        if cell.value and str(cell.value).strip():
                            cell.number_format = '@'  # 文本格式

    # 5. 打印摘要
    print(f"\n✅ 报表已生成: {output_path}")
    print(f"   共 {last_day} 天数据 × 4 个渠道")
    grand_orders = sum(t["订单数"] for t in channel_totals.values())
    grand_gmv = sum(t["GMV（元）"] for t in channel_totals.values())
    for ch_name in channel_names:
        t = channel_totals[ch_name]
        extra = ""
        if "新增注册" in t:
            extra += f", 新增注册={t['新增注册']}"
        if "CPA分成（元）" in t:
            extra += f", CPA分成={t['CPA分成（元）']}元"
        if "CPS分成（元）" in t:
            extra += f", CPS分成={t['CPS分成（元）']}元"
        print(f"   {ch_name}: 订单数={t['订单数']}, GMV={t['GMV（元）']}元{extra}")
    print(f"   合计: 订单数={grand_orders}, GMV={grand_gmv}元")

    return output_path


if __name__ == "__main__":
    month_input = sys.argv[1] if len(sys.argv) > 1 else datetime.now().strftime("%Y-%m")
    month_str = parse_month(month_input)
    run(month_str)
