src/dataset.py

import json
import os

SCHEMA_PROMPT = """You are a Text-to-SQL generator. Given a database schema and a natural language question, write a valid SQLite query that answers the question. Do not explain the query, just output the raw SQL.

Database Schema:
Table departments:
  - id: INTEGER (PRIMARY KEY)
  - name: TEXT (NOT NULL, UNIQUE, e.g., 'Sales', 'Engineering', etc.)
  - manager_id: INTEGER (references employees.id)

Table employees:
  - id: INTEGER (PRIMARY KEY)
  - name: TEXT (NOT NULL)
  - department_id: INTEGER (references departments.id)
  - salary: INTEGER (NOT NULL)
  - hire_date: TEXT (NOT NULL, YYYY-MM-DD)
  - manager_id: INTEGER (references employees.id)

Table products:
  - id: INTEGER (PRIMARY KEY)
  - name: TEXT (NOT NULL, UNIQUE, e.g., 'Enterprise CRM')
  - category: TEXT (NOT NULL, e.g., 'Software', 'Hardware', 'Cloud')
  - price: REAL (NOT NULL)

Table sales:
  - id: INTEGER (PRIMARY KEY)
  - employee_id: INTEGER (NOT NULL, references employees.id)
  - product_id: INTEGER (NOT NULL, references products.id)
  - amount: REAL (NOT NULL)
  - quantity: INTEGER (NOT NULL)
  - sale_date: TEXT (NOT NULL, YYYY-MM-DD)"""

