src/database.py

import sqlite3
import os
import random
from datetime import datetime, timedelta

def setup_database(db_path="data/company_sales.db"):
    # Ensure directory exists
    os.makedirs(os.path.dirname(db_path), exist_ok=True)
    
    # Connect to SQLite
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    # Enable foreign keys
    cursor.execute("PRAGMA foreign_keys = ON;")
    
    # Drop tables if they exist
    cursor.execute("DROP TABLE IF EXISTS sales;")
    cursor.execute("DROP TABLE IF EXISTS products;")
    cursor.execute("DROP TABLE IF EXISTS employees;")
    cursor.execute("DROP TABLE IF EXISTS departments;")
    
    # Create tables
    cursor.execute("""
    CREATE TABLE departments (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        name TEXT NOT NULL UNIQUE,
        manager_id INTEGER
    );
    """)
    
    cursor.execute("""
    CREATE TABLE employees (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        name TEXT NOT NULL,
        department_id INTEGER,
        salary INTEGER NOT NULL,
        hire_date TEXT NOT NULL,
        manager_id INTEGER,
        FOREIGN KEY (department_id) REFERENCES departments(id),
        FOREIGN KEY (manager_id) REFERENCES employees(id)
    );
    """)
    
    cursor.execute("""
    CREATE TABLE products (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        name TEXT NOT NULL UNIQUE,
        category TEXT NOT NULL,
        price REAL NOT NULL
    );
    """)
    
    cursor.execute("""
    CREATE TABLE sales (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        employee_id INTEGER NOT NULL,
        product_id INTEGER NOT NULL,
        amount REAL NOT NULL,
        quantity INTEGER NOT NULL,
        sale_date TEXT NOT NULL,
        FOREIGN KEY (employee_id) REFERENCES employees(id),
        FOREIGN KEY (product_id) REFERENCES products(id)
    );
    """)
    
    # Seed Departments
    departments = [
        ("Sales", None),
        ("Engineering", None),
        ("Marketing", None),
        ("Human Resources", None),
        ("Finance", None)
    ]
    for dept, _ in departments:
        cursor.execute("INSERT INTO departments (name) VALUES (?);", (dept,))
    
    # Seed Products
    products = [
        ("Enterprise CRM", "Software", 1200.00),
        ("Analytics Suite", "Software", 800.00),
        ("Collaboration Tools", "Software", 300.00),
        ("Security Endpoint", "Hardware", 150.00),
        ("Cloud Backup", "Cloud", 50.00),
        ("AI Copilot Plug-in", "Software", 250.00)
    ]
    for name, cat, price in products:
        cursor.execute("INSERT INTO products (name, category, price) VALUES (?, ?, ?);", (name, cat, price))
        
    # Seed Employees
    employees = [
        # Name, Department ID, Salary, Hire Date, Manager ID
        ("Alice Vance", 1, 95000, "2023-01-15", None),      # Manager of Sales (ID: 1)
        ("Bob Smith", 2, 110000, "2022-06-01", None),       # Manager of Engineering (ID: 2)
        ("Charlie Brown", 3, 85000, "2023-03-10", None),    # Manager of Marketing (ID: 3)
        ("Diana Prince", 4, 80000, "2021-11-20", None),     # Manager of HR (ID: 4)
        ("Evan Wright", 5, 90000, "2022-01-10", None),      # Manager of Finance (ID: 5)
        
        ("Frank Miller", 1, 65000, "2024-02-15", 1),
        ("Grace Hopper", 2, 105000, "2023-08-01", 2),
        ("Henry Cavill", 2, 95000, "2024-01-10", 2),
        ("Ivy Queen", 1, 72000, "2024-05-12", 1),
        ("Jack Sparrow", 3, 58000, "2024-09-01", 3),
        ("Karen Smith", 4, 60000, "2025-01-15", 4),
        ("Leo Messi", 5, 75000, "2024-11-01", 5),
        ("Mona Lisa", 1, 80000, "2023-05-20", 1)
    ]
    
    for name, dept_id, salary, hire_date, mgr_id in employees:
        cursor.execute("INSERT INTO employees (name, department_id, salary, hire_date, manager_id) VALUES (?, ?, ?, ?, ?);", 
                       (name, dept_id, salary, hire_date, mgr_id))
        
    # Update managers in departments table now that employee records exist
    cursor.execute("UPDATE departments SET manager_id = 1 WHERE name = 'Sales';")
    cursor.execute("UPDATE departments SET manager_id = 2 WHERE name = 'Engineering';")
    cursor.execute("UPDATE departments SET manager_id = 3 WHERE name = 'Marketing';")
    cursor.execute("UPDATE departments SET manager_id = 4 WHERE name = 'Human Resources';")
    cursor.execute("UPDATE departments SET manager_id = 5 WHERE name = 'Finance';")
    
    # Seed Sales
    # We will generate 60 realistic sales records across 2024 and 2025
    # Keep random seed stable for reproducibility
    random.seed(42)
    
    # Sales reps: Alice (1), Frank (6), Ivy (9), Mona (13)
    sales_reps = [1, 6, 9, 13]
    product_ids = [1, 2, 3, 4, 5, 6]
    product_prices = {1: 1200.0, 2: 800.0, 3: 300.0, 4: 150.0, 5: 50.0, 6: 250.0}
    
    start_date = datetime(2024, 1, 1)
    for i in range(75):
        emp_id = random.choice(sales_reps)
        prod_id = random.choice(product_ids)
        qty = random.choice([1, 2, 3, 5, 10])
        price = product_prices[prod_id]
        
        # Calculate amount with random bulk discounts for large quantities
        amount = qty * price
        if qty >= 5:
            amount *= 0.9 # 10% discount
            
        days_offset = random.randint(0, 700) # Covers 2024 to late 2025
        sale_date = (start_date + timedelta(days=days_offset)).strftime("%Y-%m-%d")
        
        cursor.execute("INSERT INTO sales (employee_id, product_id, amount, quantity, sale_date) VALUES (?, ?, ?, ?, ?);",
                       (emp_id, prod_id, round(amount, 2), qty, sale_date))
        
    conn.commit()
    print(f"Database created successfully at {db_path} with:")
    print(f"  - 5 Departments")
    print(f"  - 6 Products")
    print(f"  - 13 Employees")
    print(f"  - 75 Sales Records")
    
    # Print database schema details to verify
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
    tables = cursor.fetchall()
    print("Tables created:", [t[0] for t in tables])
    
    conn.close()

if __name__ == "__main__":
    setup_database()