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()
        
        # 1. Create users table
        print(f"Creating 'users' table in {db_name}...")
        cur.execute("""
        CREATE TABLE IF NOT EXISTS users (
            id INT AUTO_INCREMENT PRIMARY KEY,
            username VARCHAR(50) UNIQUE NOT NULL,
            password VARCHAR(255) NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        """)
        
        # 2. Create genealogy_records table
        print(f"Creating 'genealogy_records' table in {db_name}...")
        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
        )
        """)
        
        # 3. Add a default user
        print("Adding default user 'admin'...")
        cur.execute("SELECT * FROM users WHERE username='admin'")
        if not cur.fetchone():
            cur.execute("INSERT INTO users (username, password) VALUES (%s, %s)", ('admin', 'admin123'))
            print("Added default user 'admin' with password 'admin123'")
        else:
            print("User 'admin' already exists.")
        
        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:
        print("FAILURE: Database initialization failed.")
        sys.exit(1)