"""Caching utilities for Releso. To be used in SPOR steps."""

import json
import sqlite3

import numpy as np


class RelesoSporCache:
    """A simple SQLite cache for storing SPOR results.

    This cache is designed for use in Releso's SPOR steps, allowing for
    efficient storage and retrieval of key-value pairs where keys are strings
    and values are JSON-serializable dictionaries. The cache is backed by an
    SQLite database, which is created if it does not already exist. It is designed
    to be used both in single as well as multi environment setups. In addition,
    it can allow caching between different training runs. As long as the key and
    value pair is consistent, the cache can be reused across different runs.

    It is advisable to only cache the input and output if the most computationally
    expensive part of the SPOR step and not cache the objective and reward since
    these are usually not expensive to compute and can change between runs.

    Ensure that the key that you use a string that uniquely identifies the
    input parameters of the SPOR step. You can use the :py:meth:`make_cache_key`
    method to create a consistent key based on the input parameters. The function
    allows list and numpy array as inputs. The input will be flattened and rounded
    to a specified number of decimal places to ensure that the key is unique and
    consistent across runs and floating point arithmetic does not produce additional
    superfluous keys.

    Example usage:

    .. code-block:: python

        from releso.util.caching import RelesoSporCache

        def main(args, logger, func_data):
            # first call of function
            if func_data is None:
                # the data_keys and values are not the input parameters of the
                # function, but rather the data that you want to cache.
                func_data["cache"] = RelesoSporCache(
                    db_path="spor_cache.db",
                    example_data={"data_key1": "value1", "data_key2": "value2"},
                )
            if args.reset:
                # don't touch the cache, only reset the other func_data entries

            input_params = args.json_object["info"]["geometry_information"]
            key = RelesoSporCache.make_cache_key(input_params)
            cached_value = func_data["cache"].get(key)
            if cached_value is not None:
                # Use cached value
                pass
            else:
                # Perform expensive computation
                result = expensive_computation(input_params)
                # Store result in cache
                # The data you send here needs to be serializable by JSON.
                func_data["cache"].set(
                    key,
                    {"data_key1": result[0], "data_key2": result[1]},
                )


    Params:
        db_path (str): Path to the SQLite database file.
        This file will be created if it does not exist.
        The cache will store key-value pairs where keys are strings and values are JSON-serializable
        dictionaries.
    """

    @staticmethod
    def make_cache_key(
        key_array: np.ndarray | list[list[float]] | list[float],
        rounding: int = 6,
    ) -> str:
        """Create a cache key based on the provided arguments.

        This function was initially generated by AI.

        Args:
            key_array (np.ndarray | list[list[float]]): Array or list of control
                point y-coordinates to be used as a key for caching.
            rounding (int): Number of decimal places to round the values in the key.
                Default is 6.
        Returns:
            str: A JSON-serialized string that serves as a unique key for caching.
        """
        # Convert args and kwargs to a JSON-serializable string
        key_array = np.asarray(key_array)
        normalized_args = key_array.flatten().round(rounding).tolist()
        key = json.dumps(normalized_args, sort_keys=True)
        return key

    def __init__(self, db_path: str, example_data: dict):
        self.db_path = db_path

        self.keys = example_data.keys()
        self.value_accessor = ", ".join(f"{key}" for key in self.keys)
        self.value_question_mark = ", ".join("?" for _ in self.keys)

        self._initialize_db()

    def _initialize_db(self):
        """Initialize the SQLite database and create the cache table if it doesn't exist."""
        table_definitions = ", ".join(f"{key} TEXT" for key in self.keys)
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute(
                f"CREATE TABLE IF NOT EXISTS spor_cache (key TEXT PRIMARY KEY, {table_definitions})"
            )
            conn.commit()

    def set(self, key: str, value: dict):
        """Store a value in the cache.

        Args:
            key (str): The key to store the value under.
            value (dict): The value to store, must be a dictionary with keys matching
                          the example data keys.
        """
        if not isinstance(value, dict) or value.keys() != self.keys:
            raise ValueError(
                f"Value must be a dictionary with keys: {self.keys}"
            )
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute(
                f"INSERT OR REPLACE INTO spor_cache (key, {self.value_accessor}) VALUES (?, {self.value_question_mark})",
                (key, *[json.dumps(value[key]) for key in self.keys]),
            )
            conn.commit()

    def get(self, key: str) -> list[dict] | None:
        """Retrieve a value from the cache.

        Args:
            key (str): The key to retrieve the value for.

        Returns:
            list[dict] | None: The cached value as a list of dictionaries, or
            None if not found.
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute(
                    f"SELECT {self.value_accessor} FROM spor_cache WHERE key = ?",
                    (key,),
                )
                rows = cursor.fetchone()
                return [json.loads(row) for row in rows] if rows else None
        except sqlite3.OperationalError:
            return None
