#!/usr/bin/env python3
"""
薪资拆分明细表 + 考勤数据表 导入脚本
遵循现有 import_log + records 模式
"""
import hashlib
import psycopg2
import os
import sys
import argparse
import pandas as pd
from psycopg2.extras import execute_values

DB_CONFIG = {
    'host': 'localhost',
    'port': 5432,
    'dbname': 'bill_query',
    'user': 'freedak',
    'password': '',
}

BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
SALARY_FILE = os.path.join(BASE_DIR, '数据', '4月', '薪资拆分明细表_脱敏.xlsx')
ATTENDANCE_FILE = os.path.join(BASE_DIR, '数据', '4月', '考勤数据表-脱敏.xlsx')

def file_sha256(filepath):
    h = hashlib.sha256()
    with open(filepath, 'rb') as f:
        for chunk in iter(lambda: f.read(8192), b''):
            h.update(chunk)
    return h.hexdigest()

def to_num(val):
    if val is None or val == '':
        return None
    if isinstance(val, float) and pd.isna(val):
        return None
    try:
        return float(val)
    except (ValueError, TypeError):
        return None

def to_text(val):
    if val is None:
        return None
    if isinstance(val, float) and pd.isna(val):
        return None
    s = str(val).strip()
    if s.lower() == 'nan' or s.lower() == 'none':
        return None
    return s if s else None

