#!/usr/bin/env python3
"""
小龙 sheet update: match phones via XXTEA encryption, fill H (UID), D (trial count), I (reg date), J (channel)
"""
import sys
sys.path.insert(0, '/root/.openclaw/workspace/scripts')
from phone_encrypt import encrypt_phone
import psycopg2
import json
import requests
import time

FEISHU_TOKEN = "t-g10464c0UK5L67JVXSDDT3EWM4DPLSDY5C7R7NS6"
SPREADSHEET_TOKEN = "NoZqsFi47hIOHEt9j8WcfRtbnug"
SHEET_ID = "qJF4I"

# PostgreSQL connection
PG_CONFIG = {
    "host": "bj-postgres-16pob4sg.sql.tencentcdb.com",
    "port": 28591,
    "user": "ai_member",
    "password": "LdfjdjL83h3h3^$&**YGG*",
    "database": "vala_bi",
}

def get_sheet_data():
    """Read all data from the sheet"""
    url = f"https://open.feishu.cn/open-apis/sheets/v2/spreadsheets/{SPREADSHEET_TOKEN}/values/{SHEET_ID}!A3:J2502?valueRenderOption=ToString"
    headers = {"Authorization": f"Bearer {FEISHU_TOKEN}"}
    r = requests.get(url, headers=headers)
    data = r.json()
    values = data.get("data", {}).get("valueRange", {}).get("values", [])
    return values

def write_cell_range(range_str, values_list):
    """Write values to a range"""
    url = f"https://open.feishu.cn/open-apis/sheets/v2/spreadsheets/{SPREADSHEET_TOKEN}/values"
    headers = {
        "Authorization": f"Bearer {FEISHU_TOKEN}",
        "Content-Type": "application/json"
    }
    body = {
        "valueRange": {
            "range": f"{SHEET_ID}!{range_str}",
            "values": values_list
        }
    }
    r = requests.put(url, headers=headers, json=body)
    resp = r.json()
    if resp.get("code") != 0:
        print(f"  ERROR writing {range_str}: {resp}")
    else:
        print(f"  OK: {range_str}")
    return resp

