#!/usr/bin/env python3
"""
导出「河南钟馗科技有限公司」学历信息高级版(IVYZ4Y27)成功调用记录为 CSV。
解密 request_params，仅输出脱敏后的姓名、身份证。

依赖: pip install pycryptodome psycopg2-binary

用法 A - 直连数据库（在可访问生产/测试库的环境执行）:
  python export_ivyz4y27_calls.py \\
    --from-db \\
    --db-host tyapi-postgres-prod --db-name tyapi --db-user postgres --db-password <pwd> \\
    -o ivyz4y27_zhongkui_calls.csv

  或使用环境变量:
  set PGHOST=... PGDATABASE=... PGUSER=... PGPASSWORD=...
  python export_ivyz4y27_calls.py --from-db -o ivyz4y27_zhongkui_calls.csv

用法 B - 先 SQL 导出 JSON，再本地转 CSV:
  python export_ivyz4y27_calls.py -i ivyz4y27_zhongkui_raw.json -o ivyz4y27_zhongkui_calls.csv
"""

from __future__ import annotations

import argparse
import base64
import csv
import json
import os
import sys
from pathlib import Path
from typing import Any

try:
    from Crypto.Cipher import AES
except ImportError:
    print("请先安装: pip install pycryptodome", file=sys.stderr)
    sys.exit(1)

COMPANY_NAME = "河南钟馗科技有限公司"
PRODUCT_CODE = "IVYZ4Y27"

SQL = """
SELECT
    ac.id,
    ac.transaction_id,
    ac.user_id,
    ei.company_name,
    p.code          AS product_code,
    p.name          AS product_name,
    ac.request_params,
    au.secret_key,
    ac.status,
    ac.cost,
    ac.start_at,
    ac.created_at
FROM api_calls ac
INNER JOIN products p ON ac.product_id = p.id
INNER JOIN api_users au ON ac.user_id = au.user_id
INNER JOIN enterprise_infos ei ON ac.user_id = ei.user_id
WHERE ei.company_name = %s
  AND ei.deleted_at IS NULL
  AND p.code = %s
  AND ac.status = 'success'
ORDER BY ac.created_at DESC
"""


def aes_decrypt(cipher_b64: str, key_hex: str) -> dict[str, Any]:
    key = bytes.fromhex(key_hex.strip())
    raw = base64.b64decode(cipher_b64.strip())
    iv, ciphertext = raw[:16], raw[16:]
    plain = AES.new(key, AES.MODE_CBC, iv).decrypt(ciphertext)
    pad = plain[-1]
    plain = plain[:-pad]
    return json.loads(plain.decode("utf-8"))


def mask_name(name: str) -> str:
    """与 dwbg8b4d maskName 一致：张*三 / 李*"""
    name = (name or "").strip()
    if not name:
        return ""
    chars = list(name)
    if len(chars) <= 1:
        return name
    if len(chars) == 2:
        return chars[0] + "*"
    return chars[0] + "*" + chars[-1]


def mask_id_card(id_card: str) -> str:
    """与 query_whitelist_helpers MaskIDCard 一致：350681********0611"""
    id_card = (id_card or "").strip().upper()
    if len(id_card) <= 10:
        return id_card
    return id_card[:6] + "********" + id_card[-4:]


def extract_name_id_card(params: dict[str, Any]) -> tuple[str, str]:
    name = params.get("name") or params.get("Name") or ""
    id_card = params.get("id_card") or params.get("idCard") or params.get("IDCard") or ""
    if isinstance(name, str):
        name = name.strip()
    else:
        name = str(name)
    if isinstance(id_card, str):
        id_card = id_card.strip().upper()
    else:
        id_card = str(id_card).strip().upper()
    return name, id_card


def load_rows_from_json(path: Path) -> list[dict]:
    text = path.read_text(encoding="utf-8-sig").strip()
    data = json.loads(text)
    if isinstance(data, list):
        return data
    if isinstance(data, dict):
        for key in ("records", "rows", "data", "items"):
            if key in data and isinstance(data[key], list):
                return data[key]
    raise ValueError("JSON 应为数组，或 export SQL 的 json_agg 结果")


