Autoparts-DB/fix_arra_data.py

"""
Script to fix the ARRA vehicle data in the database.
This will remove the incorrect entry under Aston Martin and add the correct ARRA entry.
"""
import sqlite3


def connect_db(db_path: str = "vehicle_database/vehicle_database.db") -> sqlite3.Connection:
    """Connect to the vehicle database."""
    conn = sqlite3.connect(db_path)
    conn.execute("PRAGMA foreign_keys = ON")  # Enable foreign key constraints
    return conn


def get_or_insert_brand(conn: sqlite3.Connection, brand_name: str) -> int:
    """Get brand ID or insert if it doesn't exist."""
    cursor = conn.cursor()
    
    # Check if brand already exists
    cursor.execute("SELECT id FROM brands WHERE name = ?", (brand_name,))
    result = cursor.fetchone()
    
    if result:
        return result[0]  # Return existing ID
    
    # Insert new brand
    cursor.execute("INSERT INTO brands (name) VALUES (?)", (brand_name,))
    conn.commit()
    return cursor.lastrowid


def get_or_insert_model(conn: sqlite3.Connection, brand_id: int, model_name: str) -> int:
    """Get model ID or insert if it doesn't exist."""
    cursor = conn.cursor()
    
    # Check if model already exists for this brand
    cursor.execute(
        "SELECT id FROM models WHERE brand_id = ? AND name = ?", 
        (brand_id, model_name)
    )
    result = cursor.fetchone()
    
    if result:
        return result[0]  # Return existing ID
    
    # Insert new model
    cursor.execute(
        "INSERT INTO models (brand_id, name) VALUES (?, ?)", 
        (brand_id, model_name)
    )
    conn.commit()
    return cursor.lastrowid


def get_or_insert_engine(conn: sqlite3.Connection, engine_name: str) -> int:
    """Get engine ID or insert if it doesn't exist."""
    cursor = conn.cursor()
    
    # Check if engine already exists
    cursor.execute("SELECT id FROM engines WHERE name = ?", (engine_name,))
    result = cursor.fetchone()
    
    if result:
        return result[0]  # Return existing ID
    
    # Insert new engine
    cursor.execute("INSERT INTO engines (name) VALUES (?)", (engine_name,))
    conn.commit()
    return cursor.lastrowid


def get_or_insert_year(conn: sqlite3.Connection, year: int) -> int:
    """Get year ID or insert if it doesn't exist."""
    cursor = conn.cursor()
    
    # Check if year already exists
    cursor.execute("SELECT id FROM years WHERE year = ?", (year,))
    result = cursor.fetchone()
    
    if result:
        return result[0]  # Return existing ID
    
    # Insert new year
    cursor.execute("INSERT INTO years (year) VALUES (?)", (year,))
    conn.commit()
    return cursor.lastrowid


def remove_incorrect_arr_entry(conn: sqlite3.Connection) -> None:
    """Remove the incorrect ARRA entry under Aston Martin."""
    cursor = conn.cursor()
    
    # Find the model ID for ARRA under Aston Martin
    cursor.execute("""
        SELECT m.id 
        FROM models m
        JOIN brands b ON m.brand_id = b.id
        WHERE b.name = 'ASTON MARTIN' AND m.name = 'ARRA'
    """)
    arra_model_result = cursor.fetchone()
    
    if arra_model_result:
        arra_model_id = arra_model_result[0]
        
        # Remove the incorrect relationship from model_year_engine
        cursor.execute(
            "DELETE FROM model_year_engine WHERE model_id = ?", 
            (arra_model_id,)
        )
        
        # Remove the ARRA model from models table
        cursor.execute(
            "DELETE FROM models WHERE id = ?", 
            (arra_model_id,)
        )
        
        print(f"Removed incorrect ARRA model with ID {arra_model_id} from Aston Martin")
    
    conn.commit()


def insert_correct_arr_entry(conn: sqlite3.Connection) -> None:
    """Insert the correct ARRA EW-1 entry."""
    cursor = conn.cursor()
    
    # Get or create ARRA brand
    arra_brand_id = get_or_insert_brand(conn, "ARRA")
    print(f"ARRA brand ID: {arra_brand_id}")
    
    # Get or create EW-1 model under ARRA brand
    ew1_model_id = get_or_insert_model(conn, arra_brand_id, "EW-1")
    print(f"EW-1 model ID: {ew1_model_id}")
    
    # Get or create Electric engine
    engine_id = get_or_insert_engine(conn, "Electric")
    print(f"Electric engine ID: {engine_id}")
    
    # Get or create 2024 year
    year_id = get_or_insert_year(conn, 2024)
    print(f"Year 2024 ID: {year_id}")
    
    # Check if the correct relationship already exists
    cursor.execute(
        "SELECT id FROM model_year_engine WHERE model_id = ? AND year_id = ? AND engine_id = ?", 
        (ew1_model_id, year_id, engine_id)
    )
    result = cursor.fetchone()
    
    if not result:
        # Insert the correct relationship
        cursor.execute(
            "INSERT INTO model_year_engine (model_id, year_id, engine_id) VALUES (?, ?, ?)", 
            (ew1_model_id, year_id, engine_id)
        )
        conn.commit()
        print("Added correct ARRA EW-1 2024 Electric entry")
    else:
        print("Correct ARRA EW-1 2024 Electric entry already exists")


def main():
    """Main function to fix ARRA data."""
    print("Fixing ARRA vehicle data...")
    
    # Connect to database
    conn = connect_db()
    
    # Remove the incorrect entry
    remove_incorrect_arr_entry(conn)
    
    # Insert the correct entry
    insert_correct_arr_entry(conn)
    
    # Close connection
    conn.close()
    print("ARRA data fixed successfully!")


if __name__ == "__main__":
    main()