# -*- coding: utf-8 -*-
"""
从本地 JSON 示例目录重建 billiards_ods.* 表，并导入样例数据。
用法：
  PYTHONPATH=. python -m etl_billiards.scripts.rebuild_ods_from_json [--dsn ...] [--json-dir ...] [--include ...] [--drop-schema-first]

依赖环境变量：
  PG_DSN         PostgreSQL 连接串（必填）
  PG_CONNECT_TIMEOUT 可选，秒，默认 10
  JSON_DOC_DIR   可选，JSON 目录，默认 export/test-json-doc
  ODS_INCLUDE_FILES 可选，逗号分隔文件名（不含 .json）
  ODS_DROP_SCHEMA_FIRST 可选，true/false，默认 true
"""
from __future__ import annotations

import argparse
import os
import re
import sys
import json
from pathlib import Path
from typing import Iterable, List, Tuple

import psycopg2
from psycopg2 import sql
from psycopg2.extras import Json, execute_values


DEFAULT_JSON_DIR = "export/test-json-doc"
SPECIAL_LIST_PATHS: dict[str, tuple[str, ...]] = {
    "assistant_accounts_master": ("data", "assistantInfos"),
    "assistant_cancellation_records": ("data", "abolitionAssistants"),
    "assistant_service_records": ("data", "orderAssistantDetails"),
    "goods_stock_movements": ("data", "queryDeliveryRecordsList"),
    "goods_stock_summary": ("data",),
    "group_buy_packages": ("data", "packageCouponList"),
    "group_buy_redemption_records": ("data", "siteTableUseDetailsList"),
    "member_balance_changes": ("data", "tenantMemberCardLogs"),
    "member_profiles": ("data", "tenantMemberInfos"),
    "member_stored_value_cards": ("data", "tenantMemberCards"),
    "recharge_settlements": ("data", "settleList"),
    "settlement_records": ("data", "settleList"),
    "site_tables_master": ("data", "siteTables"),
    "stock_goods_category_tree": ("data", "goodsCategoryList"),
    "store_goods_master": ("data", "orderGoodsList"),
    "store_goods_sales_records": ("data", "orderGoodsLedgers"),
    "table_fee_discount_records": ("data", "taiFeeAdjustInfos"),
    "table_fee_transactions": ("data", "siteTableUseDetailsList"),
    "tenant_goods_master": ("data", "tenantGoodsList"),
}


def sanitize_identifier(name: str) -> str:
    """将任意字符串转为可用的 SQL identifier（小写、非字母数字转下划线）。"""
    cleaned = re.sub(r"[^0-9a-zA-Z_]", "_", name.strip())
    if not cleaned:
        cleaned = "col"
    if cleaned[0].isdigit():
        cleaned = f"_{cleaned}"
    return cleaned.lower()


def _extract_list_via_path(node, path: tuple[str, ...]):
    cur = node
    for key in path:
        if isinstance(cur, dict):
            cur = cur.get(key)
        else:
            return []
    return cur if isinstance(cur, list) else []


def load_records(payload, list_path: tuple[str, ...] | None = None) -> list:
    """
    尝试从 JSON 结构中提取记录列表：
    - 直接是 list -> 返回
    - dict 中 data 是 list -> 返回
    - dict 中 data 是 dict，取第一个 list 字段
    - dict 中任意值是 list -> 返回
    - 其余情况，包装为单条记录
    """
    if list_path:
        if isinstance(payload, list):
            merged: list = []
            for item in payload:
                merged.extend(_extract_list_via_path(item, list_path))
            if merged:
                return merged
        elif isinstance(payload, dict):
            lst = _extract_list_via_path(payload, list_path)
            if lst:
                return lst

    if isinstance(payload, list):
        return payload
    if isinstance(payload, dict):
        data_node = payload.get("data")
        if isinstance(data_node, list):
            return data_node
        if isinstance(data_node, dict):
            for v in data_node.values():
                if isinstance(v, list):
                    return v
        for v in payload.values():
            if isinstance(v, list):
                return v
    return [payload]


def collect_columns(records: Iterable[dict]) -> List[str]:
    """汇总所有顶层键，作为表字段；仅处理 dict 记录。"""
    cols: set[str] = set()
    for rec in records:
        if isinstance(rec, dict):
            cols.update(rec.keys())
    return sorted(cols)


def create_table(cur, schema: str, table: str, columns: List[Tuple[str, str]]):
    """
    创建表：字段全部 jsonb，外加 source_file、record_index、payload、ingested_at。
    columns: [(col_name, original_key)]
    """
    fields = [sql.SQL("{} jsonb").format(sql.Identifier(col)) for col, _ in columns]
    constraint_name = f"uq_{table}_source_record"
    ddl = sql.SQL(
        "CREATE TABLE IF NOT EXISTS {schema}.{table} ("
        "source_file text,"
        "record_index integer,"
        "{cols},"
        "payload jsonb,"
        "ingested_at timestamptz default now(),"
        "CONSTRAINT {constraint} UNIQUE (source_file, record_index)"
        ");"
    ).format(
        schema=sql.Identifier(schema),
        table=sql.Identifier(table),
        cols=sql.SQL(",").join(fields),
        constraint=sql.Identifier(constraint_name),
    )
    cur.execute(ddl)


