from __future__ import annotations

import os
import sqlite3
import tempfile

# imports
import pytest  # used for our unit tests
from codeflash.benchmarking.plugin.plugin import CodeFlashBenchmarkPlugin

# ---- Helper functions for setting up the test DB ----

def create_benchmark_db(path):
    """Create a SQLite database at path with the correct schema."""
    conn = sqlite3.connect(path)
    cur = conn.cursor()
    cur.execute("""
        CREATE TABLE benchmark_timings (
            benchmark_module_path TEXT,
            benchmark_function_name TEXT,
            benchmark_line_number INTEGER,
            benchmark_time_ns INTEGER
        )
    """)
    conn.commit()
    conn.close()

def fetch_benchmark_timings(path):
    """Fetch all rows from the benchmark_timings table."""
    conn = sqlite3.connect(path)
    cur = conn.cursor()
    cur.execute("SELECT benchmark_module_path, benchmark_function_name, benchmark_line_number, benchmark_time_ns FROM benchmark_timings")
    rows = cur.fetchall()
    conn.close()
    return rows

# ---- Unit Tests ----

# ----------- BASIC TEST CASES -----------

def test_write_single_benchmark_timing(tmp_path):
    """Test writing a single benchmark timing entry."""
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    plugin.benchmark_timings = [("mod.py", "func", 42, 1000)]
    plugin.write_benchmark_timings()
    # Check DB contents
    rows = fetch_benchmark_timings(str(db_path))

def test_write_multiple_benchmark_timings(tmp_path):
    """Test writing multiple benchmark timing entries."""
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    plugin.benchmark_timings = [
        ("mod1.py", "func1", 10, 100),
        ("mod2.py", "func2", 20, 200),
        ("mod3.py", "func3", 30, 300),
    ]
    plugin.write_benchmark_timings()
    rows = fetch_benchmark_timings(str(db_path))

def test_write_benchmark_timings_twice(tmp_path):
    """Test writing benchmark timings in two separate calls."""
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    # First call
    plugin.benchmark_timings = [("mod.py", "func", 1, 111)]
    plugin.write_benchmark_timings()
    # Second call
    plugin.benchmark_timings = [("mod2.py", "func2", 2, 222)]
    plugin.write_benchmark_timings()
    rows = fetch_benchmark_timings(str(db_path))

def test_no_benchmark_timings_no_write(tmp_path):
    """Test that nothing is written if benchmark_timings is empty."""
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    plugin.benchmark_timings = []
    plugin.write_benchmark_timings()
    rows = fetch_benchmark_timings(str(db_path))

# ----------- EDGE TEST CASES -----------

def test_write_benchmark_with_null_like_values(tmp_path):
    """Test writing entries with empty strings and zero/negative numbers."""
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    plugin.benchmark_timings = [
        ("", "", 0, 0),
        ("mod.py", "func", -1, -100),
    ]
    plugin.write_benchmark_timings()
    rows = fetch_benchmark_timings(str(db_path))

def test_write_benchmark_timings_with_long_strings(tmp_path):
    """Test writing entries with very long strings."""
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    long_str = "x" * 1000
    plugin.benchmark_timings = [
        (long_str, long_str, 123, 456789),
    ]
    plugin.write_benchmark_timings()
    rows = fetch_benchmark_timings(str(db_path))

def test_write_benchmark_timings_invalid_db_path(tmp_path):
    """Test error handling when the DB path is invalid."""
    plugin = CodeFlashBenchmarkPlugin()
    # Set to a path in a non-existent directory
    plugin._trace_path = str(tmp_path / "nonexistent" / "bench.db")
    plugin.benchmark_timings = [("mod.py", "func", 1, 2)]
    with pytest.raises(sqlite3.OperationalError):
        plugin.write_benchmark_timings()

def test_write_benchmark_timings_with_existing_connection(tmp_path):
    """Test that an existing _connection is used and not replaced."""
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    plugin._connection = sqlite3.connect(str(db_path))
    plugin.benchmark_timings = [("mod.py", "func", 1, 2)]
    plugin.write_benchmark_timings()
    # Check that the connection is still open (should not be closed)
    try:
        plugin._connection.execute("SELECT 1")
    except sqlite3.ProgrammingError:
        pass
    rows = fetch_benchmark_timings(str(db_path))

