#!/usr/bin/env python3
"""
批量行课记录查询 - 查用户当前进度、最近行课、学习时长、付费状态

输入: /tmp/pending_rows.json (从飞书表格拉取的待查询行)
输出: /tmp/results.jsonl (每行一个JSON, 包含 row + result 字段)
"""

import json
import os
import re
import sys
from collections import defaultdict

try:
    import psycopg2
    import psycopg2.extras
except ImportError:
    print("ERROR: 需要 psycopg2")
    sys.exit(1)

DB_HOST = "bj-postgres-16pob4sg.sql.tencentcdb.com"
DB_PORT = 28591
DB_USER = "ai_member"
SCRIPTS_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, SCRIPTS_DIR)
from phone_encrypt import encrypt_phone

COURSE_LEVEL_MAP = {"A1": "L1", "A2": "L2"}


def get_password():
    pw = os.environ.get("PG_ONLINE_PASSWORD", "")
    if pw:
        return pw
    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="):
                    return line.strip().split("=", 1)[1].strip("'\"")
    raise RuntimeError("PG_ONLINE_PASSWORD not found")


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


def batch_query_in(cur, sql_template, params, chunk_size=500):
    """分块执行 IN 查询，将 sql_template 中的第一个 %%s 替换为占位符"""
    results = []
    for i in range(0, len(params), chunk_size):
        chunk = params[i:i + chunk_size]
        placeholders = ",".join(["%s"] * len(chunk))
        sql = sql_template % placeholders
        cur.execute(sql, chunk)
        results.extend(cur.fetchall())
    return results


def match_phones(phones, conn):
    if not phones:
        return {}, [], []

    # 清理手机号：移除空格、小数点等
    def clean_phone(p):
        p = str(p).strip()
        # 处理浮点数格式: 13003632311.0 -> 13003632311
        if '.' in p:
            parts = p.split('.')
            if parts[1] in ('0', '00'):
                p = parts[0]
        # 验证是11位手机号
        if re.match(r'^1\d{10}$', p):
            return p
        return None

    cleaned = [(clean_phone(p), p) for p in phones]
    valid_phones = [(cp, op) for cp, op in cleaned if cp]
    if not valid_phones:
        return {}, [], []

    encrypt_to_phones = {}
    for cp, op in valid_phones:
        enc = encrypt_phone(cp)
        encrypt_to_phones.setdefault(enc, []).append(op)

    enc_list = list(encrypt_to_phones.keys())
    cur = conn.cursor()
    placeholders = ",".join(["%s"] * len(enc_list))
    cur.execute(f"""
        SELECT id AS account_id, tel_encrypt
        FROM bi_vala_app_account
        WHERE tel_encrypt IN ({placeholders}) AND status = 1 AND deleted_at IS NULL
    """, enc_list)
    rows = cur.fetchall()
    cur.close()

    # tel_encrypt → account_id
    enc_to_account = {}
    for aid, enc in rows:
        if enc not in enc_to_account:
            enc_to_account[enc] = aid

    # original phone → account_id
    phone_to_account = {}
    for op in phones:
        cp = clean_phone(op)
        if cp:
            enc = encrypt_phone(cp)
            if enc in enc_to_account:
                phone_to_account[op] = enc_to_account[enc]

    account_ids = list(set(phone_to_account.values()))
    unmatched = [p for p in phones if p not in phone_to_account]
    return phone_to_account, account_ids, unmatched


def query_course_details(cur, account_ids):
    if not account_ids:
        return {}
    rows = batch_query_in(cur, """
        SELECT account_id, course_level,
               CASE WHEN expire_time IS NULL THEN '体验课' ELSE '正式课' END AS course_type
        FROM bi_user_course_detail
        WHERE account_id IN (%s) AND deleted_at IS NULL
    """, account_ids)
    result = {}
    for aid, cl, ct in rows:
        if aid not in result:
            result[int(aid)] = (cl, ct)
    return result


