"""
Script to add extended Audi vehicle data (2017-2021) to the database.
This includes additional models and year/motor combinations.
"""
import sqlite3
from typing import List, Tuple


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 insert_version_relationship(conn: sqlite3.Connection, model_id: int, year_id: int, engine_id: int) -> None:
    """Insert a version relationship if it doesn't exist."""
    cursor = conn.cursor()
    
    # Check if version relationship already exists
    cursor.execute(
        "SELECT id FROM model_year_engine WHERE model_id = ? AND year_id = ? AND engine_id = ?", 
        (model_id, year_id, engine_id)
    )
    result = cursor.fetchone()
    
    if not result:
        # Insert new version relationship
        cursor.execute(
            "INSERT INTO model_year_engine (model_id, year_id, engine_id) VALUES (?, ?, ?)", 
            (model_id, year_id, engine_id)
        )
        conn.commit()


def main():
    """Main function to add extended Audi data."""
    print("Adding extended Audi data (2017-2021)...")
    
    # Connect to database
    conn = connect_db()
    
    # Get brand ID (should already exist)
    brand_id = get_or_insert_brand(conn, "AUDI")
    print(f"Brand 'AUDI' has ID: {brand_id}")
    
    # Define additional models that might need to be added
    additional_models = [
        'A4 ALLROAD',
        'A6 ALLROAD',
        'Q2 QUATTRO'
    ]
    
    # Insert additional models if they don't exist
    for model_name in additional_models:
        model_id = get_or_insert_model(conn, brand_id, model_name)
    
    # Define versions (model, year, engine) for years 2017-2021
    versions = [
        # 2021
        ('A1', 2021, '1.5L L4 Turbocharged'),
        ('A1', 2021, '2.0L L4 Turbocharged'),
        ('A1', 2021, '999cc L3 Turbocharged'),
        ('A4', 2021, '2.0L L4 ELECTRIC/GAS Turbocharged'),
        ('A4', 2021, '2.0L L4 Turbocharged'),
        ('A4 ALLROAD', 2021, '2.0L L4 Turbocharged'),
        ('A5', 2021, '2.0L L4 Turbocharged'),
        ('A6', 2021, '2.0L L4 Turbocharged'),
        ('A6', 2021, '3.0L V6 Turbocharged'),
        ('A6 ALLROAD', 2021, '3.0L V6 Turbocharged'),
        ('A7', 2021, '2.0L L4 ELECTRIC/GAS Turbocharged'),
        ('A7', 2021, '2.0L L4 Turbocharged'),
        ('A7', 2021, '3.0L V6 Turbocharged'),
        ('A8', 2021, '3.0L V6 ELECTRIC/GAS Turbocharged'),
        ('A8', 2021, '3.0L V6 Turbocharged'),
        ('A8', 2021, '4.0L V8 Turbocharged'),
        ('E-TRON', 2021, 'ELECTRIC'),
        ('Q2', 2021, '1.4L L4 Turbocharged'),
        ('Q3', 2021, '1.4L L4 Turbocharged'),
        ('Q3', 2021, '2.0L L4 Turbocharged'),
        ('Q5', 2021, '2.0L L4 ELECTRIC/GAS Turbocharged'),
        ('Q5', 2021, '2.0L L4 Turbocharged'),
        ('Q7', 2021, '2.0L L4 Turbocharged'),
        ('Q7', 2021, '3.0L V6 Turbocharged'),
        ('Q8', 2021, '3.0L V6 Turbocharged'),
        ('R8', 2021, '5.2L V10'),
        ('RS3', 2021, '2.5L L5 Turbocharged'),
        ('RS5', 2021, '2.9L V6 Turbocharged'),
        ('RS6', 2021, '4.0L V8 Turbocharged'),
        ('RS7', 2021, '4.0L V8 Turbocharged'),
        ('RS Q3', 2021, '2.5L L5 Turbocharged'),
        ('RS Q8', 2021, '4.0L V8 Turbocharged'),
        ('S3', 2021, '2.0L L4 Turbocharged'),
        ('S4', 2021, '3.0L V6 Turbocharged'),
        ('S5', 2021, '3.0L V6 Turbocharged'),
        ('S6', 2021, '2.9L V6 Turbocharged'),
        ('S7', 2021, '2.9L V6 Turbocharged'),
        ('S8', 2021, '4.0L V8 Turbocharged'),
        ('SQ5', 2021, '3.0L V6 Turbocharged'),
        ('SQ7', 2021, '4.0L V8 Turbocharged'),
        ('SQ8', 2021, '4.0L V8 Turbocharged'),
        ('TT', 2021, '2.0L L4 Turbocharged'),
        ('TT', 2021, '2.5L L5 Turbocharged'),

        # 2020
        ('A1', 2020, '1.5L L4 Turbocharged'),
        ('A1', 2020, '2.0L L4 Turbocharged'),
        ('A1', 2020, '999cc L3 Turbocharged'),
        ('A3', 2020, '1.4L L4 Turbocharged'),
        ('A3', 2020, '2.0L L4 Turbocharged'),
        ('A4', 2020, '2.0L L4 Turbocharged'),
        ('A4 ALLROAD', 2020, '2.0L L4 Turbocharged'),
        ('A5', 2020, '2.0L L4 Turbocharged'),
        ('A6', 2020, '2.0L L4 Turbocharged'),
        ('A6', 2020, '3.0L V6 Turbocharged'),
        ('A6 ALLROAD', 2020, '3.0L V6 Turbocharged'),
        ('A7', 2020, '3.0L V6 Turbocharged'),
        ('A8', 2020, '3.0L V6 ELECTRIC/GAS Turbocharged'),
        ('A8', 2020, '3.0L V6 Turbocharged'),
        ('A8', 2020, '4.0L V8 Turbocharged'),
        ('E-TRON', 2020, 'ELECTRIC'),
        ('Q2', 2020, '1.4L L4 Turbocharged'),
        ('Q2 QUATTRO', 2020, '2.0L L4 Turbocharged'),
        ('Q3', 2020, '1.4L L4 Turbocharged'),
        ('Q3', 2020, '2.0L L4 Turbocharged'),
        ('Q5', 2020, '2.0L L4 ELECTRIC/GAS Turbocharged'),
        ('Q5', 2020, '2.0L L4 Turbocharged'),
        ('Q7', 2020, '2.0L L4 Turbocharged'),
        ('Q7', 2020, '3.0L V6 Turbocharged'),
        ('Q8', 2020, '3.0L V6 Turbocharged'),
        ('R8', 2020, '5.2L V10'),
        ('RS3', 2020, '2.5L L5 Turbocharged'),
        ('RS5', 2020, '2.9L V6 Turbocharged'),
        ('RS6', 2020, '4.0L V8 Turbocharged'),
        ('RS7', 2020, '4.0L V8 Turbocharged'),
        ('RS Q3', 2020, '2.5L L5 Turbocharged'),
        ('RS Q8', 2020, '4.0L V8 Turbocharged'),
        ('S3', 2020, '2.0L L4 Turbocharged'),
        ('S4', 2020, '3.0L V6 Turbocharged'),
        ('S5', 2020, '3.0L V6 Turbocharged'),
        ('S6', 2020, '2.9L V6 Turbocharged'),
        ('S7', 2020, '2.9L V6 Turbocharged'),
        ('S8', 2020, '4.0L V8 Turbocharged'),
        ('SQ5', 2020, '3.0L V6 Turbocharged'),
        ('SQ7', 2020, '4.0L V8 Turbocharged'),
        ('SQ8', 2020, '4.0L V8 Turbocharged'),
        ('TT', 2020, '2.0L L4 Turbocharged'),
        ('TT', 2020, '2.5L L5 Turbocharged'),

        # 2019
        ('A3', 2019, '1.4L L4 Turbocharged'),
        ('A3', 2019, '2.0L L4 Turbocharged'),
        ('A4', 2019, '2.0L L4 Turbocharged'),
        ('A4 ALLROAD', 2019, '2.0L L4 Turbocharged'),
        ('A5', 2019, '2.0L L4 Turbocharged'),
        ('A6', 2019, '2.0L L4 Turbocharged'),
        ('A6', 2019, '3.0L V6 Turbocharged'),
        ('A7', 2019, '3.0L V6 Turbocharged'),
        ('A8', 2019, '3.0L V6 Turbocharged'),
        ('A8', 2019, '4.0L V8 Turbocharged'),
        ('E-TRON', 2019, 'ELECTRIC'),
        ('Q2', 2019, '1.4L L4 Turbocharged'),
        ('Q2', 2019, '2.0L L4 Turbocharged'),
        ('Q3', 2019, '1.4L L4 Turbocharged'),
        ('Q3', 2019, '2.0L L4 Turbocharged'),
        ('Q5', 2019, '2.0L L4 Turbocharged'),
        ('Q7', 2019, '2.0L L4 Turbocharged'),
        ('Q7', 2019, '3.0L V6 DIESEL Turbocharged'),
        ('Q7', 2019, '3.0L V6 Supercharged'),
        ('Q8', 2019, '3.0L V6 Turbocharged'),
        ('R8', 2019, '5.2L V10'),
        ('RS3', 2019, '2.5L L5 Turbocharged'),
        ('RS5', 2019, '2.9L V6 Turbocharged'),
        ('RS6', 2019, '4.0L V8 Turbocharged'),
        ('RS7', 2019, '4.0L V8 Turbocharged'),
        ('RS Q3', 2019, '2.5L L5 Turbocharged'),
        ('RS Q8', 2019, '4.0L V8 Turbocharged'),
        ('S3', 2019, '2.0L L4 Turbocharged'),
        ('S4', 2019, '3.0L V6 Turbocharged'),
        ('S5', 2019, '3.0L V6 Turbocharged'),
        ('S6', 2019, '2.9L V6 Turbocharged'),
        ('S7', 2019, '2.9L V6 Turbocharged'),
        ('S8', 2019, '4.0L V8 Turbocharged'),
        ('SQ5', 2019, '3.0L V6 Turbocharged'),
        ('SQ7', 2019, '4.0L V8 Turbocharged'),
        ('SQ8', 2019, '4.0L V8 Turbocharged'),
        ('TT', 2019, '2.0L L4 Turbocharged'),
        ('TT', 2019, '2.5L L5 Turbocharged'),

        # 2018
        ('A1', 2018, '1.4L L4 Turbocharged'),
        ('A1', 2018, '1.8L L4 Turbocharged'),
        ('A3', 2018, '1.4L L4 ELECTRIC/GAS Turbocharged'),
        ('A3', 2018, '1.4L L4 Turbocharged'),
        ('A3', 2018, '1.8L L4 Turbocharged'),
        ('A3', 2018, '2.0L L4 Turbocharged'),
        ('A4', 2018, '2.0L L4 Turbocharged'),
        ('A4 ALLROAD', 2018, '2.0L L4 Turbocharged'),
        ('A5', 2018, '2.0L L4 Turbocharged'),
        ('A6', 2018, '1.8L L4 Turbocharged'),
        ('A6', 2018, '2.0L L4 Turbocharged'),
        ('A6', 2018, '3.0L V6 Supercharged'),
        ('A7', 2018, '2.0L L4 Turbocharged'),
        ('A7', 2018, '3.0L V6 Supercharged'),
        ('A8', 2018, '3.0L V6 Supercharged'),
        ('A8', 2018, '4.0L V8 Turbocharged'),
        ('A8', 2018, '6.3L W12'),
        ('Q2', 2018, '1.4L L4 Turbocharged'),
        ('Q2', 2018, '2.0L L4 Turbocharged'),
        ('Q3', 2018, '1.4L L4 Turbocharged'),
        ('Q3', 2018, '2.0L L4 Turbocharged'),
        ('Q5', 2018, '2.0L L4 Turbocharged'),
        ('Q7', 2018, '2.0L L4 Turbocharged'),
        ('Q7', 2018, '3.0L V6 DIESEL Turbocharged'),
        ('Q7', 2018, '3.0L V6 Supercharged'),
        ('Q8', 2018, '3.0L V6 Turbocharged'),
        ('R8', 2018, '5.2L V10'),
        ('RS3', 2018, '2.5L L5 Turbocharged'),
        ('RS5', 2018, '2.9L V6 Turbocharged'),
        ('RS6', 2018, '4.0L V8 Turbocharged'),
        ('RS7', 2018, '4.0L V8 Turbocharged'),
        ('RS Q3', 2018, '2.5L L5 Turbocharged'),
        ('RS Q8', 2018, '4.0L V8 Turbocharged'),
        ('S3', 2018, '2.0L L4 Turbocharged'),
        ('S4', 2018, '3.0L V6 Supercharged'),
        ('S4', 2018, '3.0L V6 Turbocharged'),
        ('S5', 2018, '3.0L V6 Supercharged'),
        ('S5', 2018, '3.0L V6 Turbocharged'),
        ('S6', 2018, '4.0L V8 Turbocharged'),
        ('S7', 2018, '4.0L V8 Turbocharged'),
        ('S8', 2018, '4.0L V8 Turbocharged'),
        ('SQ5', 2018, '3.0L V6 Supercharged'),
        ('SQ5', 2018, '3.0L V6 Turbocharged'),
        ('SQ7', 2018, '4.0L V8 Turbocharged'),
        ('SQ8', 2018, '4.0L V8 Turbocharged'),
        ('TT', 2018, '1.8L L4 Turbocharged'),
        ('TT', 2018, '2.0L L4 Turbocharged'),
        ('TT', 2018, '2.5L L5 Turbocharged'),

        # 2017
        ('A1', 2017, '1.4L L4 Turbocharged'),
        ('A1', 2017, '1.8L L4 Turbocharged'),
        ('A3', 2017, '1.4L L4 ELECTRIC/GAS Turbocharged'),
        ('A3', 2017, '1.4L L4 Turbocharged'),
        ('A3', 2017, '2.0L L4 Turbocharged'),
        ('A4', 2017, '2.0L L4 Turbocharged'),
        ('A4 ALLROAD', 2017, '2.0L L4 Turbocharged'),
        ('A5', 2017, '2.0L L4 Turbocharged'),
        ('A6', 2017, '1.8L L4 Turbocharged'),
        ('A6', 2017, '2.0L L4 Turbocharged'),
        ('A6', 2017, '3.0L V6 Supercharged'),
        ('A7', 2017, '2.0L L4 Turbocharged'),
        ('A7', 2017, '3.0L V6 Supercharged'),
        ('A8', 2017, '3.0L V6 Supercharged'),
        ('A8', 2017, '4.0L V8 Turbocharged'),
        ('A8', 2017, '6.3L W12'),
        ('Q3', 2017, '1.4L L4 Turbocharged'),
        ('Q3', 2017, '2.0L L4 Turbocharged'),
        ('Q5', 2017, '2.0L L4 Turbocharged'),
        ('Q5', 2017, '3.0L V6 DIESEL Turbocharged'),
        ('Q5', 2017, '3.0L V6 Supercharged'),
        ('Q7', 2017, '2.0L L4 Turbocharged'),
        ('Q7', 2017, '3.0L V6 DIESEL Turbocharged'),
        ('Q7', 2017, '3.0L V6 Supercharged'),
        ('Q8', 2017, '3.0L V6 Turbocharged'),
        ('R8', 2017, '5.2L V10'),
        ('RS3', 2017, '2.5L L5 Turbocharged'),
        ('RS5', 2017, '2.9L V6 Turbocharged'),
        ('RS6', 2017, '4.0L V8 Turbocharged'),
        ('RS7', 2017, '4.0L V8 Turbocharged'),
        ('RS Q3', 2017, '2.5L L5 Turbocharged'),
        ('RS Q8', 2017, '4.0L V8 Turbocharged'),
        ('S3', 2017, '2.0L L4 Turbocharged'),
        ('S4', 2017, '3.0L V6 Supercharged'),
        ('S4', 2017, '3.0L V6 Turbocharged'),
        ('S5', 2017, '3.0L V6 Supercharged'),
        ('S5', 2017, '3.0L V6 Turbocharged'),
        ('S6', 2017, '4.0L V8 Turbocharged'),
        ('S7', 2017, '4.0L V8 Turbocharged'),
        ('S8', 2017, '4.0L V8 Turbocharged'),
        ('SQ5', 2017, '3.0L V6 Supercharged'),
        ('SQ5', 2017, '3.0L V6 Turbocharged'),
        ('SQ7', 2017, '4.0L V8 Turbocharged'),
        ('SQ8', 2017, '4.0L V8 Turbocharged'),
        ('TT', 2017, '1.8L L4 Turbocharged'),
        ('TT', 2017, '2.0L L4 Turbocharged'),
        ('TT', 2017, '2.5L L5 Turbocharged')
    ]
    
    # Insert version relationships
    for model_name, year, engine in versions:
        model_id = get_or_insert_model(conn, brand_id, model_name)
        year_id = get_or_insert_year(conn, year)
        engine_id = get_or_insert_engine(conn, engine)
        insert_version_relationship(conn, model_id, year_id, engine_id)
    
    print(f"Inserted {len(versions)} additional version relationships for AUDI vehicles (2017-2021)")
    
    # Close connection
    conn.close()
    print("Extended Audi data added successfully!")


if __name__ == "__main__":
    main()