"""Database schemas and trace data access."""

from __future__ import annotations

import re
import sqlite3
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Generator
    from typing import Any

FUNCTION_CALLS_SCHEMA: str = (
    "CREATE TABLE function_calls("
    "type TEXT, function TEXT, classname TEXT, "
    "filename TEXT, line_number INTEGER, "
    "last_frame_address INTEGER, "
    "time_ns INTEGER, args BLOB)"
)

TOTAL_TIME_SCHEMA: str = "CREATE TABLE total_time (time_ns INTEGER)"


def sanitize_to_filename(arg: str) -> str:
    """Sanitize a string for use as a filename."""
    arg = arg.replace("\n", "_").replace("\r", "_")
    parts = re.split(r"\s+", arg)
    if len(parts) > 5:  # noqa: PLR2004
        parts = parts[:5]
    arg = "_".join(parts)
    arg = re.sub(r"[^\w._]", "", arg)
    arg = arg.strip("._")
    arg = arg[:100]
    return arg or "untitled"


def get_traced_arguments(
    trace_file: str | Path,
    function_name: str,
    file_name: str,
    class_name: str | None = None,
    num_to_get: int = 25,
) -> Generator[Any, None, None]:
    """Yield pickled argument blobs from *trace_file*."""
    db = sqlite3.connect(str(trace_file))
    try:
        cur = db.cursor()
        if class_name is not None:
            cursor = cur.execute(
                "SELECT * FROM function_calls "
                "WHERE function = ? AND filename = ? "
                "AND classname = ? "
                "ORDER BY time_ns ASC LIMIT ?",
                (function_name, file_name, class_name, num_to_get),
            )
        else:
            cursor = cur.execute(
                "SELECT * FROM function_calls "
                "WHERE function = ? AND filename = ? "
                "ORDER BY time_ns ASC LIMIT ?",
                (function_name, file_name, num_to_get),
            )
        while (val := cursor.fetchone()) is not None:
            event_type = val[0]
            if event_type == "call":
                yield val[7]
            else:
                msg = "Invalid Trace event type"
                raise ValueError(msg)
    finally:
        db.close()


def get_trace_total_run_time_ns(
    trace_file_path: str | Path,
) -> int:
    """Return total run time in nanoseconds from a trace database."""
    trace_file_path = Path(trace_file_path)
    if not trace_file_path.is_file():
        return 0
    con = sqlite3.connect(str(trace_file_path))
    try:
        cur = con.cursor()
        try:
            time_data = cur.execute(
                "SELECT time_ns FROM total_time"
            ).fetchone()
        except sqlite3.OperationalError:
            return 0
    finally:
        con.close()
    time_data = time_data[0] if time_data else 0
    return int(time_data)
