"""Tests for benchmarking trace database access."""

from __future__ import annotations

import pickle
import sqlite3

import pytest

from codeflash_python.benchmarking._trace_db import (
    FUNCTION_CALLS_SCHEMA,
    TOTAL_TIME_SCHEMA,
    get_trace_total_run_time_ns,
    get_traced_arguments,
    sanitize_to_filename,
)


class TestSanitizeToFilename:
    """Tests for sanitize_to_filename."""

    def test_simple_string(self) -> None:
        """Simple alphanumeric string passes through."""
        assert "hello" == sanitize_to_filename("hello")

    def test_strips_special_characters(self) -> None:
        """Special characters are removed."""
        assert "ab" == sanitize_to_filename("a!@#$%^&*()b")

    def test_replaces_newlines(self) -> None:
        """Newlines become underscores."""
        assert "a_b" == sanitize_to_filename("a\nb")
        assert "a_b" == sanitize_to_filename("a\rb")

    def test_truncates_long_words(self) -> None:
        """More than 5 whitespace-separated words are truncated."""
        assert "a_b_c_d_e" == sanitize_to_filename("a b c d e f g")

    def test_truncates_at_100_chars(self) -> None:
        """Output is capped at 100 characters."""
        long = "a" * 200
        assert 100 == len(sanitize_to_filename(long))

    def test_empty_returns_untitled(self) -> None:
        """Empty or all-special input returns 'untitled'."""
        assert "untitled" == sanitize_to_filename("")
        assert "untitled" == sanitize_to_filename("!@#$")

    def test_strips_leading_trailing_dots_underscores(self) -> None:
        """Leading/trailing dots and underscores are removed."""
        assert "hello" == sanitize_to_filename("..hello__")

    def test_preserves_dots_in_middle(self) -> None:
        """Dots inside the string are preserved."""
        assert "file.py" == sanitize_to_filename("file.py")


class TestGetTracedArguments:
    """Tests for get_traced_arguments."""

    @pytest.fixture
    def _trace_db(self, tmp_path):
        """Create a trace database with sample call data."""
        db_path = tmp_path / "trace.db"
        con = sqlite3.connect(str(db_path))
        con.execute(FUNCTION_CALLS_SCHEMA)
        for i in range(5):
            blob = pickle.dumps({"arg": i})
            con.execute(
                "INSERT INTO function_calls VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                ("call", "my_func", None, "mod.py", 10, 0, i * 1000, blob),
            )
        con.commit()
        con.close()
        return db_path

    def test_yields_argument_blobs(self, _trace_db) -> None:
        """Call events yield their argument blobs."""
        results = list(
            get_traced_arguments(_trace_db, "my_func", "mod.py", num_to_get=3)
        )

        assert 3 == len(results)
        assert {"arg": 0} == pickle.loads(results[0])

    def test_class_filter(self, tmp_path) -> None:
        """class_name parameter filters by classname column."""
        db_path = tmp_path / "trace.db"
        con = sqlite3.connect(str(db_path))
        con.execute(FUNCTION_CALLS_SCHEMA)
        blob = pickle.dumps({"cls": True})
        con.execute(
            "INSERT INTO function_calls VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
            ("call", "method", "MyClass", "mod.py", 5, 0, 100, blob),
        )
        blob2 = pickle.dumps({"cls": False})
        con.execute(
            "INSERT INTO function_calls VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
            ("call", "method", "OtherClass", "mod.py", 5, 0, 200, blob2),
        )
        con.commit()
        con.close()

        results = list(
            get_traced_arguments(
                db_path,
                "method",
                "mod.py",
                class_name="MyClass",
            )
        )

        assert 1 == len(results)
        assert {"cls": True} == pickle.loads(results[0])

    def test_invalid_event_type_raises(self, tmp_path) -> None:
        """Non-call event types raise ValueError."""
        db_path = tmp_path / "trace.db"
        con = sqlite3.connect(str(db_path))
        con.execute(FUNCTION_CALLS_SCHEMA)
        con.execute(
            "INSERT INTO function_calls VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
            ("return", "my_func", None, "mod.py", 10, 0, 100, b""),
        )
        con.commit()
        con.close()

        with pytest.raises(ValueError, match="Invalid Trace event type"):
            list(
                get_traced_arguments(db_path, "my_func", "mod.py"),
            )

    def test_empty_db_yields_nothing(self, tmp_path) -> None:
        """An empty database yields no results."""
        db_path = tmp_path / "trace.db"
        con = sqlite3.connect(str(db_path))
        con.execute(FUNCTION_CALLS_SCHEMA)
        con.commit()
        con.close()

        results = list(
            get_traced_arguments(db_path, "my_func", "mod.py"),
        )

        assert [] == results


class TestGetTraceTotalRunTimeNs:
    """Tests for get_trace_total_run_time_ns."""

    def test_returns_stored_time(self, tmp_path) -> None:
        """Total time is read from the database."""
        db_path = tmp_path / "trace.db"
        con = sqlite3.connect(str(db_path))
        con.execute(TOTAL_TIME_SCHEMA)
        con.execute("INSERT INTO total_time VALUES (?)", (123456789,))
        con.commit()
        con.close()

        assert 123456789 == get_trace_total_run_time_ns(db_path)

    def test_missing_file_returns_zero(self, tmp_path) -> None:
        """Non-existent file returns 0."""
        assert 0 == get_trace_total_run_time_ns(
            tmp_path / "nonexistent.db",
        )

    def test_missing_table_returns_zero(self, tmp_path) -> None:
        """Database without total_time table returns 0."""
        db_path = tmp_path / "trace.db"
        con = sqlite3.connect(str(db_path))
        con.execute(FUNCTION_CALLS_SCHEMA)
        con.commit()
        con.close()

        assert 0 == get_trace_total_run_time_ns(db_path)

    def test_empty_table_returns_zero(self, tmp_path) -> None:
        """Empty total_time table returns 0."""
        db_path = tmp_path / "trace.db"
        con = sqlite3.connect(str(db_path))
        con.execute(TOTAL_TIME_SCHEMA)
        con.commit()
        con.close()

        assert 0 == get_trace_total_run_time_ns(db_path)