def query_chapter_play_records(cur, char_ids):
    """查询每个角色最近一次行课记录和总学习时间"""
    if not char_ids:
        return {}, {}

    # 构建 UNION ALL SQL - 使用 CTE 方式先创建 ID 列表
    # 创建临时ID数组
    cur.execute("CREATE TEMP TABLE IF NOT EXISTS _tmp_char_ids (id int PRIMARY KEY) ON COMMIT DROP")
    cur.execute("DELETE FROM _tmp_char_ids")
    psycopg2.extras.execute_values(cur,
        "INSERT INTO _tmp_char_ids (id) VALUES %s",
        [(cid,) for cid in char_ids]
    )

    # 查询每个角色的最近行课记录
    union_parts = []
    for i in range(SHARD_COUNT):
        union_parts.append(
            f"SELECT cpr.user_id, cpr.chapter_id, cpr.chapter_unique_id, cpr.created_at, cpr.play_status "
            f"FROM bi_user_chapter_play_record_{i} cpr "
            f"JOIN _tmp_char_ids t ON cpr.user_id = t.id "
            f"WHERE cpr.deleted_at IS NULL"
        )

    all_plays_sql = " UNION ALL ".join(union_parts)

    cur.execute(f"""
        WITH all_plays AS ({all_plays_sql})
        , latest AS (
            SELECT DISTINCT ON (user_id) user_id, chapter_id, chapter_unique_id, created_at AS latest_date
            FROM all_plays
            ORDER BY user_id, created_at DESC
        )
        SELECT user_id, chapter_id, chapter_unique_id, latest_date FROM latest
    """)

    char_latest = {}
    chapter_ids = set()
    for uid, ch_id, cu_id, ldate in cur.fetchall():
        char_latest[uid] = (ch_id, ldate)
        chapter_ids.add(ch_id)

    # 查询每个角色所有 play_status=1 的 chapter_unique_id
    cur.execute(f"""
        WITH all_plays AS ({all_plays_sql})
        SELECT user_id, chapter_unique_id
        FROM all_plays
        WHERE play_status = 1
    """)

    char_cuids = defaultdict(set)
    all_cuids = set()
    for uid, cuid in cur.fetchall():
        char_cuids[uid].add(cuid)
        all_cuids.add(cuid)

    # 查询组件耗时
    char_total_ms = defaultdict(int)
    if all_cuids:
        cuid_list = list(all_cuids)

        # 创建临时表存 cuids
        cur.execute("CREATE TEMP TABLE IF NOT EXISTS _tmp_cuids (cuid text PRIMARY KEY) ON COMMIT DROP")
        cur.execute("DELETE FROM _tmp_cuids")
        psycopg2.extras.execute_values(cur,
            "INSERT INTO _tmp_cuids (cuid) VALUES %s",
            [(c,) for c in cuid_list]
        )

        comp_parts = []
        for i in range(SHARD_COUNT):
            comp_parts.append(
                f"SELECT comp.chapter_unique_id, SUM(comp.interval_time) AS total_ms "
                f"FROM bi_user_component_play_record_{i} comp "
                f"JOIN _tmp_cuids t ON comp.chapter_unique_id = t.cuid "
                f"GROUP BY comp.chapter_unique_id"
            )
        comp_union = " UNION ALL ".join(comp_parts)

        cur.execute(f"""
            SELECT chapter_unique_id, SUM(total_ms) AS total_ms
            FROM ({comp_union}) t2
            GROUP BY chapter_unique_id
        """)

        cuid_time = {}
        for cuid, tms in cur.fetchall():
            cuid_time[cuid] = (tms or 0)

        for uid, cuids in char_cuids.items():
            total = sum(cuid_time.get(c, 0) for c in cuids)
            char_total_ms[uid] = total

    return char_latest, char_total_ms


def query_chapter_course_map(cur, chapter_ids):
    if not chapter_ids:
        return {}
    rows = batch_query_in(cur, """
        SELECT id, course_level, course_season, course_unit, course_lesson
        FROM bi_level_unit_lesson WHERE id IN (%s)
    """, list(chapter_ids))
    return {r[0]: (r[1], r[2], r[3], r[4]) for r in rows}


def query_payment_status(cur, account_ids):
    if not account_ids:
        return {}
    rows = batch_query_in(cur, """
        SELECT account_id, MAX(CASE WHEN order_status IN (3,4) THEN 1 ELSE 0 END) AS has_order
        FROM bi_vala_order WHERE account_id IN (%s) AND deleted_at IS NULL
        GROUP BY account_id
    """, account_ids)
    status = {int(aid): "已付费" if ho == 1 else "未付费" for aid, ho in rows}
    for aid in account_ids:
        if int(aid) not in status:
            status[int(aid)] = "未付费"
    return status


def query_characters(cur, account_ids):
    if not account_ids:
        return {}
    rows = batch_query_in(cur, """
        SELECT account_id, id AS character_id
        FROM bi_vala_app_character
        WHERE account_id IN (%s) AND nickname IS NOT NULL AND nickname != '' AND deleted_at IS NULL
    """, account_ids)
    result = defaultdict(list)
    for aid, cid in rows:
        result[int(aid)].append(int(cid))
    return result