# Training dataset pairs (50 high-quality pairs covering diverse SQL capabilities)
TRAIN_DATA = [
    # Basic select and filter
    ("List all employees and their salaries.", "SELECT name, salary FROM employees;"),
    ("Show all departments.", "SELECT * FROM departments;"),
    ("Find all products in the Software category.", "SELECT * FROM products WHERE category = 'Software';"),
    ("List products that cost more than 500 dollars.", "SELECT * FROM products WHERE price > 500;"),
    ("Find the salary of Bob Smith.", "SELECT salary FROM employees WHERE name = 'Bob Smith';"),
    ("Who was hired after January 1st, 2024?", "SELECT name, hire_date FROM employees WHERE hire_date > '2024-01-01';"),
    ("Show the name and price of the product with ID 3.", "SELECT name, price FROM products WHERE id = 3;"),
    ("Find the sales records for product ID 1.", "SELECT * FROM sales WHERE product_id = 1;"),
    ("List employees who earn less than 70000.", "SELECT name, salary FROM employees WHERE salary < 70000;"),
    ("Show the department name of department with ID 2.", "SELECT name FROM departments WHERE id = 2;"),

    # Aggregations & Group By
    ("What is the average salary of all employees?", "SELECT AVG(salary) AS average_salary FROM employees;"),
    ("How many employees are there in total?", "SELECT COUNT(*) AS total_employees FROM employees;"),
    ("What is the highest salary among employees?", "SELECT MAX(salary) AS max_salary FROM employees;"),
    ("What is the lowest salary among employees?", "SELECT MIN(salary) AS min_salary FROM employees;"),
    ("Find the total revenue generated from all sales.", "SELECT SUM(amount) AS total_revenue FROM sales;"),
    ("How many items were sold in total?", "SELECT SUM(quantity) AS total_quantity_sold FROM sales;"),
    ("What is the average price of a software product?", "SELECT AVG(price) AS avg_price FROM products WHERE category = 'Software';"),
    ("How many products are in each category?", "SELECT category, COUNT(*) AS product_count FROM products GROUP BY category;"),
    ("Find the average salary by department.", "SELECT department_id, AVG(salary) AS avg_salary FROM employees GROUP BY department_id;"),
    ("Count the number of sales made by employee 6.", "SELECT COUNT(*) FROM sales WHERE employee_id = 6;"),

    # Joins
    ("List the names of employees alongside their department names.", "SELECT e.name, d.name AS department FROM employees e JOIN departments d ON e.department_id = d.id;"),
    ("Find the total quantity of CRM products sold.", "SELECT SUM(s.quantity) FROM sales s JOIN products p ON s.product_id = p.id WHERE p.name = 'Enterprise CRM';"),
    ("Show the manager's name for each department.", "SELECT d.name AS department, e.name AS manager FROM departments d JOIN employees e ON d.manager_id = e.id;"),
    ("Show the names of employees in the Engineering department.", "SELECT e.name FROM employees e JOIN departments d ON e.department_id = d.id WHERE d.name = 'Engineering';"),
    ("Who made a sale on Collaboration Tools?", "SELECT DISTINCT e.name FROM employees e JOIN sales s ON e.id = s.employee_id JOIN products p ON s.product_id = p.id WHERE p.name = 'Collaboration Tools';"),
    ("List all products sold by Alice Vance.", "SELECT DISTINCT p.name FROM products p JOIN sales s ON p.id = s.product_id JOIN employees e ON s.employee_id = e.id WHERE e.name = 'Alice Vance';"),
    ("Show all sales details, including employee name and product name.", "SELECT s.id, e.name AS employee_name, p.name AS product_name, s.amount, s.quantity, s.sale_date FROM sales s JOIN employees e ON s.employee_id = e.id JOIN products p ON s.product_id = p.id;"),
    ("Which employees generated sales of more than 1000 in a single transaction?", "SELECT DISTINCT e.name FROM employees e JOIN sales s ON e.id = s.employee_id WHERE s.amount > 1000;"),

    # Advanced aggregates & JOINs
    ("Find the total sales revenue generated by each employee in the Sales department.", "SELECT e.name, SUM(s.amount) AS total_revenue FROM employees e JOIN sales s ON e.id = s.employee_id JOIN departments d ON e.department_id = d.id WHERE d.name = 'Sales' GROUP BY e.name;"),
    ("Show the total amount of sales for each product, ordered by total revenue descending.", "SELECT p.name, SUM(s.amount) AS total_revenue FROM products p JOIN sales s ON p.id = s.product_id GROUP BY p.name ORDER BY total_revenue DESC;"),
    ("Which department has the highest average salary?", "SELECT d.name, AVG(e.salary) AS avg_salary FROM departments d JOIN employees e ON d.id = e.department_id GROUP d.name ORDER BY avg_salary DESC LIMIT 1;"),
    ("Find the employee who has made the most sales transactions.", "SELECT e.name, COUNT(s.id) AS transaction_count FROM employees e JOIN sales s ON e.id = s.employee_id GROUP BY e.name ORDER BY transaction_count DESC LIMIT 1;"),
    ("List departments that have an average employee salary greater than 80000.", "SELECT d.name, AVG(e.salary) AS avg_salary FROM departments d JOIN employees e ON d.id = e.department_id GROUP BY d.name HAVING avg_salary > 80000;"),
    
    # Dates
    ("Find all sales made in 2024.", "SELECT * FROM sales WHERE sale_date BETWEEN '2024-01-01' AND '2024-12-31';"),
    ("Find all sales made in the month of May 2024.", "SELECT * FROM sales WHERE sale_date BETWEEN '2024-05-01' AND '2024-05-31';"),
    ("List the names of employees hired in 2023.", "SELECT name FROM employees WHERE hire_date BETWEEN '2023-01-01' AND '2023-12-31';"),
    ("Show the sales revenue generated in the second half of 2024.", "SELECT SUM(amount) FROM sales WHERE sale_date BETWEEN '2024-07-01' AND '2024-12-31';"),

    # Subqueries & Complex
    ("List the products that have never been sold.", "SELECT p.name FROM products p LEFT JOIN sales s ON p.id = s.product_id WHERE s.id IS NULL;"),
    ("Which employees earn more than their manager?", "SELECT e.name FROM employees e JOIN employees m ON e.manager_id = m.id WHERE e.salary > m.salary;"),
    ("List all employees who report to Bob Smith.", "SELECT e.name FROM employees e JOIN employees m ON e.manager_id = m.id WHERE m.name = 'Bob Smith';"),
    ("Show the names of employees who have not made any sales.", "SELECT e.name FROM employees e LEFT JOIN sales s ON e.id = s.employee_id WHERE s.id IS NULL AND e.department_id = 1;"),
    ("Find products with a price higher than the average product price.", "SELECT name, price FROM products WHERE price > (SELECT AVG(price) FROM products);"),
    ("Which product had the highest single sales transaction amount?", "SELECT p.name, s.amount FROM products p JOIN sales s ON p.id = s.product_id ORDER BY s.amount DESC LIMIT 1;"),
    ("Show all employees who have the same manager as Grace Hopper.", "SELECT e.name FROM employees e WHERE e.manager_id = (SELECT manager_id FROM employees WHERE name = 'Grace Hopper') AND e.name != 'Grace Hopper';"),
    ("Find the total revenue by product category in 2025.", "SELECT p.category, SUM(s.amount) AS total_revenue FROM products p JOIN sales s ON p.id = s.product_id WHERE s.sale_date BETWEEN '2025-01-01' AND '2025-12-31' GROUP BY p.category;"),
    ("Find the department manager for Alice Vance's department.", "SELECT m.name FROM employees e JOIN departments d ON e.department_id = d.id JOIN employees m ON d.manager_id = m.id WHERE e.name = 'Alice Vance';"),
    ("Show the products whose total sales quantity exceeds 15 items.", "SELECT p.name, SUM(s.quantity) AS total_qty FROM products p JOIN sales s ON p.id = s.product_id GROUP BY p.name HAVING total_qty > 15;"),
    ("Find the most expensive product in the Hardware category.", "SELECT name, price FROM products WHERE category = 'Hardware' ORDER BY price DESC LIMIT 1;"),
    ("What is the average salary of employees hired in 2024?", "SELECT AVG(salary) FROM employees WHERE hire_date BETWEEN '2024-01-01' AND '2024-12-31';"),
    ("List all employees along with their manager's name.", "SELECT e.name AS employee, m.name AS manager FROM employees e LEFT JOIN employees m ON e.manager_id = m.id;")
]

