SwiftDBAI/Tests/SwiftDBAITests/PromptBuilderTests.swift

// PromptBuilderTests.swift
// SwiftDBAI

import Testing
@testable import SwiftDBAI

@Suite("PromptBuilder")
struct PromptBuilderTests {

    // MARK: - Helpers

    /// Creates a sample schema for testing.
    private func makeSampleSchema() -> DatabaseSchema {
        let usersTable = TableSchema(
            name: "users",
            columns: [
                ColumnSchema(cid: 0, name: "id", type: "INTEGER", isNotNull: true, defaultValue: nil, isPrimaryKey: true),
                ColumnSchema(cid: 1, name: "name", type: "TEXT", isNotNull: true, defaultValue: nil, isPrimaryKey: false),
                ColumnSchema(cid: 2, name: "email", type: "TEXT", isNotNull: false, defaultValue: nil, isPrimaryKey: false),
                ColumnSchema(cid: 3, name: "created_at", type: "TEXT", isNotNull: false, defaultValue: "CURRENT_TIMESTAMP", isPrimaryKey: false),
            ],
            primaryKey: ["id"],
            foreignKeys: [],
            indexes: [
                IndexSchema(name: "idx_users_email", isUnique: true, columns: ["email"])
            ]
        )

        let ordersTable = TableSchema(
            name: "orders",
            columns: [
                ColumnSchema(cid: 0, name: "id", type: "INTEGER", isNotNull: true, defaultValue: nil, isPrimaryKey: true),
                ColumnSchema(cid: 1, name: "user_id", type: "INTEGER", isNotNull: true, defaultValue: nil, isPrimaryKey: false),
                ColumnSchema(cid: 2, name: "total", type: "REAL", isNotNull: true, defaultValue: nil, isPrimaryKey: false),
                ColumnSchema(cid: 3, name: "status", type: "TEXT", isNotNull: true, defaultValue: "'pending'", isPrimaryKey: false),
            ],
            primaryKey: ["id"],
            foreignKeys: [
                ForeignKeySchema(fromColumn: "user_id", toTable: "users", toColumn: "id", onUpdate: "NO ACTION", onDelete: "CASCADE")
            ],
            indexes: []
        )

        return DatabaseSchema(
            tables: ["users": usersTable, "orders": ordersTable],
            tableNames: ["users", "orders"]
        )
    }

    private func makeEmptySchema() -> DatabaseSchema {
        DatabaseSchema(tables: [:], tableNames: [])
    }

    // MARK: - System Instructions Tests

    @Test("System instructions contain role section")
    func systemInstructionsContainRole() {
        let builder = PromptBuilder(schema: makeSampleSchema())
        let instructions = builder.buildSystemInstructions()

        #expect(instructions.contains("ROLE"))
        #expect(instructions.contains("SQL assistant"))
        #expect(instructions.contains("SQLite database"))
    }

    @Test("System instructions contain schema")
    func systemInstructionsContainSchema() {
        let builder = PromptBuilder(schema: makeSampleSchema())
        let instructions = builder.buildSystemInstructions()

        #expect(instructions.contains("DATABASE SCHEMA"))
        #expect(instructions.contains("TABLE users"))
        #expect(instructions.contains("TABLE orders"))
        #expect(instructions.contains("name TEXT"))
        #expect(instructions.contains("email TEXT"))
    }

    @Test("System instructions contain foreign keys from schema")
    func systemInstructionsContainForeignKeys() {
        let builder = PromptBuilder(schema: makeSampleSchema())
        let instructions = builder.buildSystemInstructions()

        #expect(instructions.contains("FOREIGN KEY"))
        #expect(instructions.contains("REFERENCES users(id)"))
    }

    @Test("System instructions contain SQL generation rules")
    func systemInstructionsContainRules() {
        let builder = PromptBuilder(schema: makeSampleSchema())
        let instructions = builder.buildSystemInstructions()

        #expect(instructions.contains("SQL GENERATION RULES"))
        #expect(instructions.contains("Use ONLY the tables and columns"))
        #expect(instructions.contains("Never generate DDL"))
    }

    @Test("System instructions contain output format section")
    func systemInstructionsContainOutputFormat() {
        let builder = PromptBuilder(schema: makeSampleSchema())
        let instructions = builder.buildSystemInstructions()

        #expect(instructions.contains("OUTPUT FORMAT"))
    }

    @Test("Default allowlist is read-only")
    func defaultAllowlistIsReadOnly() {
        let builder = PromptBuilder(schema: makeSampleSchema())
        let instructions = builder.buildSystemInstructions()

        #expect(instructions.contains("ONLY generate SELECT queries"))
        #expect(instructions.contains("No data modifications"))
    }