def main():
    print("加载待查询行...")
    with open("/tmp/pending_rows.json") as f:
        pending_rows = json.load(f)
    print(f"共 {len(pending_rows)} 条待查询")

    with_id = []
    without_id = []
    for row in pending_rows:
        uid_str = (row.get("user_id") or "").strip()
        if uid_str:
            try:
                row["account_id"] = int(float(uid_str))
                with_id.append(row)
            except (ValueError, TypeError):
                phone = (row.get("phone") or "").strip()
                if phone:
                    row["phone_clean"] = phone
                    without_id.append(row)
        else:
            phone = (row.get("phone") or "").strip()
            if phone:
                row["phone_clean"] = phone
                without_id.append(row)

    print(f"有用户ID: {len(with_id)}, 需手机号匹配: {len(without_id)}")

    conn = get_conn()
    cur = conn.cursor()

    # Step 1: 手机号匹配
    phone_to_account = {}
    if without_id:
        phones = [r["phone_clean"] for r in without_id]
        phone_to_account, matched_ids, unmatched = match_phones(phones, conn)
        print(f"手机号匹配: {len(phone_to_account)} 成功, {len(unmatched)} 未匹配")

        for row in without_id:
            p = row["phone_clean"]
            if p in phone_to_account:
                row["account_id"] = int(phone_to_account[p])
                with_id.append(row)

    print(f"总计可查询用户: {len(with_id)}")

    all_account_ids = list(set(r["account_id"] for r in with_id))
    print(f"去重 account_id: {len(all_account_ids)}")

    if not all_account_ids:
        print("没有可查询的用户，退出")
        conn.close()
        return

    # Step 2: 查询课程分配
    print("查询课程分配...")
    course_details = query_course_details(cur, all_account_ids)
    print(f"  有课程: {len(course_details)}")

    # Step 3: 查询角色
    print("查询角色...")
    account_chars = query_characters(cur, all_account_ids)
    all_char_ids = []
    for chars in account_chars.values():
        all_char_ids.extend(chars)
    print(f"  角色数: {len(all_char_ids)}")

    # Step 4: 查询行课记录
    print("查询行课记录 (可能需要几分钟)...")
    char_latest, char_total_ms = query_chapter_play_records(cur, all_char_ids)
    print(f"  有行课记录: {len(char_latest)}")

    # Step 5: chapter 映射
    all_chapter_ids = set()
    for ch_id, _ in char_latest.values():
        all_chapter_ids.add(ch_id)
    print("查询课程映射...")
    course_chapter_map = query_chapter_course_map(cur, all_chapter_ids)
    print(f"  chapter: {len(course_chapter_map)}")

    # Step 6: 付费状态
    print("查询付费状态...")
    pay_status = query_payment_status(cur, all_account_ids)
    print(f"  付费: {sum(1 for v in pay_status.values() if v == '已付费')}")

    conn.close()

    # Step 7: 组装结果
    print("组装结果...")
    results = []

    # 按 row 去重，同一用户多行合并
    row_aid_map = {}
    for row in with_id:
        rn = row["row"]
        aid = row["account_id"]
        sales = row.get("sales", "").strip()
        if rn not in row_aid_map:
            row_aid_map[rn] = (aid, sales)

    for row in with_id:
        rn = row["row"]
        aid = row["account_id"]
        sales = row.get("sales", "").strip()

        # 课程描述
        cd = course_details.get(aid)
        if cd:
            cl, ct = cd
            level_display = COURSE_LEVEL_MAP.get(cl, cl)
            course_desc = f"{level_display}{ct}"
        else:
            course_desc = "暂无课程"

        # 最近行课和当前章节
        chars = account_chars.get(aid, [])
        latest_date = None
        chapter_pos = ""
        for cid in chars:
            if cid in char_latest:
                ch_id, ldate = char_latest[cid]
                if ldate:
                    if latest_date is None or ldate > latest_date:
                        latest_date = ldate
                    if ch_id in course_chapter_map:
                        lvl, season, unit, lesson = course_chapter_map[ch_id]
                        chapter_pos = f"{lvl}-{season}-{unit}-{lesson}"

        if not chapter_pos:
            chapter_pos = "无记录"
        date_str = latest_date.strftime("%Y-%m-%d") if latest_date else "无记录"

        # 学习时长
        total_min = 0
        for cid in chars:
            total_min += (char_total_ms.get(cid, 0) or 0)
        total_min = round(total_min / 1000 / 60)

        # 付费状态
        pay = pay_status.get(aid, "未付费")

        result_text = f"销售:{sales} | 用户:{aid} | 当前:{course_desc}-{chapter_pos} | 最近行课:{date_str} | 学习{total_min}min | {pay}"

        results.append({
            "row": rn,
            "result": result_text,
            "aid": aid
        })

    # 未匹配的
    for row in without_id:
        p = row.get("phone_clean", "")
        if p not in phone_to_account:
            results.append({
                "row": row["row"],
                "result": f"未匹配到用户",
                "aid": 0
            })

    # 输出
    with open("/tmp/results.jsonl", "w") as f:
        for r in results:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")

    print(f"\n✅ 查询完成: {len(results)} 条结果")
    unrecorded = sum(1 for r in results if "无记录" in r.get("result","") and "未匹配" not in r.get("result",""))
    unmatched_count = sum(1 for r in results if "未匹配" in r.get("result",""))
    paid_count = sum(1 for r in results if "已付费" in r.get("result",""))
    print(f"   有记录: {len(results) - unrecorded - unmatched_count}")
    print(f"   无记录: {unrecorded}")
    print(f"   未匹配: {unmatched_count}")
    print(f"   已付费: {paid_count}")


if __name__ == "__main__":
    main()