def test_write_benchmark_timings_rollback_on_failure(tmp_path):
    """Test that rollback is called and no data is written if insert fails."""
    db_path = tmp_path / "bench.db"
    # Create a DB with the wrong schema (missing a column)
    conn = sqlite3.connect(str(db_path))
    cur = conn.cursor()
    cur.execute("""
        CREATE TABLE benchmark_timings (
            benchmark_module_path TEXT,
            benchmark_function_name TEXT
        )
    """)
    conn.commit()
    conn.close()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    plugin.benchmark_timings = [("mod.py", "func", 1, 2)]
    with pytest.raises(sqlite3.OperationalError):
        plugin.write_benchmark_timings()
    # Table should still be empty
    conn = sqlite3.connect(str(db_path))
    cur = conn.cursor()
    cur.execute("SELECT * FROM benchmark_timings")
    rows = cur.fetchall()
    conn.close()

# ----------- LARGE SCALE TEST CASES -----------

def test_write_large_number_of_benchmark_timings(tmp_path):
    """Test writing a large number (1000) of benchmark timings."""
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    N = 1000
    data = [(f"mod{i}.py", f"func{i}", i, i*100) for i in range(N)]
    plugin.benchmark_timings = list(data)
    plugin.write_benchmark_timings()
    rows = fetch_benchmark_timings(str(db_path))

def test_write_large_benchmark_timings_multiple_batches(tmp_path):
    """Test multiple large writes to the DB."""
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    N = 500
    # First batch
    data1 = [(f"modA{i}.py", f"funcA{i}", i, i*10) for i in range(N)]
    plugin.benchmark_timings = list(data1)
    plugin.write_benchmark_timings()
    # Second batch
    data2 = [(f"modB{i}.py", f"funcB{i}", i, i*20) for i in range(N)]
    plugin.benchmark_timings = list(data2)
    plugin.write_benchmark_timings()
    rows = fetch_benchmark_timings(str(db_path))

def test_write_benchmark_timings_performance(tmp_path):
    """Test that writing 1000 entries does not take excessive time."""
    import time
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    N = 1000
    plugin.benchmark_timings = [(f"mod{i}.py", f"func{i}", i, i*123) for i in range(N)]
    start = time.time()
    plugin.write_benchmark_timings()
    duration = time.time() - start
    rows = fetch_benchmark_timings(str(db_path))
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from __future__ import annotations

import os
import sqlite3
import tempfile

# imports
import pytest  # used for our unit tests
from codeflash.benchmarking.plugin.plugin import CodeFlashBenchmarkPlugin


# Helper function to create a temporary SQLite DB with the required table
def create_temp_db_with_benchmark_table():
    temp_db = tempfile.NamedTemporaryFile(delete=False)
    conn = sqlite3.connect(temp_db.name)
    cur = conn.cursor()
    # Create the table with the expected schema
    cur.execute(
        """
        CREATE TABLE benchmark_timings (
            benchmark_module_path TEXT,
            benchmark_function_name TEXT,
            benchmark_line_number INTEGER,
            benchmark_time_ns INTEGER
        )
        """
    )
    conn.commit()
    conn.close()
    return temp_db

# Helper function to read all rows from the benchmark_timings table
def read_benchmark_timings_from_db(db_path):
    conn = sqlite3.connect(db_path)
    cur = conn.cursor()
    cur.execute("SELECT benchmark_module_path, benchmark_function_name, benchmark_line_number, benchmark_time_ns FROM benchmark_timings")
    rows = cur.fetchall()
    conn.close()
    return rows

# -------------------------
# Basic Test Cases
# -------------------------

def test_write_single_benchmark_timing_basic():
    """Test writing a single benchmark timing entry to the database."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    plugin.benchmark_timings = [
        ("module.py", "func", 42, 123456789)
    ]
    plugin.write_benchmark_timings()
    # Check the DB contents
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

def test_write_multiple_benchmark_timings_basic():
    """Test writing multiple benchmark timings in one call."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    plugin.benchmark_timings = [
        ("a.py", "f1", 1, 100),
        ("b.py", "f2", 2, 200),
        ("c.py", "f3", 3, 300)
    ]
    plugin.write_benchmark_timings()
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

