"""Utility functions for interacting with the SQLite database."""
import io
import logging
import sqlite3
from typing import Any, List, Optional

import torch


def select_tensors(
        db_path: str,
        table_name: str,
        keys: List[str] = ['layer', 'pooling_method', 'tensor_dim', 'tensor'],
        sql_where: Optional[str] = None,
        ) -> List[Any]:
    """Select and return all tensors from the specified SQLite database and table.

    Args:
        db_path (str): Path to the SQLite database file.
        table_name (str): Name of the table to query.
        keys (List[str]): List of keys to select from the database.
        sql_where (str): Optional SQL WHERE clause to filter results.

    Returns:
        List[Any]: A list of tensors retrieved from the database.
    """
    if 'tensor' not in keys:
        logging.warning("'tensor' key should be included to retrieve tensors; automatically adding it.")
        keys.append('tensor')
    final_results = []
    with sqlite3.connect(db_path) as connection:
        cursor = connection.cursor()
        query = f'SELECT {", ".join(keys)} FROM {table_name}'
        if sql_where:
            assert sql_where.strip().lower().startswith('where'), "sql_where should start with 'WHERE'"
            query += f' {sql_where}'
        cursor.execute(query)
        results = cursor.fetchall()
        for row in results:
            result_item = {key: value for key, value in zip(keys, row)}
            result_item['tensor'] = torch.load(io.BytesIO(result_item['tensor']), map_location='cpu')
            final_results.append(result_item)
    return final_results