import pymysql
import sys

db_host = "rm-f8ze60yirdj8786u2.mysql.rds.aliyuncs.com "
db_user = "root"
db_pass = "csqz@20255"
db_name = "csqz-client"

def initialize():
    try:
        conn = pymysql.connect(
            host=db_host,
            user=db_user,
            password=db_pass,
            db=db_name,
            port=3306
        )
        cur = conn.cursor()
        
        # Ensure genealogy_records table exists
        cur.execute("""
        CREATE TABLE IF NOT EXISTS genealogy_records (
            id INT AUTO_INCREMENT PRIMARY KEY,
            file_name VARCHAR(255) NOT NULL,
            oss_url TEXT NOT NULL,
            page_number INT,
            upload_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        """)
        
        # Add a test user if not exists
        # We need to match the existing schema
        cur.execute("SELECT * FROM users WHERE username='genealogy_admin'")
        if not cur.fetchone():
            sql = """
            INSERT INTO users (username, password, password_hash, full_name, role, is_active) 
            VALUES (%s, %s, %s, %s, %s, %s)
            """
            # Using plain password for 'password' column and a dummy hash for 'password_hash'
            cur.execute(sql, ('genealogy_admin', 'admin123', 'dummy_hash', 'Genealogy Admin', 'admin', 1))
            print("Added user 'genealogy_admin' with password 'admin123'")
        
        conn.commit()
        print("Initialization complete.")
        return True
        
    except Exception as e:
        print(f"Error during initialization: {e}")
        return False
    finally:
        if 'conn' in locals() and conn.open:
            conn.close()

if __name__ == "__main__":
    if initialize():
        print("SUCCESS: Database initialized.")
    else:
        sys.exit(1)