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

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

BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
SALARY_FILE = os.path.join(BASE_DIR, '马兰拉面数据分析', '薪资拆分明细表_脱敏.xlsx')
ATTENDANCE_FILE = os.path.join(BASE_DIR, '马兰拉面数据分析', '考勤数据表-脱敏.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 0
    try:
        return float(val)
    except (ValueError, TypeError):
        return 0

def to_text(val):
    if val is None:
        return None
    s = str(val).strip()
    return s if s else None

def import_salary(conn):
    print(f'导入薪资拆分明细表: {SALARY_FILE}')
    sha = file_sha256(SALARY_FILE)
    wb = openpyxl.load_workbook(SALARY_FILE, read_only=True, data_only=True)
    ws = wb[wb.sheetnames[0]]

    # Row 1: title, Row 2: merged title, Row 3: headers, Row 4+: data
    rows = list(ws.iter_rows(min_row=4, values_only=True))
    data_rows = [r for r in rows if r[0] is not None and r[9] is not None and str(r[9]).strip()]
    print(f'  数据行数: {len(data_rows)}')

    report_month = '2026-04-01'
    source_file = os.path.basename(SALARY_FILE)

    cur = conn.cursor()
    # Check if already imported
    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]}, 跳过')
        wb.close()
        return

    # Insert import log
    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}')

    # Batch insert
    insert_sql = '''INSERT INTO public.salary_detail_records (
        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
    ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)'''

    batch = []
    batch_size = 500
    for idx, r in enumerate(data_rows):
        source_row = idx + 4  # Excel row number
        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(vals)
        if len(batch) >= batch_size:
            cur.executemany(insert_sql, batch)
            conn.commit()
            print(f'  已导入 {idx + 1}/{len(data_rows)} 行')
            batch = []

    if batch:
        cur.executemany(insert_sql, batch)
        conn.commit()
        print(f'  已导入 {len(data_rows)}/{len(data_rows)} 行')

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

def import_attendance(conn):
    print(f'\n导入考勤数据表: {ATTENDANCE_FILE}')
    sha = file_sha256(ATTENDANCE_FILE)
    wb = openpyxl.load_workbook(ATTENDANCE_FILE, read_only=True, data_only=True)
    ws = wb[wb.sheetnames[0]]

    # Row 1: headers, Row 2+: data
    rows = list(ws.iter_rows(min_row=2, values_only=True))
    data_rows = [r for r in rows if r[0] is not None]
    print(f'  数据行数: {len(data_rows)}')

    report_month = '2026-04-01'
    source_file = os.path.basename(ATTENDANCE_FILE)

    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]}, 跳过')
        wb.close()
        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}')

    # 33 columns: 工号, 岗位, 所在部门, day_01..day_30
    cols = ['import_id', 'source_row', 'employee_code', 'position', 'department']
    cols += [f'day_{str(i).zfill(2)}' for i in range(1, 31)]
    placeholders = ', '.join(['%s'] * len(cols))
    col_names = ', '.join(cols)
    insert_sql = f'INSERT INTO public.attendance_records ({col_names}) VALUES ({placeholders})'

    batch = []
    batch_size = 500
    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-30 (columns 3-32)
        for d in range(30):
            vals.append(to_text(r[3 + d]) if 3 + d < len(r) else None)
        batch.append(vals)
        if len(batch) >= batch_size:
            cur.executemany(insert_sql, batch)
            conn.commit()
            print(f'  已导入 {idx + 1}/{len(data_rows)} 行')
            batch = []

    if batch:
        cur.executemany(insert_sql, batch)
        conn.commit()
        print(f'  已导入 {len(data_rows)}/{len(data_rows)} 行')

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

def main():
    conn = psycopg2.connect(**DB_CONFIG)
    conn.autocommit = False
    try:
        import_salary(conn)
        import_attendance(conn)
        print('\n=== 导入完成 ===')
    except Exception as e:
        conn.rollback()
        print(f'错误: {e}', file=sys.stderr)
        raise
    finally:
        conn.close()

if __name__ == '__main__':
    main()