package com.github.catatafishen.agentbridge.session.db;

import com.github.catatafishen.agentbridge.bridge.ContextFileRef;
import com.github.catatafishen.agentbridge.bridge.EntryData;
import com.github.catatafishen.agentbridge.bridge.FileRef;
import com.github.catatafishen.agentbridge.services.hooks.HookStageResult;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.List;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;

class ConversationWriterTest {

    private Connection conn;
    private ConversationDatabase database;
    private ConversationWriter writer;

    @BeforeEach
    void setUp() throws Exception {
        Class.forName("org.sqlite.JDBC");
        conn = DriverManager.getConnection("jdbc:sqlite::memory:");
        try (Statement s = conn.createStatement()) {
            s.execute("PRAGMA foreign_keys = ON");
        }
        database = new ConversationDatabase();
        database.initializeWithConnection(conn);
        writer = new ConversationWriter(database);
    }

    @AfterEach
    void tearDown() throws Exception {
        if (conn != null) conn.close();
    }

    @Test
    void writesPromptAsTurnAndSession() throws Exception {
        EntryData.Prompt prompt = new EntryData.Prompt(
            "Fix the bug", "2026-01-01T10:00:00Z", null, "turn-1", "turn-1");

        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(prompt));

        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery("SELECT id, agent_name, client_id FROM sessions")) {
            assertTrue(rs.next());
            assertEquals("sess-1", rs.getString(1));
            assertEquals("Copilot", rs.getString(2));
            assertEquals("copilot", rs.getString(3));
        }
        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT id, session_id, prompt_text FROM turns")) {
            assertTrue(rs.next());
            assertEquals("turn-1", rs.getString(1));
            assertEquals("sess-1", rs.getString(2));
            assertEquals("Fix the bug", rs.getString(3));
        }
    }

    @Test
    void writesEventsLinkedToCurrentTurnWithSequenceNumbers() throws Exception {
        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(
            new EntryData.Prompt("Hi", "2026-01-01T10:00:00Z", null, "turn-1", "turn-1"),
            new EntryData.Text("hello", "2026-01-01T10:00:01Z", "Copilot", "gpt-5", "ev-text"),
            new EntryData.Thinking("hmm", "2026-01-01T10:00:02Z", "Copilot", "gpt-5", "ev-think")
        ));

        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT id, turn_id, sequence_num, event_type FROM events ORDER BY sequence_num")) {
            assertTrue(rs.next());
            assertEquals("ev-text", rs.getString(1));
            assertEquals("turn-1", rs.getString(2));
            assertEquals(0, rs.getInt(3));
            assertEquals("text", rs.getString(4));

            assertTrue(rs.next());
            assertEquals("ev-think", rs.getString(1));
            assertEquals(1, rs.getInt(3));
            assertEquals("thinking", rs.getString(4));
        }
        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery("SELECT content FROM text_events")) {
            assertTrue(rs.next());
            assertEquals("hello", rs.getString(1));
        }
    }

    @Test
    void writesToolCallWithIsMcpNullWhenNotYetCorrelated() throws Exception {
        EntryData.ToolCall tc = new EntryData.ToolCall(
            "read_file", "{\"path\":\"a.txt\"}", "read", "ok", "completed",
            null, "a.txt", false, null, null,
            "2026-01-01T10:00:01Z", "Copilot", "gpt-5", "ev-tc");

        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(
            new EntryData.Prompt("Hi", "2026-01-01T10:00:00Z", null, "turn-1", "turn-1"),
            tc
        ));

        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT tool_name, arguments, status, file_path, is_mcp, client_id "
                     + "FROM tool_call_events WHERE event_id = 'ev-tc'")) {
            assertTrue(rs.next());
            assertEquals("read_file", rs.getString(1));
            assertEquals("{\"path\":\"a.txt\"}", rs.getString(2));
            assertEquals("completed", rs.getString(3));
            assertEquals("a.txt", rs.getString(4));
            // pluginTool = null → not yet correlated with MCP, is_mcp unknown
            assertNull(rs.getObject(5));
            assertEquals("copilot", rs.getString(6));
        }

        writer.markToolCallMcp("ev-tc");
        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT is_mcp FROM tool_call_events WHERE event_id = 'ev-tc'")) {
            assertTrue(rs.next());
            assertEquals(1, rs.getInt(1));
        }
    }

    @Test
    void finaliseTurnFillsTotalsAndCommits() throws Exception {
        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(
            new EntryData.Prompt("Build feature", "2026-01-01T10:00:00Z", null,
                "turn-1", "turn-1"),
            new EntryData.TurnStats("turn-1", 5000L, 100L, 200L, 0.05, 3,
                12, 4, "gpt-5", "1.0", 0L, 0L, 0L, 0.0, 0, 0, 0,
                "2026-01-01T10:00:05Z", "stats-1",
                List.of("abc123", "def456"))
        ));

        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT ended_at, model, input_tokens, output_tokens, duration_ms, "
                     + "tool_call_count, lines_added, lines_removed FROM turns")) {
            assertTrue(rs.next());
            assertEquals("2026-01-01T10:00:05Z", rs.getString(1));
            assertEquals("gpt-5", rs.getString(2));
            assertEquals(100L, rs.getLong(3));
            assertEquals(200L, rs.getLong(4));
            assertEquals(5000L, rs.getLong(5));
            assertEquals(3, rs.getInt(6));
            assertEquals(12, rs.getInt(7));
            assertEquals(4, rs.getInt(8));
        }
        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT commit_hash FROM commits ORDER BY id")) {
            assertTrue(rs.next());
            assertEquals("abc123", rs.getString(1));
            assertTrue(rs.next());
            assertEquals("def456", rs.getString(1));
        }
    }

    @Test
    void writesContextFilesForCurrentTurn() throws Exception {
        EntryData.Prompt prompt = new EntryData.Prompt(
            "Look at this", "2026-01-01T10:00:00Z",
            List.of(new ContextFileRef("Foo.kt", "src/Foo.kt", 42)),
            "turn-1", "turn-1");

        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(
            prompt,
            new EntryData.ContextFiles(
                List.of(new FileRef("Bar.kt", "src/Bar.kt")), "ctx-1")
        ));

        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT file_name, file_path, file_line FROM turn_context_files "
                     + "WHERE turn_id = 'turn-1' ORDER BY id")) {
            assertTrue(rs.next());
            assertEquals("Foo.kt", rs.getString(1));
            assertEquals("src/Foo.kt", rs.getString(2));
            assertEquals(42, rs.getInt(3));
            assertTrue(rs.next());
            assertEquals("Bar.kt", rs.getString(1));
            assertEquals(0, rs.getInt(3));
        }
    }

    @Test
    void appendsTimestamplessEntryToExistingSession() throws Exception {
        // First call: create session with a timestamped prompt
        EntryData.Prompt prompt = new EntryData.Prompt(
            "Hello", "2026-01-01T10:00:00Z", null, "turn-1", "turn-1");
        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(prompt));

        // Second call: append only a ContextFiles entry (no timestamp by design).
        // Before the fix, this threw IllegalArgumentException because
        // extractStartedAt required a non-null timestamp in every batch.
        EntryData.ContextFiles contextFiles = new EntryData.ContextFiles(
            List.of(new FileRef("Image.png", "img/Image.png")), "ctx-1");
        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(contextFiles));

        // Verify the context file was written successfully
        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT file_name FROM turn_context_files WHERE turn_id = 'turn-1'")) {
            assertTrue(rs.next());
            assertEquals("Image.png", rs.getString(1));
        }
    }

    @Test
    void recordsHookExecution() throws Exception {
        // Standalone tool call — no turn.
        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(
            new EntryData.ToolCall(
                "delete_file", "{}", "fs", null, "pending",
                null, null, false, null, null,
                "2026-01-01T10:00:01Z", "", "", "ev-tc")
        ));
        // First entry was a tool call without prior prompt → cursor.turnId is null.
        // This proves standalone events work.

        writer.recordHookExecution(new ConversationWriter.HookExecutionRecord(
            "ev-tc", "permission", "block-deletes",
            "/usr/bin/policy --check", 0, 25L,
            "{\"tool\":\"delete_file\"}", "{\"allow\":false}",
            "deny", "delete blocked by policy"));

        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT trigger_kind, entry_id, exit_code, duration_ms, outcome, outcome_reason "
                     + "FROM hook_executions")) {
            assertTrue(rs.next());
            assertEquals("permission", rs.getString(1));
            assertEquals("block-deletes", rs.getString(2));
            assertEquals(0, rs.getInt(3));
            assertEquals(25L, rs.getLong(4));
            assertEquals("deny", rs.getString(5));
            assertEquals("delete blocked by policy", rs.getString(6));
        }

        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT turn_id FROM events WHERE id = 'ev-tc'")) {
            assertTrue(rs.next());
            // Standalone event: turn_id is NULL.
            rs.getString(1);
            assertTrue(rs.wasNull(), "Standalone tool call should have NULL turn_id");
        }
    }

    @Test
    void skipsUnsentNudges() throws Exception {
        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(
            new EntryData.Prompt("Hi", "2026-01-01T10:00:00Z", null, "turn-1", "turn-1"),
            new EntryData.Nudge("don't forget X", "n-1", false,
                "2026-01-01T10:00:01Z", "ev-nudge-pending"),
            new EntryData.Nudge("send help", "n-2", true,
                "2026-01-01T10:00:02Z", "ev-nudge-sent")
        ));

        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT event_id, text FROM nudge_events")) {
            assertTrue(rs.next());
            assertEquals("ev-nudge-sent", rs.getString(1));
            assertEquals("send help", rs.getString(2));
            org.junit.jupiter.api.Assertions.assertFalse(rs.next(),
                "Pending nudges must not be persisted");
        }
    }

    @Test
    void restoresCursorFromExistingRows() throws Exception {
        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(
            new EntryData.Prompt("Hi", "2026-01-01T10:00:00Z", null, "turn-1", "turn-1"),
            new EntryData.Text("a", "2026-01-01T10:00:01Z", "", "", "ev-1"),
            new EntryData.Text("b", "2026-01-01T10:00:02Z", "", "", "ev-2")
        ));

        // Fresh writer (simulates restart): cursor empty.
        ConversationWriter fresh = new ConversationWriter(database);
        fresh.restoreCursor("sess-1");
        fresh.recordEntries("sess-1", "Copilot", "copilot", List.of(
            new EntryData.Text("c", "2026-01-01T10:00:03Z", "", "", "ev-3")
        ));

        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT sequence_num, turn_id FROM events WHERE id = 'ev-3'")) {
            assertTrue(rs.next());
            assertEquals(2, rs.getInt(1), "Cursor should resume sequence at 2");
            assertEquals("turn-1", rs.getString(2));
        }
    }

    @Test
    void ignoresStatusAndSessionSeparator() throws Exception {
        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(
            new EntryData.Prompt("Hi", "2026-01-01T10:00:00Z", null, "turn-1", "turn-1"),
            new EntryData.Status("⚠", "stuck", "ev-status"),
            new EntryData.SessionSeparator("2026-01-01T10:00:00Z", "Copilot", "ev-sep")
        ));
        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery("SELECT COUNT(*) FROM events")) {
            assertTrue(rs.next());
            assertEquals(0, rs.getInt(1), "Status and SessionSeparator should not be persisted");
        }
    }

    @Test
    void mcpFlagSurvivesReReads() throws Exception {
        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(
            new EntryData.Prompt("Hi", "2026-01-01T10:00:00Z", null, "turn-1", "turn-1"),
            new EntryData.ToolCall("write_file", null, "fs", null, null, null, null,
                false, null, null, "2026-01-01T10:00:01Z", "", "", "ev-tc-mcp")
        ));
        writer.markToolCallMcp("ev-tc-mcp");

        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT COUNT(*) FROM tool_call_events WHERE is_mcp = 1")) {
            assertTrue(rs.next());
            assertEquals(1, rs.getInt(1));
        }
        // Sanity: marking unknown id is a no-op, not an error.
        writer.markToolCallMcp("does-not-exist");
        assertNotNull(database.getConnection(), "Connection must survive a no-op markToolCallMcp call");
    }

    @Test
    void isMcpSetWhenPluginToolSet_nullOtherwise() throws Exception {
        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(
            new EntryData.Prompt("Hi", "2026-01-01T10:00:00Z", null, "turn-1", "turn-1"),
            // Correlated MCP tool — pluginTool holds the confirmed tool name
            new EntryData.ToolCall("agentbridge-read_file", null, "fs", null, null, null, null,
                false, null, "read_file", "2026-01-01T10:00:01Z", "", "", "ev-mcp"),
            // Also correlated (OpenCode style client prefix, different pluginTool)
            new EntryData.ToolCall("agentbridge_search_text", null, "search", null, null, null, null,
                false, null, "search_text", "2026-01-01T10:00:02Z", "", "", "ev-mcp2"),
            // Built-in / not yet correlated — pluginTool is null
            new EntryData.ToolCall("bash", null, "shell", null, null, null, null,
                false, null, null, "2026-01-01T10:00:03Z", "", "", "ev-builtin")
        ));

        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT event_id, is_mcp FROM tool_call_events ORDER BY event_id")) {
            assertTrue(rs.next());
            assertEquals("ev-builtin", rs.getString(1));
            assertNull(rs.getObject(2), "bash with null pluginTool → is_mcp unknown (NULL)");
            assertTrue(rs.next());
            assertEquals("ev-mcp", rs.getString(1));
            assertEquals(1, rs.getInt(2), "correlated MCP tool should be flagged as is_mcp=1");
            assertTrue(rs.next());
            assertEquals("ev-mcp2", rs.getString(1));
            assertEquals(1, rs.getInt(2), "correlated MCP tool should be flagged as is_mcp=1");
        }
    }

    @Test
    void enrichToolCallStatsUpdatesExistingRow() throws Exception {
        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(
            new EntryData.Prompt("Hi", "2026-01-01T10:00:00Z", null, "turn-1", "turn-1"),
            new EntryData.ToolCall("agentbridge-read_file", null, "fs", null, null, null, null,
                false, null, null, "2026-01-01T10:00:01Z", "", "", "ev-enrich")
        ));

        writer.enrichToolCallStats(new ToolCallStatsEnrichment(
            "ev-enrich", 256, 1024, 150, true, null, "file", "Read File"));

        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT input_size_bytes, output_size_bytes, duration_ms, success, category, "
                     + "tool_name, display_name "
                     + "FROM tool_call_events WHERE event_id = 'ev-enrich'")) {
            assertTrue(rs.next());
            assertEquals(256, rs.getLong(1));
            assertEquals(1024, rs.getLong(2));
            assertEquals(150, rs.getLong(3));
            assertEquals(1, rs.getInt(4));
            assertEquals("file", rs.getString(5));
            // tool_name must be stripped of the ACP prefix; display_name set by enrichment
            assertEquals("read_file", rs.getString(6));
            assertEquals("Read File", rs.getString(7));
        }
    }

    @Test
    void enrichToolCallStatsWithFailure() throws Exception {
        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(
            new EntryData.Prompt("Hi", "2026-01-01T10:00:00Z", null, "turn-1", "turn-1"),
            new EntryData.ToolCall("agentbridge-run_command", null, "shell", null, null, null, null,
                false, null, null, "2026-01-01T10:00:01Z", "", "", "ev-fail")
        ));

        writer.enrichToolCallStats(new ToolCallStatsEnrichment(
            "ev-fail", 100, 500, 3000, false, "timeout expired", "shell", null));

        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT success, error_message FROM tool_call_events WHERE event_id = 'ev-fail'")) {
            assertTrue(rs.next());
            assertEquals(0, rs.getInt(1));
            assertEquals("timeout expired", rs.getString(2));
        }
    }

    @Test
    void recordHookStagesInsertsRows() throws Exception {
        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(
            new EntryData.Prompt("Hi", "2026-01-01T10:00:00Z", null, "turn-1", "turn-1"),
            new EntryData.ToolCall("agentbridge-run_command", null, "shell", null, null, null, null,
                false, null, null, "2026-01-01T10:00:01Z", "", "", "ev-hooks")
        ));

        List<HookStageResult> stages = List.of(
            new HookStageResult("permission", "check-auth.sh", "allowed", 5, null),
            new HookStageResult("pre", "inject-env.sh", "modified", 12, "added GH_TOKEN")
        );
        writer.recordHookStages("ev-hooks", stages);

        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT trigger_kind, entry_id, outcome, duration_ms, outcome_reason "
                     + "FROM hook_executions WHERE tool_event_id = 'ev-hooks' "
                     + "ORDER BY trigger_kind")) {
            assertTrue(rs.next());
            assertEquals("permission", rs.getString(1));
            assertEquals("check-auth.sh", rs.getString(2));
            assertEquals("allowed", rs.getString(3));
            assertEquals(5, rs.getLong(4));
            assertNull(rs.getString(5));

            assertTrue(rs.next());
            assertEquals("pre", rs.getString(1));
            assertEquals("inject-env.sh", rs.getString(2));
            assertEquals("modified", rs.getString(3));
            assertEquals(12, rs.getLong(4));
            assertEquals("added GH_TOKEN", rs.getString(5));
        }
    }

    @Test
    void updateToolCallCompletionPersistsResultAfterEarlyInsert() throws Exception {
        // Simulate the race: tool call inserted early with null result
        EntryData.ToolCall tc = new EntryData.ToolCall(
            "read_file", "{\"path\":\"a.txt\"}", "read", null, "running",
            null, null, false, null, null,
            "2026-01-01T10:00:01Z", "Copilot", "gpt-5", "ev-tc-update");

        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(
            new EntryData.Prompt("Hi", "2026-01-01T10:00:00Z", null, "turn-1", "turn-1"),
            tc
        ));

        // Verify it was inserted with null result and running status
        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT result, status FROM tool_call_events WHERE event_id = 'ev-tc-update'")) {
            assertTrue(rs.next());
            assertNull(rs.getString(1));
            assertEquals("running", rs.getString(2));
        }

        // Now the tool call completes — call the UPDATE method
        writer.updateToolCallCompletion("ev-tc-update", "File content here", "completed", false, null);

        // Verify result and status are now persisted
        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT result, status, auto_denied FROM tool_call_events WHERE event_id = 'ev-tc-update'")) {
            assertTrue(rs.next());
            assertEquals("File content here", rs.getString(1));
            assertEquals("completed", rs.getString(2));
            assertEquals(0, rs.getInt(3));
        }
    }

    @Test
    void updateSubAgentCompletionPersistsResultAfterEarlyInsert() throws Exception {
        // Simulate the race: sub-agent inserted early with null result
        EntryData.SubAgent sa = new EntryData.SubAgent(
            "explore", "Find implementations", "Search for X in the codebase",
            null, "running", 0, null, false, null,
            "2026-01-01T10:00:01Z", "Copilot", "gpt-5", "ev-sa-update");

        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(
            new EntryData.Prompt("Hi", "2026-01-01T10:00:00Z", null, "turn-1", "turn-1"),
            sa
        ));

        // Verify it was inserted with null result
        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT result_text, status FROM sub_agent_events WHERE event_id = 'ev-sa-update'")) {
            assertTrue(rs.next());
            assertNull(rs.getString(1));
            assertEquals("running", rs.getString(2));
        }

        // Sub-agent completes with actual result text
        writer.updateSubAgentCompletion(
            "ev-sa-update", "Found 3 implementations in src/", "completed", false, null);

        // Verify result and status are persisted
        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT result_text, status FROM sub_agent_events WHERE event_id = 'ev-sa-update'")) {
            assertTrue(rs.next());
            assertEquals("Found 3 implementations in src/", rs.getString(1));
            assertEquals("completed", rs.getString(2));
        }
    }

    @Test
    void updateToolCallCompletionIsNoOpWhenRowMissing() {
        // UPDATE on non-existent row should silently succeed (0 rows affected)
        writer.updateToolCallCompletion("nonexistent-id", "result", "completed", false, null);
        // If we get here without exception, the test passes
        assertNotNull(database.getConnection(), "Connection should survive no-op update");
    }

    @Test
    void updateToolCallCompletionPreservesExistingResult() throws Exception {
        // Insert with a result already set
        EntryData.ToolCall tc = new EntryData.ToolCall(
            "read_file", "{}", "read", "original content", "completed",
            null, null, false, null, null,
            "2026-01-01T10:00:01Z", "Copilot", "gpt-5", "ev-tc-existing");

        writer.recordEntries("sess-1", "Copilot", "copilot", List.of(
            new EntryData.Prompt("Hi", "2026-01-01T10:00:00Z", null, "turn-1", "turn-1"),
            tc
        ));

        // Calling update with null result should preserve existing via COALESCE
        writer.updateToolCallCompletion("ev-tc-existing", null, "completed", false, null);

        try (Statement s = conn.createStatement();
             ResultSet rs = s.executeQuery(
                 "SELECT result FROM tool_call_events WHERE event_id = 'ev-tc-existing'")) {
            assertTrue(rs.next());
            assertEquals("original content", rs.getString(1));
        }
    }

    @Test
    void updateSessionTitleUpdatesDisplayName() throws Exception {
        writer.recordEntries("sess-title", "Copilot", "copilot", List.of(
            new EntryData.Prompt("Hello", "2026-01-01T10:00:00Z", null, "t1", "t1")
        ));

        writer.updateSessionTitle("sess-title", "My Custom Title");

        try (var ps = conn.prepareStatement("SELECT display_name FROM sessions WHERE id = ?")) {
            ps.setString(1, "sess-title");
            try (var rs = ps.executeQuery()) {
                assertTrue(rs.next());
                assertEquals("My Custom Title", rs.getString(1));
            }
        }
    }

    @Test
    void updateSessionTitleSkipsBlankTitle() throws Exception {
        writer.recordEntries("sess-blank", "Copilot", "copilot", List.of(
            new EntryData.Prompt("Hello", "2026-01-01T10:00:00Z", null, "t1", "t1")
        ));
        // First set a real title
        writer.updateSessionTitle("sess-blank", "Real Title");
        // Then try to blank it — should be a no-op
        writer.updateSessionTitle("sess-blank", "   ");

        try (var ps = conn.prepareStatement("SELECT display_name FROM sessions WHERE id = ?")) {
            ps.setString(1, "sess-blank");
            try (var rs = ps.executeQuery()) {
                assertTrue(rs.next());
                assertEquals("Real Title", rs.getString(1));
            }
        }
    }

    @Test
    void markToolCallNonMcpSetsIsMcpToZero() throws Exception {
        writer.recordEntries("sess-mcp", "Copilot", "copilot", List.of(
            new EntryData.Prompt("Hi", "2026-01-01T10:00:00Z", null, "t1", "t1"),
            new EntryData.ToolCall(
                "custom_tool", "{}", "custom", null, "completed",
                null, null, false, null, null,
                "2026-01-01T10:00:01Z", "Copilot", "gpt-5", "ev-non-mcp")
        ));

        // Initially is_mcp should be NULL (pending correlation)
        try (var ps = conn.prepareStatement(
            "SELECT is_mcp FROM tool_call_events WHERE event_id = ?")) {
            ps.setString(1, "ev-non-mcp");
            try (var rs = ps.executeQuery()) {
                assertTrue(rs.next());
                rs.getInt(1);
                assertTrue(rs.wasNull(), "is_mcp should be NULL initially");
            }
        }

        writer.markToolCallNonMcp("ev-non-mcp");

        try (var ps = conn.prepareStatement(
            "SELECT is_mcp FROM tool_call_events WHERE event_id = ?")) {
            ps.setString(1, "ev-non-mcp");
            try (var rs = ps.executeQuery()) {
                assertTrue(rs.next());
                assertEquals(0, rs.getInt(1));
            }
        }
    }

    @Test
    void markToolCallNonMcpDoesNotOverwriteExistingMcp() throws Exception {
        writer.recordEntries("sess-mcp2", "Copilot", "copilot", List.of(
            new EntryData.Prompt("Hi", "2026-01-01T10:00:00Z", null, "t2", "t2"),
            new EntryData.ToolCall(
                "read_file", "{}", "read", null, "completed",
                null, null, false, null, null,
                "2026-01-01T10:00:01Z", "Copilot", "gpt-5", "ev-already-mcp")
        ));

        // Mark it as MCP first
        writer.markToolCallMcp("ev-already-mcp");

        // Now try to mark it as non-MCP — the SQL WHERE clause prevents overwriting (is_mcp IS NULL)
        writer.markToolCallNonMcp("ev-already-mcp");

        try (var ps = conn.prepareStatement(
            "SELECT is_mcp FROM tool_call_events WHERE event_id = ?")) {
            ps.setString(1, "ev-already-mcp");
            try (var rs = ps.executeQuery()) {
                assertTrue(rs.next());
                assertEquals(1, rs.getInt(1), "markToolCallNonMcp should not overwrite existing MCP=1");
            }
        }
    }
}
