import psycopg2

DB_CONFIG = {
    "host": "192.168.110.252",
    "port": 15432,
    "database": "postgresql",
    "user": "postgresql",
    "password": "Jchl1528",
    "options": "-c search_path=hisdev",
}


def fix():
    conn = None
    try:
        conn = psycopg2.connect(**DB_CONFIG)
        cursor = conn.cursor()

        print("=" * 80)
        print("Fixing remaining surgery records")
        print("=" * 80)

        # Count before
        cursor.execute("""
            SELECT COUNT(*)
            FROM wor_service_request
            WHERE category_enum = 4
            AND delete_flag = '0'
            AND (content_json::jsonb->>'surgeryName' IS NULL OR content_json::jsonb->>'surgeryName' = '')
        """)
        before = cursor.fetchone()[0]
        print(f"\nBefore fix: {before} records without surgeryName")

        # Execute fix
        cursor.execute("""
            UPDATE wor_service_request wsr
            SET content_json = 
                COALESCE(
                    (wsr.content_json::jsonb || jsonb_build_object(
                        'surgeryName', cs.surgery_name,
                        'surgeryCode', COALESCE(cs.surgery_code, '')
                    ))::text,
                    jsonb_build_object(
                        'surgeryName', cs.surgery_name,
                        'surgeryCode', COALESCE(cs.surgery_code, '')
                    )::text
                )
            FROM cli_surgery cs
            WHERE wsr.category_enum = 4
            AND wsr.delete_flag = '0'
            AND (wsr.content_json::jsonb->>'surgeryName' IS NULL OR wsr.content_json::jsonb->>'surgeryName' = '')
            AND cs.id = wsr.activity_id
            AND cs.surgery_name IS NOT NULL
            AND cs.delete_flag = '0'
        """)

        fixed = cursor.rowcount
        conn.commit()

        print(f"Fixed: {fixed} records")

        # Count after
        cursor.execute("""
            SELECT COUNT(*)
            FROM wor_service_request
            WHERE category_enum = 4
            AND delete_flag = '0'
            AND (content_json::jsonb->>'surgeryName' IS NULL OR content_json::jsonb->>'surgeryName' = '')
        """)
        after = cursor.fetchone()[0]
        print(f"After fix: {after} records without surgeryName")

        if after == 0:
            print("\n*** ALL RECORDS FIXED! ***")
        else:
            print(f"\n*** {after} records still need attention ***")

        cursor.close()

    except Exception as e:
        print(f"Error: {e}")
        if conn:
            conn.rollback()
    finally:
        if conn:
            conn.close()


if __name__ == "__main__":
    fix()