#!/usr/bin/env python3
import argparse
import datetime as dt
import json
import os
import pathlib
import sys
from typing import Dict, List, Tuple

import pytds

BUSINESS_TABLES = [
    "jst_goods_profit_kpi_items_afs",
    "jst_sku",
    "fys_shopInfo",
    "jst_order_express_details",
    "fys_SkuDataLibrary_SpuConfig",
    "op_operManage",
    "op_TargetManagement",
    "fys_Ticketing_install",
    "fys_Ticketing_install_item",
    "fys_jstFreightCost",
    "fys_jstFreightCostDetail",
    "pc_kfChatLog",
]

ENV_FILE = "/root/.openclaw/credentials/company-sql.env"


def load_env_file(path: str) -> None:
    p = pathlib.Path(path)
    if not p.exists():
        return
    for raw in p.read_text(encoding="utf-8").splitlines():
        line = raw.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, val = line.split("=", 1)
        key = key.strip()
        val = val.strip().strip('"').strip("'")
        if key and key not in os.environ:
            os.environ[key] = val


def conn_params() -> Dict[str, object]:
    load_env_file(ENV_FILE)
    return {
        "server": os.environ.get("COMPANY_DB_HOST", "218.93.191.16"),
        "port": int(os.environ.get("COMPANY_DB_PORT", "1433")),
        "database": os.environ.get("COMPANY_DB_NAME", "JSDataMiddlePlatform"),
        "user": os.environ.get("COMPANY_DB_USER", "nanting"),
        "password": os.environ.get("COMPANY_DB_PASSWORD", ""),
        "login_timeout": 15,
        "timeout": 300,
    }


def get_conn():
    params = conn_params()
    if not params["password"]:
        raise RuntimeError("COMPANY_DB_PASSWORD is empty")
    return pytds.connect(**params)


def _fmt(value):
    if isinstance(value, dt.datetime):
        return value.isoformat(sep=" ")
    if isinstance(value, (dt.date, dt.time)):
        return value.isoformat()
    return value


def print_rows(columns: List[str], rows: List[Tuple[object, ...]], truncated: bool = False) -> None:
    print("COLUMNS:", ", ".join(columns))
    if not rows:
        print("ROWS: 0")
        return
    if truncated:
        print(f"ROWS_SHOWN: {len(rows)} (truncated)")
    else:
        print(f"ROWS: {len(rows)}")
    for row in rows:
        obj = {columns[i]: _fmt(row[i]) for i in range(len(columns))}
        print(json.dumps(obj, ensure_ascii=False, default=str))


def split_go_batches(sql_text: str) -> List[str]:
    batches: List[str] = []
    current: List[str] = []
    for line in sql_text.splitlines():
        if line.strip().upper() == "GO":
            block = "\n".join(current).strip()
            if block:
                batches.append(block)
            current = []
        else:
            current.append(line)
    block = "\n".join(current).strip()
    if block:
        batches.append(block)
    return batches


def fetch_with_limit(cursor, row_limit: int) -> Tuple[List[Tuple[object, ...]], bool]:
    if row_limit <= 0:
        return cursor.fetchall(), False

    rows: List[Tuple[object, ...]] = []
    truncated = False
    while len(rows) <= row_limit:
        batch = cursor.fetchmany(256)
        if not batch:
            break
        rows.extend(batch)
        if len(rows) > row_limit:
            rows = rows[:row_limit]
            truncated = True
            break
    return rows, truncated


def run_sql(sql_text: str, row_limit: int) -> int:
    batches = split_go_batches(sql_text)
    with get_conn() as conn:
        cur = conn.cursor()
        set_index = 0
        for batch in batches:
            cur.execute(batch)
            while True:
                if cur.description:
                    set_index += 1
                    cols = [c[0] for c in cur.description]
                    rows, truncated = fetch_with_limit(cur, row_limit)
                    print(f"RESULT_SET {set_index}")
                    print_rows(cols, rows, truncated=truncated)
                if not cur.nextset():
                    break
        cur.close()
    return set_index


def cmd_health(_args: argparse.Namespace) -> int:
    with get_conn() as conn:
        cur = conn.cursor()
        cur.execute("SELECT @@VERSION AS version, DB_NAME() AS db_name, GETDATE() AS server_time")
        cols = [c[0] for c in cur.description]
        row = cur.fetchone()
        data = {cols[i]: _fmt(row[i]) for i in range(len(cols))}
        print(json.dumps(data, ensure_ascii=False, default=str, indent=2))
        cur.close()
    return 0