def insert_records(cur, schema: str, table: str, columns: List[Tuple[str, str]], records: list, source_file: str):
    """批量插入记录。"""
    col_idents = [sql.Identifier(col) for col, _ in columns]
    col_names = [col for col, _ in columns]
    orig_keys = [orig for _, orig in columns]
    all_cols = [sql.Identifier("source_file"), sql.Identifier("record_index")] + col_idents + [
        sql.Identifier("payload")
    ]

    rows = []
    for idx, rec in enumerate(records):
        if not isinstance(rec, dict):
            rec = {"value": rec}
        row_values = [source_file, idx]
        for key in orig_keys:
            row_values.append(Json(rec.get(key)))
        row_values.append(Json(rec))
        rows.append(row_values)

    insert_sql = sql.SQL("INSERT INTO {}.{} ({}) VALUES %s ON CONFLICT DO NOTHING").format(
        sql.Identifier(schema),
        sql.Identifier(table),
        sql.SQL(",").join(all_cols),
    )
    execute_values(cur, insert_sql, rows, page_size=500)


def rebuild(schema: str = "billiards_ods", data_dir: str | Path = DEFAULT_JSON_DIR):
    parser = argparse.ArgumentParser(description="重建 billiards_ods.* 表并导入 JSON 样例")
    parser.add_argument("--dsn", dest="dsn", help="PostgreSQL DSN（默认读取环境变量 PG_DSN）")
    parser.add_argument("--json-dir", dest="json_dir", help=f"JSON 目录，默认 {DEFAULT_JSON_DIR}")
    parser.add_argument(
        "--include",
        dest="include_files",
        help="限定导入的文件名（逗号分隔，不含 .json），默认全部",
    )
    parser.add_argument(
        "--drop-schema-first",
        dest="drop_schema_first",
        action="store_true",
        help="先删除并重建 schema（默认 true）",
    )
    parser.add_argument(
        "--no-drop-schema-first",
        dest="drop_schema_first",
        action="store_false",
        help="保留现有 schema，仅按冲突去重导入",
    )
    parser.set_defaults(drop_schema_first=None)
    args = parser.parse_args()

    dsn = args.dsn or os.environ.get("PG_DSN")
    if not dsn:
        print("缺少参数/环境变量 PG_DSN，无法连接数据库。")
        sys.exit(1)
    timeout = max(1, min(int(os.environ.get("PG_CONNECT_TIMEOUT", 10)), 60))
    env_drop = os.environ.get("ODS_DROP_SCHEMA_FIRST") or os.environ.get("DROP_SCHEMA_FIRST")
    drop_schema_first = (
        args.drop_schema_first
        if args.drop_schema_first is not None
        else str(env_drop or "true").lower() in ("1", "true", "yes")
    )
    include_files_env = args.include_files or os.environ.get("ODS_INCLUDE_FILES") or os.environ.get("INCLUDE_FILES")
    include_files = set()
    if include_files_env:
        include_files = {p.strip().lower() for p in include_files_env.split(",") if p.strip()}

    base_dir = Path(args.json_dir or data_dir or DEFAULT_JSON_DIR)
    if not base_dir.exists():
        print(f"JSON 目录不存在: {base_dir}")
        sys.exit(1)

    conn = psycopg2.connect(dsn, connect_timeout=timeout)
    conn.autocommit = False
    cur = conn.cursor()

    if drop_schema_first:
        print(f"Dropping schema {schema} ...")
        cur.execute(sql.SQL("DROP SCHEMA IF EXISTS {} CASCADE;").format(sql.Identifier(schema)))
        cur.execute(sql.SQL("CREATE SCHEMA {};").format(sql.Identifier(schema)))
    else:
        cur.execute(
            sql.SQL("SELECT schema_name FROM information_schema.schemata WHERE schema_name=%s"),
            (schema,),
        )
        if not cur.fetchone():
            cur.execute(sql.SQL("CREATE SCHEMA {};").format(sql.Identifier(schema)))

    json_files = sorted(base_dir.glob("*.json"))
    for path in json_files:
        stem_lower = path.stem.lower()
        if include_files and stem_lower not in include_files:
            continue

        print(f"Processing {path.name} ...")
        payload = json.loads(path.read_text(encoding="utf-8"))
        list_path = SPECIAL_LIST_PATHS.get(stem_lower)
        records = load_records(payload, list_path=list_path)
        columns_raw = collect_columns(records)
        columns = [(sanitize_identifier(c), c) for c in columns_raw]

        table_name = sanitize_identifier(path.stem)
        create_table(cur, schema, table_name, columns)
        if records:
            insert_records(cur, schema, table_name, columns, records, path.name)
        print(f"  -> rows: {len(records)}, columns: {len(columns)}")

    conn.commit()
    cur.close()
    conn.close()
    print("Rebuild done.")


if __name__ == "__main__":
    rebuild()