# Validation/Test dataset (12 high-quality out-of-sample pairs)
TEST_DATA = [
    ("List all employees in department 1.", "SELECT * FROM employees WHERE department_id = 1;"),
    ("Which software products cost less than 500 dollars?", "SELECT * FROM products WHERE category = 'Software' AND price < 500;"),
    ("How many departments do not have a manager?", "SELECT COUNT(*) FROM departments WHERE manager_id IS NULL;"),
    ("Find the total salary budget for the Engineering department.", "SELECT SUM(e.salary) FROM employees e JOIN departments d ON e.department_id = d.id WHERE d.name = 'Engineering';"),
    ("What is the total quantity of CRM products sold in 2024?", "SELECT SUM(s.quantity) FROM sales s JOIN products p ON s.product_id = p.id WHERE p.name = 'Enterprise CRM' AND s.sale_date BETWEEN '2024-01-01' AND '2024-12-31';"),
    ("Show the employee with the highest salary in department 2.", "SELECT name, salary FROM employees WHERE department_id = 2 ORDER BY salary DESC LIMIT 1;"),
    ("List employees who report directly to Alice Vance.", "SELECT e.name FROM employees e JOIN employees m ON e.manager_id = m.id WHERE m.name = 'Alice Vance';"),
    ("Find the average sale amount for employee ID 9.", "SELECT AVG(amount) FROM sales WHERE employee_id = 9;"),
    ("Find the total revenue generated from software category products.", "SELECT SUM(s.amount) FROM sales s JOIN products p ON s.product_id = p.id WHERE p.category = 'Software';"),
    ("Which product has the lowest price?", "SELECT name, price FROM products ORDER BY price ASC LIMIT 1;"),
    ("Find the employee who generated the highest total revenue across all sales.", "SELECT e.name, SUM(s.amount) AS total_revenue FROM employees e JOIN sales s ON e.id = s.employee_id GROUP BY e.name ORDER BY total_revenue DESC LIMIT 1;"),
    ("List products sold in November 2024.", "SELECT DISTINCT p.name FROM products p JOIN sales s ON p.id = s.product_id WHERE s.sale_date BETWEEN '2024-11-01' AND '2024-11-30';")
]

def format_phi3_chat(question, sql_query):
    """
    Format data in the standard Phi-3 Prompt Template:
    <|user|>
    [Schema + Instruction]
    Question: {question}
    <|end|>
    <|assistant|>
    {sql_query}<|end|>
    """
    prompt = f"{SCHEMA_PROMPT}\n\nQuestion: {question}"
    
    # Return formatted chat message dictionary
    return {
        "messages": [
            {"role": "user", "content": prompt},
            {"role": "assistant", "content": sql_query}
        ]
    }

def generate_datasets(output_dir="data"):
    os.makedirs(output_dir, exist_ok=True)
    
    train_formatted = [format_phi3_chat(q, sql) for q, sql in TRAIN_DATA]
    test_formatted = [format_phi3_chat(q, sql) for q, sql in TEST_DATA]
    
    # Save training
    train_path = os.path.join(output_dir, "train_dataset.jsonl")
    with open(train_path, "w", encoding="utf-8") as f:
        for item in train_formatted:
            f.write(json.dumps(item) + "\n")
            
    # Save test
    test_path = os.path.join(output_dir, "test_dataset.jsonl")
    with open(test_path, "w", encoding="utf-8") as f:
        for item in test_formatted:
            f.write(json.dumps(item) + "\n")
            
    print(f"Datasets generated successfully:")
    print(f"  - Training dataset: {train_path} ({len(train_formatted)} records)")
    print(f"  - Testing dataset: {test_path} ({len(test_formatted)} records)")

if __name__ == "__main__":
    generate_datasets()