def cmd_tables(_args: argparse.Namespace) -> int:
    names = ", ".join([f"'{t}'" for t in BUSINESS_TABLES])
    sql = f"""
    SELECT c.TABLE_NAME, COUNT(*) AS column_count
    FROM INFORMATION_SCHEMA.COLUMNS c
    WHERE c.TABLE_SCHEMA='dbo'
      AND c.TABLE_NAME IN ({names})
    GROUP BY c.TABLE_NAME
    ORDER BY c.TABLE_NAME
    """
    with get_conn() as conn:
        cur = conn.cursor()
        cur.execute(sql)
        rows = cur.fetchall()
        cols = [c[0] for c in cur.description]
        print_rows(cols, rows)
        cur.close()
    return 0


def cmd_columns(args: argparse.Namespace) -> int:
    table = args.table
    sql = f"""
    SELECT
      COLUMN_NAME,
      DATA_TYPE,
      CHARACTER_MAXIMUM_LENGTH,
      NUMERIC_PRECISION,
      NUMERIC_SCALE,
      IS_NULLABLE,
      ORDINAL_POSITION
    FROM INFORMATION_SCHEMA.COLUMNS
    WHERE TABLE_SCHEMA='dbo' AND TABLE_NAME='{table}'
    ORDER BY ORDINAL_POSITION
    """
    with get_conn() as conn:
        cur = conn.cursor()
        cur.execute(sql)
        rows = cur.fetchall()
        cols = [c[0] for c in cur.description]
        print_rows(cols, rows)
        cur.close()
    return 0


def _collect_schema() -> Dict[str, List[Tuple[object, ...]]]:
    names = ", ".join([f"'{t}'" for t in BUSINESS_TABLES])
    sql = f"""
    SELECT
      TABLE_NAME,
      COLUMN_NAME,
      DATA_TYPE,
      CHARACTER_MAXIMUM_LENGTH,
      NUMERIC_PRECISION,
      NUMERIC_SCALE,
      IS_NULLABLE,
      ORDINAL_POSITION
    FROM INFORMATION_SCHEMA.COLUMNS
    WHERE TABLE_SCHEMA='dbo'
      AND TABLE_NAME IN ({names})
    ORDER BY TABLE_NAME, ORDINAL_POSITION
    """
    with get_conn() as conn:
        cur = conn.cursor()
        cur.execute(sql)
        rows = cur.fetchall()
        cur.close()
    by_table: Dict[str, List[Tuple[object, ...]]] = {t: [] for t in BUSINESS_TABLES}
    for row in rows:
        by_table[str(row[0])].append(row)
    return by_table


