#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
订单均摊结算统计脚本
用法：
  python3 run.py --start 2026-03-01 --end 2026-03-31
  python3 run.py 2026-03-01 2026-03-31
"""
import os
import sys
import re
import argparse
import json
import psycopg2
import pandas as pd

def main():
    parser = argparse.ArgumentParser(description='订单均摊结算统计脚本')
    parser.add_argument('--start', type=str, help='账期起始日（YYYY-MM-DD）')
    parser.add_argument('--end', type=str, help='账期结束日（YYYY-MM-DD）')
    parser.add_argument('pos_start', nargs='?', type=str, help='账期起始日（位置参数）')
    parser.add_argument('pos_end', nargs='?', type=str, help='账期结束日（位置参数）')
    args = parser.parse_args()

    start_date = args.start or args.pos_start
    end_date = args.end or args.pos_end

    if not start_date or not end_date:
        print("错误：请提供账期起始日和结束日", file=sys.stderr)
        print("用法1：python3 run.py --start 2026-03-01 --end 2026-03-31", file=sys.stderr)
        print("用法2：python3 run.py 2026-03-01 2026-03-31", file=sys.stderr)
        sys.exit(1)

    # ── 1. 读取数据库密码 ──
    secrets_path = '/root/.openclaw/workspace/secrets.env'
    if not os.path.exists(secrets_path):
        print(f"错误：secrets.env文件不存在 {secrets_path}", file=sys.stderr)
        sys.exit(1)
    with open(secrets_path, 'r', encoding='utf-8') as f:
        secrets_content = f.read()
    pg_password_match = re.search(r"PG_ONLINE_PASSWORD='(.*)'", secrets_content)
    if not pg_password_match:
        print("错误：未找到PG_ONLINE_PASSWORD配置", file=sys.stderr)
        sys.exit(1)
    pg_password = pg_password_match.group(1)

    # ── 2. 读取 SQL 模板 ──
    script_dir = os.path.dirname(os.path.abspath(__file__))
    sql_dir = os.path.join(script_dir, 'sql')
    common_sql_path = os.path.join(sql_dir, '_common.sql')
    detail_sql_path = os.path.join(sql_dir, 'detail.sql')
    summary_sql_path = os.path.join(sql_dir, 'summary.sql')
    prepaid_sql_path = os.path.join(sql_dir, 'prepaid.sql')

    for p in [common_sql_path, detail_sql_path, summary_sql_path]:
        if not os.path.exists(p):
            print(f"错误：SQL模板文件不存在 {p}", file=sys.stderr)
            sys.exit(1)

    with open(common_sql_path, 'r', encoding='utf-8') as f:
        common_sql_raw = f.read()
    with open(detail_sql_path, 'r', encoding='utf-8') as f:
        detail_output_raw = f.read()
    with open(summary_sql_path, 'r', encoding='utf-8') as f:
        summary_output_raw = f.read()

    # 拼接公共CTE + 各自SELECT → 统一替换参数（避免多次格式化）
    detail_sql = (common_sql_raw + "\n" + detail_output_raw).format(
        period_start=start_date, period_end=end_date)
    summary_sql = (common_sql_raw + "\n" + summary_output_raw).format(
        period_start=start_date, period_end=end_date)

    prepaid_sql = None
    if os.path.exists(prepaid_sql_path):
        with open(prepaid_sql_path, 'r', encoding='utf-8') as f:
            prepaid_sql = f.read().format(period_start=start_date, period_end=end_date)

    # ── 3. 连接数据库并执行查询 ──
    try:
        conn = psycopg2.connect(
            host="bj-postgres-16pob4sg.sql.tencentcdb.com",
            port=28591,
            user="ai_member",
            password=pg_password,
            database="vala_bi"
        )
    except Exception as e:
        print(f"数据库连接失败：{str(e)}", file=sys.stderr)
        sys.exit(1)

    try:
        summary_df = pd.read_sql(summary_sql, conn)
        detail_df = pd.read_sql(detail_sql, conn)
        prepaid_df = pd.read_sql(prepaid_sql, conn) if prepaid_sql else None
    except Exception as e:
        print(f"SQL执行失败：{str(e)}", file=sys.stderr)
        conn.close()
        sys.exit(1)
    finally:
        conn.close()

    # ── 4. 输出 Excel 报表 ──
    output_dir = '/root/.openclaw/workspace/output'
    os.makedirs(output_dir, exist_ok=True)
    output_file = os.path.join(output_dir, f'订单均摊结算报表_{start_date}_{end_date}.xlsx')

    try:
        with pd.ExcelWriter(output_file, engine='openpyxl') as writer:
            summary_df.to_excel(writer, sheet_name='汇总表', index=False)
            detail_df.to_excel(writer, sheet_name='订单明细', index=False)
            if prepaid_df is not None:
                prepaid_df.to_excel(writer, sheet_name='本月预收账款', index=False)
    except Exception as e:
        print(f"Excel生成失败：{str(e)}", file=sys.stderr)
        sys.exit(1)

    # ── 5. 控制台输出 JSON 汇总 ──
    result = summary_df.to_dict(orient='records')[0]
    print(json.dumps(result, ensure_ascii=False, indent=2))
    print(f"\n报表已保存到：{output_file}", file=sys.stderr)
    sys.exit(0)

if __name__ == "__main__":
    main()