def import_salary(conn, report_month='2026-04-01', salary_file=None):
    filepath = salary_file or SALARY_FILE
    print(f'导入薪资拆分明细表: {filepath}')
    sha = file_sha256(filepath)
    
    # Read with pandas (row 3 = headers, row 4+ = data)
    raw_df = pd.read_excel(filepath, header=None, engine='openpyxl')
    print(f'  总行数: {len(raw_df)}, 列数: {len(raw_df.columns)}')
    
    # Data starts from row 3 (0-indexed), filter rows with employee_code (col 9)
    df = raw_df.iloc[3:].copy()
    df = df[df[9].notna() & (df[9].astype(str).str.strip() != '') & (df[9].astype(str).str.strip() != 'nan')]
    data_rows = df.values.tolist()
    print(f'  数据行数: {len(data_rows)}')
    
    source_file = os.path.basename(filepath)

    cur = conn.cursor()
    cur.execute(
        'SELECT import_id FROM public.salary_import_log WHERE report_month = %s AND source_file = %s AND file_sha256 = %s',
        (report_month, source_file, sha)
    )
    existing = cur.fetchone()
    if existing:
        print(f'  已导入过, import_id={existing[0]}, 跳过')
        return

    cur.execute(
        '''INSERT INTO public.salary_import_log (report_month, source_file, file_sha256, workbook_rows, imported_rows)
           VALUES (%s, %s, %s, %s, %s) RETURNING import_id''',
        (report_month, source_file, sha, len(data_rows), len(data_rows))
    )
    import_id = cur.fetchone()[0]
    print(f'  import_id={import_id}')

    # Column mapping: use execute_values for batch insert
    col_names = '''import_id, source_row,
        org_level1, org_level2, org_level3, org_level4, org_level5, org_level6, org_level7, org_level8,
        employee_code, salary_period, position, work_type, employment_type, hire_date, leave_date,
        salary_standard, base_wage, overtime_subsidy, social_subsidy, position_wage, tenure_wage,
        hourly_rate, total_wage, brand_wage, rent_subsidy,
        month_days, expected_attend, calc_attend, actual_attend, actual_hours,
        expected_rest, actual_rest, expected_legal_holiday, actual_legal_holiday,
        comp_leave_days, annual_leave, marriage_leave, paid_leave, bereavement_leave,
        personal_leave_days, personal_leave_deduction, furlough_days, furlough_deduction,
        injury_leave_days, injury_leave_deduction, absent_days, absent_deduction, join_leave_absent,
        attendance_wage, rest_subsidy, overtime_pay, shift_subsidy, comp_leave_wage,
        cashier_amount, bonus, subsidy, subsidy_remark, rush_wage, total_supplement,
        late_deduction, no_punch_deduction, internal_social_total, internal_social_remark,
        loan, loan_remark, fine, compensation, total_deduction, injury_wage, singapore_subsidy,
        perf_standard, perf_score, perf_amount, phone_allowance, dorm_fee,
        net_salary, gross_pay, income_tax, net_pay,
        external_base_standard, external_overtime, external_absence, external_bonus,
        external_subsidy, external_other_deduction, external_gross,
        pension_deduction, medical_deduction, unemployment_deduction,
        external_social_total, external_tax, external_net, internal_tax, internal_net,
        external_unit, attendance_remark, employment_type_orig, salary_category'''

    batch = []
    for idx, r in enumerate(data_rows):
        source_row = idx + 4
        vals = [
            import_id, source_row,
            to_text(r[1]), to_text(r[2]), to_text(r[3]), to_text(r[4]), to_text(r[5]), to_text(r[6]), to_text(r[7]), to_text(r[8]),
            to_text(r[9]), to_text(r[10]), to_text(r[11]), to_text(r[12]), to_text(r[13]), to_text(r[14]), to_text(r[15]),
            to_num(r[16]), to_num(r[17]), to_num(r[18]), to_num(r[19]), to_num(r[20]), to_num(r[21]),
            to_num(r[22]), to_num(r[23]), to_num(r[24]), to_num(r[25]),
            to_num(r[26]), to_num(r[27]), to_num(r[28]), to_num(r[29]), to_num(r[30]),
            to_num(r[31]), to_num(r[32]), to_num(r[33]), to_num(r[34]),
            to_num(r[35]), to_num(r[36]), to_num(r[37]), to_num(r[38]), to_num(r[39]),
            to_num(r[40]), to_num(r[41]), to_num(r[42]), to_num(r[43]),
            to_num(r[44]), to_num(r[45]), to_num(r[46]), to_num(r[47]), to_num(r[48]),
            to_num(r[49]), to_num(r[50]), to_num(r[51]), to_num(r[52]), to_num(r[53]),
            to_num(r[54]), to_num(r[55]), to_num(r[56]), to_text(r[57]), to_num(r[58]), to_num(r[59]),
            to_num(r[60]), to_num(r[61]), to_num(r[62]), to_text(r[63]),
            to_num(r[64]), to_text(r[65]), to_num(r[66]), to_num(r[67]), to_num(r[68]), to_num(r[69]), to_num(r[70]),
            to_num(r[71]), to_num(r[72]), to_num(r[73]), to_num(r[74]), to_num(r[75]),
            to_num(r[76]), to_num(r[77]), to_num(r[78]), to_num(r[79]),
            to_num(r[80]), to_num(r[81]), to_num(r[82]), to_num(r[83]),
            to_num(r[84]), to_num(r[85]), to_num(r[86]),
            to_num(r[87]), to_num(r[88]), to_num(r[89]),
            to_num(r[90]), to_num(r[91]), to_num(r[92]), to_num(r[93]), to_num(r[94]),
            to_text(r[95]), to_text(r[96]), to_text(r[97]), to_text(r[98]),
        ]
        batch.append(tuple(vals))
        if len(batch) >= 500:
            execute_values(cur, f'INSERT INTO public.salary_detail_records ({col_names}) VALUES %s', batch, page_size=500)
            conn.commit()
            print(f'  已导入 {idx + 1}/{len(data_rows)} 行')
            batch = []

    if batch:
        execute_values(cur, f'INSERT INTO public.salary_detail_records ({col_names}) VALUES %s', batch, page_size=500)
        conn.commit()
        print(f'  已导入 {len(data_rows)}/{len(data_rows)} 行')

    print(f'  薪资明细导入完成: {len(data_rows)} 行')