def fetch_rows_from_db(args: argparse.Namespace) -> list[dict]:
    try:
        import psycopg2
        import psycopg2.extras
    except ImportError:
        print("直连数据库需要: pip install psycopg2-binary", file=sys.stderr)
        sys.exit(1)

    host = args.db_host or os.getenv("PGHOST", "localhost")
    port = args.db_port or os.getenv("PGPORT", "5432")
    dbname = args.db_name or os.getenv("PGDATABASE", "tyapi_dev")
    user = args.db_user or os.getenv("PGUSER", "postgres")
    password = args.db_password or os.getenv("PGPASSWORD", "")

    conn = psycopg2.connect(
        host=host,
        port=port,
        dbname=dbname,
        user=user,
        password=password,
    )
    try:
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            cur.execute(SQL, (COMPANY_NAME, PRODUCT_CODE))
            rows = cur.fetchall()
            return [dict(r) for r in rows]
    finally:
        conn.close()


def row_to_csv_record(row: dict) -> dict[str, str]:
    encrypted = row.get("request_params") or ""
    secret_key = row.get("secret_key") or ""
    name_masked = ""
    id_card_masked = ""
    decrypt_error = ""

    if encrypted and secret_key:
        try:
            params = aes_decrypt(encrypted, secret_key)
            name, id_card = extract_name_id_card(params)
            name_masked = mask_name(name)
            id_card_masked = mask_id_card(id_card)
        except Exception as exc:
            decrypt_error = str(exc)

    return {
        "company_name": str(row.get("company_name") or COMPANY_NAME),
        "product_code": str(row.get("product_code") or PRODUCT_CODE),
        "product_name": str(row.get("product_name") or ""),
        "api_call_id": str(row.get("api_call_id") or row.get("id") or ""),
        "transaction_id": str(row.get("transaction_id") or ""),
        "call_time": str(row.get("call_time") or row.get("created_at") or ""),
        "start_at": str(row.get("start_at") or ""),
        "cost": str(row.get("cost") or ""),
        "name_masked": name_masked,
        "id_card_masked": id_card_masked,
        "decrypt_error": decrypt_error,
    }


def write_csv(records: list[dict[str, str]], output: Path) -> None:
    fields = [
        "company_name",
        "product_code",
        "product_name",
        "api_call_id",
        "transaction_id",
        "call_time",
        "start_at",
        "cost",
        "name_masked",
        "id_card_masked",
        "decrypt_error",
    ]
    with output.open("w", encoding="utf-8-sig", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fields)
        writer.writeheader()
        writer.writerows(records)


def main() -> None:
    parser = argparse.ArgumentParser(description="导出 IVYZ4Y27 学历高级版调用记录 CSV")
    parser.add_argument("-i", "--input", help="SQL 导出的 JSON 文件")
    parser.add_argument("-o", "--output", default="ivyz4y27_zhongkui_calls.csv")
    parser.add_argument("--from-db", action="store_true", help="直连 PostgreSQL 查询")
    parser.add_argument("--db-host")
    parser.add_argument("--db-port", default="5432")
    parser.add_argument("--db-name")
    parser.add_argument("--db-user")
    parser.add_argument("--db-password")
    args = parser.parse_args()

    if args.from_db:
        rows = fetch_rows_from_db(args)
    elif args.input:
        rows = load_rows_from_json(Path(args.input))
    else:
        print("请指定 --from-db 或 -i <json文件>", file=sys.stderr)
        sys.exit(1)

    records = [row_to_csv_record(r) for r in rows]
    out = Path(args.output)
    write_csv(records, out)

    failed = sum(1 for r in records if r.get("decrypt_error"))
    print(f"企业: {COMPANY_NAME}")
    print(f"产品: {PRODUCT_CODE} (学历信息高级版)")
    print(f"成功记录: {len(records)} 条, 解密失败: {failed} 条")
    print(f"输出: {out.resolve()}")


if __name__ == "__main__":
    main()