import_db.py

import argparse
import subprocess
import mysql.connector
import sys
import os

def get_db_connection(host, user, password, port=3306):
    try:
        conn = mysql.connector.connect(
            host=host,
            user=user,
            password=password,
            port=port
        )
        return conn
    except mysql.connector.Error as err:
        print(f"Error connecting to database server: {err}")
        return None

def create_database_if_not_exists(cursor, db_name):
    try:
        cursor.execute(f"CREATE DATABASE IF NOT EXISTS `{db_name}`")
        print(f"Database '{db_name}' checked/created.")
    except mysql.connector.Error as err:
        print(f"Error creating database: {err}")
        sys.exit(1)

def import_dump(host, user, password, port, db_name, sql_file):
    if not os.path.exists(sql_file):
        print(f"Error: File {sql_file} not found.")
        sys.exit(1)

    print(f"Importing {sql_file} into {db_name}...")

    cmd = [
        "mysql",
        "-h", host,
        "-P", str(port),
        "-u", user,
        f"-p{password}",
        db_name
    ]

    try:
        with open(sql_file, "r") as infile:
            subprocess.run(cmd, stdin=infile, check=True, stderr=subprocess.PIPE)
        print(f"Successfully imported {sql_file} into {db_name}")
    except subprocess.CalledProcessError as e:
        print(f"Error importing database: {e.stderr.decode()}")

def main():
    parser = argparse.ArgumentParser(description="Import a MySQL database.")
    parser.add_argument("--host", default="localhost", help="Database host")
    parser.add_argument("--user", required=True, help="Database username")
    parser.add_argument("--password", required=True, help="Database password")
    parser.add_argument("--port", default=3306, type=int, help="Database port")
    parser.add_argument("--db-name", required=True, help="Target Database Name")
    parser.add_argument("--file", required=True, help="Path to SQL dump file")

    args = parser.parse_args()

    # Prompt user check
    response = input(f"Have you created the user '{args.user}' on this server (or does it already exist)? [y/N]: ")
    if response.lower() != 'y':
        print("Please create the user first. You can use the instructions in 'restore_user.md' if you exported it previously.")
        sys.exit(0)

    conn = get_db_connection(args.host, args.user, args.password, args.port)
    if not conn:
        print("Could not connect to server. Please check credentials or ensure the user exists.")
        sys.exit(1)

    cursor = conn.cursor()

    create_database_if_not_exists(cursor, args.db_name)

    cursor.close()
    conn.close()

    # Import logic uses subprocess
    import_dump(args.host, args.user, args.password, args.port, args.db_name, args.file)

if __name__ == "__main__":
    main()