def main():
    # Step 1: Read sheet data
    print("Reading sheet data...")
    rows = get_sheet_data()
    print(f"Got {len(rows)} rows")

    # Parse rows with phones
    phone_rows = []  # (row_num, phone, existing_uid, existing_d, existing_i, existing_j)
    for i, row in enumerate(rows):
        row_num = i + 3
        if len(row) < 5:
            continue
        phone = row[4].strip() if row[4] else ""
        if not phone or len(phone) != 11 or not phone.isdigit():
            continue
        uid = row[7].strip() if len(row) > 7 and row[7] else ""
        d_val = row[3] if len(row) > 3 and row[3] else ""
        i_val = row[8].strip() if len(row) > 8 and row[8] else ""
        j_val = row[9].strip() if len(row) > 9 and row[9] else ""
        phone_rows.append((row_num, phone, uid, d_val, i_val, j_val))

    print(f"Found {len(phone_rows)} rows with valid phones")

    # Separate into two groups:
    # Group A: H is "未注册" or empty → need to match and fill H, D, I, J
    # Group B: H has valid numeric UID but D is empty → need to fill D, I, J
    group_a = []  # need phone matching
    group_b = []  # already have UID, need D/I/J

    for row_num, phone, uid, d_val, i_val, j_val in phone_rows:
        if not uid or uid == "未注册":
            group_a.append((row_num, phone, uid, d_val, i_val, j_val))
        elif uid.isdigit():
            if not d_val or not i_val or not j_val:
                group_b.append((row_num, phone, uid, d_val, i_val, j_val))

    print(f"Group A (need phone match): {len(group_a)}")
    print(f"Group B (have UID, need D/I/J): {len(group_b)}")

    # Step 2: Encrypt all phones and query PostgreSQL
    all_phones_a = [p[1] for p in group_a]
    all_uids_b = [p[2] for p in group_b]

    conn = psycopg2.connect(**PG_CONFIG)
    cur = conn.cursor()

    # Map: encrypted_phone -> account_id
    enc_to_uid = {}
    if all_phones_a:
        enc_list = [encrypt_phone(p) for p in all_phones_a]
        phone_to_enc = dict(zip(all_phones_a, enc_list))

        # Query accounts
        placeholders = ",".join(["%s"] * len(enc_list))
        cur.execute(
            f"SELECT id, tel_encrypt FROM bi_vala_app_account WHERE tel_encrypt IN ({placeholders}) AND status=1 AND deleted_at IS NULL",
            enc_list
        )
        for row in cur.fetchall():
            uid, tel_enc = row
            enc_to_uid[tel_enc] = uid

        print(f"Matched {len(enc_to_uid)} phones in bi_vala_app_account")

        # For matched accounts, get registration date and download channel
        matched_uids = list(enc_to_uid.values())
        uid_to_info = {}
        if matched_uids:
            placeholders2 = ",".join(["%s"] * len(matched_uids))
            cur.execute(
                f"SELECT id, created_at::date, download_channel FROM bi_vala_app_account WHERE id IN ({placeholders2}) AND status=1 AND deleted_at IS NULL",
                matched_uids
            )
            for row in cur.fetchall():
                uid_to_info[row[0]] = (str(row[1]) if row[1] else "", row[2] or "")

    # Step 3: Get trial lesson counts for ALL accounts (both groups)
    all_uids = list(set(
        [enc_to_uid[encrypt_phone(p[1])] for p in group_a if encrypt_phone(p[1]) in enc_to_uid] +
        [p[2] for p in group_b]
    ))
    
    uid_to_trial_count = {}
    if all_uids:
        placeholders3 = ",".join(["%s"] * len(all_uids))
        cur.execute(
            f"SELECT account_id, COUNT(*) FROM bi_user_course_detail WHERE account_id IN ({placeholders3}) AND expire_time IS NULL AND deleted_at IS NULL GROUP BY account_id",
            all_uids
        )
        for row in cur.fetchall():
            uid_to_trial_count[row[0]] = row[1]

    cur.close()
    conn.close()

    print(f"Trial counts: {len(uid_to_trial_count)} accounts have trials")
    print(f"Account info: {len(uid_to_info)} accounts have reg date/channel")

    # Step 4: Build write batches
    # For Group A: write H (UID), D (trial count), I (reg date), J (channel)
    # For Group B: write D (trial count), I (reg date), J (channel)

    writes_h = []  # (row, value)
    writes_d = []  # (row, value)
    writes_i = []  # (row, value)
    writes_j = []  # (row, value)

    matched_count = 0

    for row_num, phone, old_uid, old_d, old_i, old_j in group_a:
        enc = encrypt_phone(phone)
        uid = enc_to_uid.get(enc)
        if uid:
            matched_count += 1
            uid_str = str(uid)
            # H column: UID
            writes_h.append((row_num, uid_str))
            # D column: trial count
            trial = uid_to_trial_count.get(uid, 0)
            writes_d.append((row_num, str(trial) if trial else ""))
            # I column: reg date
            info = uid_to_info.get(uid, ("", ""))
            writes_i.append((row_num, info[0]))
            # J column: channel
            writes_j.append((row_num, info[1]))
        # If not matched, skip - don't write anything

    for row_num, phone, uid_str, old_d, old_i, old_j in group_b:
        uid = int(uid_str)
        # D column: trial count
        trial = uid_to_trial_count.get(uid, 0)
        writes_d.append((row_num, str(trial) if trial else ""))
        # I column: reg date
        info = uid_to_info.get(uid, ("", ""))
        writes_i.append((row_num, info[0]))
        # J column: channel
        writes_j.append((row_num, info[1]))

    print(f"\nMatched {matched_count} new phones")
    print(f"Total H writes: {len(writes_h)}")
    print(f"Total D writes: {len(writes_d)}")
    print(f"Total I writes: {len(writes_i)}")
    print(f"Total J writes: {len(writes_j)}")

    # Step 5: Write in batches (consecutive rows per column)
    def write_batch(writes, col_letter):
        if not writes:
            return
        # Sort by row number
        writes.sort(key=lambda x: x[0])
        # Group consecutive rows
        i = 0
        while i < len(writes):
            start = writes[i][0]
            vals = []
            j = i
            while j < len(writes) and writes[j][0] == start + (j - i):
                vals.append([writes[j][1]])
                j += 1
            end = start + len(vals) - 1
            range_str = f"{col_letter}{start}:{col_letter}{end}"
            write_cell_range(range_str, vals)
            i = j
            time.sleep(0.05)

    print("\nWriting H column...")
    write_batch(writes_h, "H")

    print("\nWriting D column...")
    write_batch(writes_d, "D")

    print("\nWriting I column...")
    write_batch(writes_i, "I")

    print("\nWriting J column...")
    write_batch(writes_j, "J")

    print("\n=== SUMMARY ===")
    print(f"Phones matched: {matched_count}")
    print(f"H (UID) written: {len(writes_h)}")
    print(f"D (trial count) written: {len(writes_d)}")
    print(f"I (reg date) written: {len(writes_i)}")
    print(f"J (channel) written: {len(writes_j)}")

if __name__ == "__main__":
    main()