def test_write_benchmark_timings_twice_accumulates():
    """Test that calling write_benchmark_timings twice appends new rows."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name

    # First write
    plugin.benchmark_timings = [("x.py", "foo", 10, 999)]
    plugin.write_benchmark_timings()
    # Second write
    plugin.benchmark_timings = [("y.py", "bar", 20, 888)]
    plugin.write_benchmark_timings()

    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

# -------------------------
# Edge Test Cases
# -------------------------

def test_write_empty_benchmark_timings_list():
    """Test that nothing is written if benchmark_timings is empty."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    plugin.benchmark_timings = []
    plugin.write_benchmark_timings()
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

def test_write_benchmark_timings_none_connection():
    """Test that the function creates a connection if _connection is None."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    plugin._connection = None
    plugin.benchmark_timings = [("mod.py", "func", 5, 1000)]
    plugin.write_benchmark_timings()
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

def test_write_benchmark_timings_existing_connection():
    """Test that the function uses an existing connection if present."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    plugin._connection = sqlite3.connect(temp_db.name)
    plugin.benchmark_timings = [("mod2.py", "func2", 6, 2000)]
    plugin.write_benchmark_timings()
    rows = read_benchmark_timings_from_db(temp_db.name)
    plugin._connection.close()
    os.unlink(temp_db.name)

def test_write_benchmark_timings_unicode_and_long_strings():
    """Test writing entries with Unicode and very long strings."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    long_str = "a" * 500
    unicode_str = "模块.py"
    plugin.benchmark_timings = [
        (long_str, "f", 1, 1),
        (unicode_str, "函数", 2, 2)
    ]
    plugin.write_benchmark_timings()
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

def test_write_benchmark_timings_negative_and_zero_values():
    """Test writing entries with zero and negative numbers."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    plugin.benchmark_timings = [
        ("mod.py", "zero", 0, 0),
        ("mod.py", "neg", -1, -100)
    ]
    plugin.write_benchmark_timings()
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

def test_write_benchmark_timings_rollback_on_failure():
    """Test that if an error occurs, no partial data is committed."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    # Insert a valid row first
    plugin.benchmark_timings = [("ok.py", "ok", 1, 1)]
    plugin.write_benchmark_timings()
    # Now, insert an invalid row (too few columns)
    plugin.benchmark_timings = [("bad.py", "bad", 2)]  # Missing one value
    with pytest.raises(sqlite3.ProgrammingError):
        plugin.write_benchmark_timings()
    # The DB should still only have the first row
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

def test_write_benchmark_timings_no_table():
    """Test that writing to a DB without the table raises an error."""
    temp_db = tempfile.NamedTemporaryFile(delete=False)
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    plugin.benchmark_timings = [("no_table.py", "func", 1, 1)]
    with pytest.raises(sqlite3.OperationalError):
        plugin.write_benchmark_timings()
    os.unlink(temp_db.name)

def test_write_benchmark_timings_non_string_module_and_func():
    """Test that non-string types for module/function raise an error."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    plugin.benchmark_timings = [
        (123, 456, 7, 8)
    ]
    # SQLite will coerce types, so this should succeed, but let's check the DB
    plugin.write_benchmark_timings()
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

# -------------------------
# Large Scale Test Cases
# -------------------------

def test_write_benchmark_timings_large_scale():
    """Test writing a large number of benchmark timings (scalability)."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    N = 1000  # Upper limit as per instructions
    plugin.benchmark_timings = [
        (f"mod{i}.py", f"func{i}", i, i * 100) for i in range(N)
    ]
    plugin.write_benchmark_timings()
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

def test_write_benchmark_timings_large_scale_multiple_batches():
    """Test writing large batches in multiple calls to ensure accumulation."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    N = 500
    # First batch
    plugin.benchmark_timings = [
        (f"A{i}.py", f"F{i}", i, i) for i in range(N)
    ]
    plugin.write_benchmark_timings()
    # Second batch
    plugin.benchmark_timings = [
        (f"B{i}.py", f"G{i}", i, i*2) for i in range(N)
    ]
    plugin.write_benchmark_timings()
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.