#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import sys
import re
import argparse
import json
import psycopg2
import pandas as pd

def main():
    # Parse command line arguments
    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()

    # Get start and end dates
    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)

    # Read PostgreSQL password from secrets.env
    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)

    # Read SQL templates
    script_dir = os.path.dirname(os.path.abspath(__file__))
    summary_sql_path = os.path.join(script_dir, 'sql', 'summary.sql')
    detail_sql_path = os.path.join(script_dir, 'sql', 'detail.sql')
    prepaid_sql_path = os.path.join(script_dir, 'sql', 'prepaid.sql')

    if not os.path.exists(summary_sql_path) or not os.path.exists(detail_sql_path):
        print("错误：SQL模板文件不存在", file=sys.stderr)
        sys.exit(1)
    
    with open(summary_sql_path, 'r', encoding='utf-8') as f:
        summary_sql_template = f.read()
    with open(detail_sql_path, 'r', encoding='utf-8') as f:
        detail_sql_template = f.read()
    prepaid_sql_template = None
    if os.path.exists(prepaid_sql_path):
        with open(prepaid_sql_path, 'r', encoding='utf-8') as f:
            prepaid_sql_template = f.read()

    # Replace placeholders
    summary_sql = summary_sql_template.format(period_start=start_date, period_end=end_date)
    detail_sql = detail_sql_template.format(period_start=start_date, period_end=end_date)
    prepaid_sql = prepaid_sql_template.format(period_start=start_date, period_end=end_date) if prepaid_sql_template else None

    # Connect to PostgreSQL
    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)

    # Execute queries
    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()

    # Write to 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)

    # Print summary result as 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()