import sqlite3
import unittest

from netscaler_triage.correlation import correlate_database


class CorrelationTests(unittest.TestCase):
    def setUp(self) -> None:
        self.db = sqlite3.connect(":memory:")
        self.db.executescript("""
        CREATE TABLE events(id INTEGER PRIMARY KEY,appliance TEXT,source TEXT,
                            line INTEGER,timestamp TEXT,kind TEXT,data TEXT);
        CREATE TABLE event_keys(event_id INTEGER,appliance TEXT,key_type TEXT,
                                key_value TEXT,timestamp TEXT);
        """)

    def tearDown(self) -> None:
        self.db.close()

    def add(
        self,
        number: int,
        stamp: str | None,
        source: str = "access",
        device: str = "one",
        key: str = "source_ip",
        value: str = "192.0.2.1",
    ) -> None:
        self.db.execute(
            "INSERT INTO events VALUES(?,?,?,?,?,?,?)",
            (number, device, source, number, stamp, "http", "{}"),
        )
        self.db.execute(
            "INSERT INTO event_keys VALUES(?,?,?,?,?)", (number, device, key, value, stamp)
        )

    def test_nearest_other_log_and_device_isolation(self) -> None:
        self.add(1, "2026-01-01T00:00:00Z", "auth")
        self.add(2, "2026-01-01T00:00:01Z")
        self.add(3, "2026-01-01T00:00:02Z")
        self.add(4, "2026-01-01T00:00:03Z", device="two")
        result = correlate_database(self.db)
        self.assertEqual(result["links"], 2)
        self.assertEqual(
            self.db.execute("SELECT left_id,right_id FROM correlations").fetchall(),
            [(1, 2), (1, 3)],
        )

    def test_times_limits_and_rerun(self) -> None:
        self.add(1, "2026-01-01T01:00:00+01:00", "auth")
        self.add(2, "2026-01-01T00:00:01Z")
        self.add(3, "2026-01-01T00:00:02Z")
        self.add(4, None)
        self.add(5, "2026-01-01T00:00:03")
        result = correlate_database(self.db, max_links=1)
        self.assertTrue(result["truncated"])
        result = correlate_database(self.db)
        self.assertEqual(result["links"], 2)
        self.assertEqual(result["invalid_time_keys"], 2)
        self.assertEqual(correlate_database(self.db, window_seconds=0)["links"], 0)

    def test_exact_keys_and_no_same_source_links(self) -> None:
        self.add(1, "2026-01-01T00:00:00Z")
        self.add(2, "2026-01-01T00:00:01Z")
        self.add(3, "2026-01-01T00:00:02Z", "auth", value="192.0.2.10")
        self.assertEqual(correlate_database(self.db)["links"], 0)


if __name__ == "__main__":
    unittest.main()