def import_attendance(conn, report_month='2026-04-01', attendance_file=None):
    filepath = attendance_file or ATTENDANCE_FILE
    print(f'\n导入考勤数据表: {filepath}')
    sha = file_sha256(filepath)
    
    # Read with pandas (row 0 = headers, row 1+ = data)
    raw_df = pd.read_excel(filepath, header=None, engine='openpyxl')
    print(f'  总行数: {len(raw_df)}, 列数: {len(raw_df.columns)}')
    
    # Data starts from row 1
    df = raw_df.iloc[1:].copy()
    df = df[df[0].notna() & (df[0].astype(str).str.strip() != '') & (df[0].astype(str).str.strip() != 'nan')]
    data_rows = df.values.tolist()
    print(f'  数据行数: {len(data_rows)}')
    
    source_file = os.path.basename(filepath)

    cur = conn.cursor()
    cur.execute(
        'SELECT import_id FROM public.attendance_import_log WHERE report_month = %s AND source_file = %s AND file_sha256 = %s',
        (report_month, source_file, sha)
    )
    existing = cur.fetchone()
    if existing:
        print(f'  已导入过, import_id={existing[0]}, 跳过')
        return

    cur.execute(
        '''INSERT INTO public.attendance_import_log (report_month, source_file, file_sha256, workbook_rows, imported_rows)
           VALUES (%s, %s, %s, %s, %s) RETURNING import_id''',
        (report_month, source_file, sha, len(data_rows), len(data_rows))
    )
    import_id = cur.fetchone()[0]
    print(f'  import_id={import_id}')

    # 34 columns: import_id, source_row, employee_code, position, department, day_01..day_31
    cols = ['import_id', 'source_row', 'employee_code', 'position', 'department']
    cols += [f'day_{str(i).zfill(2)}' for i in range(1, 32)]
    col_names = ','.join(cols)

    batch = []
    for idx, r in enumerate(data_rows):
        source_row = idx + 2
        vals = [import_id, source_row, to_text(r[0]), to_text(r[1]), to_text(r[2])]
        # Days 1-31 (columns 3-33)
        for d in range(31):
            vals.append(to_text(r[3 + d]) if 3 + d < len(r) else None)
        batch.append(tuple(vals))
        if len(batch) >= 500:
            execute_values(cur, f'INSERT INTO public.attendance_records ({col_names}) VALUES %s', batch, page_size=500)
            conn.commit()
            print(f'  已导入 {idx + 1}/{len(data_rows)} 行')
            batch = []

    if batch:
        execute_values(cur, f'INSERT INTO public.attendance_records ({col_names}) VALUES %s', batch, page_size=500)
        conn.commit()
        print(f'  已导入 {len(data_rows)}/{len(data_rows)} 行')

    print(f'  考勤数据导入完成: {len(data_rows)} 行')

def main():
    parser = argparse.ArgumentParser(description='薪资考勤数据导入')
    parser.add_argument('--month', default='2026-04-01', help='报告月份')
    parser.add_argument('--db', default='bill_query', help='数据库名')
    parser.add_argument('--skip-salary', action='store_true', help='跳过薪资导入')
    parser.add_argument('--skip-attendance', action='store_true', help='跳过考勤导入')
    parser.add_argument('--salary-file', default=None, help='薪资Excel文件路径 (默认使用内置路径)')
    parser.add_argument('--attendance-file', default=None, help='考勤Excel文件路径 (默认使用内置路径)')
    args = parser.parse_args()
    
    print(f'导入月份: {args.month}, 数据库: {args.db}')
    cfg = DB_CONFIG.copy()
    cfg['dbname'] = args.db
    conn = psycopg2.connect(**cfg)
    conn.autocommit = False
    try:
        if not args.skip_salary:
            import_salary(conn, args.month, args.salary_file)
        if not args.skip_attendance:
            import_attendance(conn, args.month, args.attendance_file)
        print('\n=== 导入完成 ===')
    except Exception as e:
        conn.rollback()
        print(f'错误: {e}', file=sys.stderr)
        raise
    finally:
        conn.close()

if __name__ == '__main__':
    main()