def cmd_catalog(args: argparse.Namespace) -> int:
    by_table = _collect_schema()
    now = dt.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    out = pathlib.Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)

    lines: List[str] = []
    lines.append("# 公司销售/售后数据库字段与维度目录")
    lines.append("")
    lines.append(f"> 自动生成时间: {now}")
    lines.append("> 数据源: SQLServer / JSDataMiddlePlatform / dbo")
    lines.append("> 适用范围: 公司销售数据、售后数据、店铺运营、物流运费、安装工单、客服问题")
    lines.append("")
    lines.append("## 一、强制路由规则（给AI）")
    lines.append("")
    lines.append("- 凡是涉及 `销售/业绩/售后/退款/费比/店铺业绩` 的需求，必须走数据库查询，禁止走旧业绩网页。")
    lines.append("- 先查字段目录再写 SQL：`/root/.openclaw/workspace/outputs/company-sales-db-catalog.md`")
    lines.append("- 优先复用周报模板：`/root/.openclaw/workspace/sqlserver/last_week_afs_report.sql`")
    lines.append("- 输出结论必须标注：`时间范围 + 使用表 + 关键字段 + 过滤条件`。")
    lines.append("")
    lines.append("## 二、可查询业务维度（语义层）")
    lines.append("")
    dimension_rows = [
        ("订单维度", "订单号、子订单号、内部订单ID、订单状态、售后状态", "jst_order_express_details, jst_goods_profit_kpi_items_afs"),
        ("商品维度", "SKU、SPU、货号、商品名称、规格、类目、品牌", "jst_sku, op_TargetManagement, op_operManage, fys_SkuDataLibrary_SpuConfig"),
        ("店铺维度", "店铺名称、平台、店铺编码、负责人、渠道类型", "fys_shopInfo, jst_order_express_details, jst_goods_profit_kpi_items_afs"),
        ("销售金额维度", "商品金额、付款金额、实发金额、成本、毛利相关费用", "jst_goods_profit_kpi_items_afs, jst_sku"),
        ("售后维度", "售后类型、售后原因、退款/退货金额、售后标签", "jst_goods_profit_kpi_items_afs, pc_kfChatLog"),
        ("物流维度", "快递公司、快递单号、物流状态、运费、重量、收货地区", "jst_order_express_details, fys_jstFreightCost, fys_jstFreightCostDetail"),
        ("安装工单维度", "工单状态、安装费用、结算金额、服务商", "fys_Ticketing_install, fys_Ticketing_install_item"),
        ("人员维度", "运营、设计师、采购、美工、店长、客服", "op_TargetManagement, op_operManage, fys_shopInfo, fys_SkuDataLibrary_SpuConfig"),
        ("时间维度", "创建时间、支付时间、发货时间、签收时间、订单完成时间", "全表通用时间字段（createdAt/createTime/pay_time/send_time 等）"),
        ("客服问题维度", "聊天主次原因、问题内容、订单关联", "pc_kfChatLog"),
    ]
    lines.append("| 维度 | 可查询内容 | 主要来源表 |")
    lines.append("|---|---|---|")
    for row in dimension_rows:
        lines.append(f"| {row[0]} | {row[1]} | {row[2]} |")
    lines.append("")
    lines.append("## 三、核心关联键（Join Key）")
    lines.append("")
    lines.append("- `order_id`：`so_id / undefined2 / orderNo / OrderNum / Sourceno`")
    lines.append("- `inner_order_id`：`inner_order_id / undefined1 / o_id`")
    lines.append("- `sku_id`：`sku_id / undefined10 / barCode`")
    lines.append("- `spu`：`spu / other_1 / undefined96`")
    lines.append("- `shop_name`：`shop_name / shopName / undefined8 / shopname`")
    lines.append("")
    lines.append("## 四、全表字段清单（真实表头）")
    lines.append("")

    total_cols = 0
    for table in BUSINESS_TABLES:
        rows = by_table.get(table, [])
        total_cols += len(rows)
        lines.append(f"### {table}（{len(rows)} 列）")
        lines.append("")
        lines.append("| 序号 | 字段名 | 类型 | 长度/精度 | 可空 |")
        lines.append("|---:|---|---|---|---|")
        for r in rows:
            _, col, dtype, clen, prec, scale, nullable, ord_pos = r
            if clen is not None:
                length = str(clen)
            elif prec is not None:
                if scale is not None:
                    length = f"{prec},{scale}"
                else:
                    length = str(prec)
            else:
                length = "-"
            lines.append(f"| {ord_pos} | {col} | {dtype} | {length} | {nullable} |")
        lines.append("")

    lines.append("## 五、统计")
    lines.append("")
    lines.append(f"- 业务表数量：{len(BUSINESS_TABLES)}")
    lines.append(f"- 字段总数：{total_cols}")
    lines.append("- 建议：高频分析优先在订单售后主表 `jst_goods_profit_kpi_items_afs` 先聚合，再按需关联其他表，避免超时。")

    out.write_text("\n".join(lines), encoding="utf-8")
    print(str(out))
    return 0


def cmd_query(args: argparse.Namespace) -> int:
    sql_text = None

    if args.file:
        sql_text = pathlib.Path(args.file).read_text(encoding="utf-8")
    elif args.sql:
        sql_text = args.sql
    elif args.sql_positional:
        sql_text = " ".join(args.sql_positional).strip()

    if not sql_text:
        raise SystemExit("query requires --sql / --file / or inline SQL text")

    set_count = run_sql(sql_text, args.limit)
    print(f"TOTAL_RESULT_SETS: {set_count}")
    return 0


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(prog="company-db", description="Company SQLServer helper")
    sub = p.add_subparsers(dest="cmd", required=True)

    sub.add_parser("health")
    sub.add_parser("tables")

    p_cols = sub.add_parser("columns")
    p_cols.add_argument("table")

    p_query = sub.add_parser("query")
    p_query.add_argument("--sql", help="SQL text")
    p_query.add_argument("--file", help="SQL file path")
    p_query.add_argument("--limit", type=int, default=200, help="row limit via SET ROWCOUNT")
    p_query.add_argument("sql_positional", nargs="*", help="inline SQL text, e.g. company-db query SELECT 1")

    p_run = sub.add_parser("run-file")
    p_run.add_argument("file", help="SQL file path")
    p_run.add_argument("--limit", type=int, default=0, help="max rows shown per result set (0 means no limit)")

    p_cat = sub.add_parser("catalog")
    p_cat.add_argument("--out", default="/root/.openclaw/workspace/outputs/company-sales-db-catalog.md")

    return p


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()

    if args.cmd == "health":
        return cmd_health(args)
    if args.cmd == "tables":
        return cmd_tables(args)
    if args.cmd == "columns":
        return cmd_columns(args)
    if args.cmd == "query":
        return cmd_query(args)
    if args.cmd == "run-file":
        args.sql = None
        args.file = args.file
        return cmd_query(args)
    if args.cmd == "catalog":
        return cmd_catalog(args)

    parser.print_help()
    return 1


if __name__ == "__main__":
    sys.exit(main())