    @Test("Standard allowlist shows correct operations")
    func standardAllowlistInstructions() {
        let builder = PromptBuilder(schema: makeSampleSchema(), allowlist: .standard)
        let instructions = builder.buildSystemInstructions()

        #expect(instructions.contains("INSERT"))
        #expect(instructions.contains("SELECT"))
        #expect(instructions.contains("UPDATE"))
    }

    @Test("Unrestricted allowlist warns about DELETE")
    func unrestrictedAllowlistWarnsAboutDelete() {
        let builder = PromptBuilder(schema: makeSampleSchema(), allowlist: .unrestricted)
        let instructions = builder.buildSystemInstructions()

        #expect(instructions.contains("DELETE"))
        #expect(instructions.contains("destructive"))
        #expect(instructions.contains("confirmation"))
    }

    @Test("Additional context is appended")
    func additionalContextAppended() {
        let builder = PromptBuilder(
            schema: makeSampleSchema(),
            additionalContext: "All dates are stored in ISO 8601 format."
        )
        let instructions = builder.buildSystemInstructions()

        #expect(instructions.contains("ADDITIONAL CONTEXT"))
        #expect(instructions.contains("ISO 8601"))
    }

    @Test("No additional context section when nil")
    func noAdditionalContextWhenNil() {
        let builder = PromptBuilder(schema: makeSampleSchema())
        let instructions = builder.buildSystemInstructions()

        #expect(!instructions.contains("ADDITIONAL CONTEXT"))
    }

    @Test("No additional context section when empty string")
    func noAdditionalContextWhenEmpty() {
        let builder = PromptBuilder(schema: makeSampleSchema(), additionalContext: "")
        let instructions = builder.buildSystemInstructions()

        #expect(!instructions.contains("ADDITIONAL CONTEXT"))
    }

    @Test("Empty schema produces valid instructions")
    func emptySchemaProducesValidInstructions() {
        let builder = PromptBuilder(schema: makeEmptySchema())
        let instructions = builder.buildSystemInstructions()

        #expect(instructions.contains("ROLE"))
        #expect(instructions.contains("SQL GENERATION RULES"))
        // Schema section should still be present, just empty
        #expect(instructions.contains("DATABASE SCHEMA"))
    }

    // MARK: - User Prompt Tests

    @Test("User prompt passes through question directly")
    func userPromptPassesThrough() {
        let builder = PromptBuilder(schema: makeSampleSchema())
        let prompt = builder.buildUserPrompt("How many users signed up this week?")

        #expect(prompt == "How many users signed up this week?")
    }

    // MARK: - Follow-up Prompt Tests

    @Test("Follow-up prompt includes previous context")
    func followUpPromptIncludesPreviousContext() {
        let builder = PromptBuilder(schema: makeSampleSchema())
        let prompt = builder.buildFollowUpPrompt(
            "Now sort them by name",
            previousSQL: "SELECT * FROM users WHERE created_at > date('now', '-7 days')",
            previousResultSummary: "Found 42 users who signed up this week"
        )

        #expect(prompt.contains("Previous query:"))
        #expect(prompt.contains("SELECT * FROM users"))
        #expect(prompt.contains("Previous result:"))
        #expect(prompt.contains("42 users"))
        #expect(prompt.contains("Follow-up question:"))
        #expect(prompt.contains("sort them by name"))
    }

    // MARK: - Schema Description Quality

    @Test("Schema includes column types and constraints")
    func schemaIncludesColumnDetails() {
        let builder = PromptBuilder(schema: makeSampleSchema())
        let instructions = builder.buildSystemInstructions()

        // Should include type info
        #expect(instructions.contains("INTEGER"))
        #expect(instructions.contains("TEXT"))
        #expect(instructions.contains("REAL"))

        // Should include constraints
        #expect(instructions.contains("NOT NULL"))
        #expect(instructions.contains("PRIMARY KEY"))
    }

    @Test("Schema includes index information")
    func schemaIncludesIndexes() {
        let builder = PromptBuilder(schema: makeSampleSchema())
        let instructions = builder.buildSystemInstructions()

        #expect(instructions.contains("INDEX"))
        #expect(instructions.contains("idx_users_email"))
    }

    // MARK: - Sendable Conformance

    @Test("PromptBuilder is Sendable")
    func promptBuilderIsSendable() async {
        let builder = PromptBuilder(schema: makeSampleSchema())

        // Verify it can be sent across concurrency boundaries
        let instructions = await Task.detached {
            builder.buildSystemInstructions()
        }.value

        #expect(instructions.contains("ROLE"))
    }

    // MARK: - Custom Allowlist

    @Test("Custom allowlist with select and delete only")
    func customAllowlist() {
        let allowlist = OperationAllowlist([.select, .delete])
        let builder = PromptBuilder(schema: makeSampleSchema(), allowlist: allowlist)
        let instructions = builder.buildSystemInstructions()

        #expect(instructions.contains("DELETE"))
        #expect(instructions.contains("SELECT"))
        #expect(instructions.contains("destructive"))
    }
}