import logging
from typing import Dict, List, Optional
from uuid import UUID

import pandas as pd
from dataiku import Dataset, SQLExecutor2
from dataiku.sql import Dialects

from solutions.graph.models import GraphId
from solutions.llm.models import LoggingDatasetModel

DIALECT_VALUES = [
    v for k, v in Dialects.__dict__.items()
    if not k.startswith('_') and isinstance(v, str)
]

logger = logging.getLogger(__name__)

class LoggingDataset:
    """
    Manages logging user queries to a Dataiku dataset, which acts as a database table.

    This class handles the initialization of a connection to a dataset and provides
    methods to log queries and manage the dataset's schema.
    """

    def __init__(self) -> None:
        """Initializes the LoggingDataset instance with default values."""
        self.__dataset: Optional[Dataset] = None
        self.__executor: Optional[SQLExecutor2] = None

    @property
    def is_initialized(self) -> bool:
        """Returns the status of the LoggingDataset."""
        return bool(self.__dataset and self.__executor)
    
    @property
    def quoted_table_name(self):
        """
        Returns the quoted resolved table name of the underlying dataset.

        Raises:
            ValueError: If the service has not been initialized correctly and the dataset is None.

        Returns:
            str or None: The quoted resolved table name if available, otherwise None.
        """
        if self.__dataset is None:
            raise ValueError("Service has not been initialized correctly.")

        return self.__dataset.get_location_info().get('info', {}).get('quotedResolvedTableName')

    def log_query(self,
        user: str,
        title: str,
        llm_name: str,
        question: str,
        answer: str,
        graph_id: GraphId
    ) -> UUID:
        """
        Logs a new query to the database.

        Creates a new log record with the provided details and commits it to the
        database session.

        Args:
            user (str): The identifier of the user who made the query.
            title (str): The title of the conversation generated by the LLM.
            llm_name (str): The ID of the LLM used.
            question (str): The question asked by the user.
            answer (str): The answer generated by the LLM.
            graph_uuid (str): The UUID for the graph used.

        Returns:
            UUID: The unique identifier of the newly created log record.

        Raises:
            ValueError: If the service has not been initialized before calling this method.
        """
        if self.__executor is None or self.__dataset is None:
            raise ValueError("Service has not been initialized correctly.")

        log = LoggingDatasetModel(
            user=user,
            title=title,
            llm_name=llm_name,
            question=question,
            answer=answer,
            graph_id=graph_id
        )

        dialect = self.__dataset.get_config()["type"]
        insert = log.to_sql_insert(self.quoted_table_name, dialect)
        self.__executor.query_to_df(insert, post_queries=['COMMIT'])

        return log.uuid


    def init(self, dataset_name: str) -> None:
        """
        Initializes the service, gets the table name, and defines the History model.

        This is the main entry point for setting up the logging feature. It connects
        to the specified Dataiku dataset and dynamically creates a SQLAlchemy model
        class with the correct table name.

        Args:
            dataset_name (str): The name of the Dataiku dataset to use for logging.
        """
        if self.is_initialized:
            return

        dataset = Dataset(dataset_name)
        # Prevent init if the dataset is not a SQL dataset
        if dataset.get_config()["type"] not in DIALECT_VALUES:
            logger.warning("Logging dataset is not a SQL dataset. LLM queries will not be stored.")
            return

        self.__dataset = dataset
        self.__executor = SQLExecutor2(dataset=dataset_name)

        self.init_schema()

    def init_schema(self) -> None:
        """
        Writes the schema to the dataset if it doesn't already exist.

        Checks if the dataset has a schema. If not, it creates an empty DataFrame
        with the correct columns (derived from the History model) and writes it
        to the dataset to establish the schema.

        Raises:
            ValueError: If the service has not been initialized correctly.
        """
        if not self.__dataset:
            raise ValueError("Service has not been initialized correctly.")

        try:
            self.__dataset.read_schema(raise_if_empty=True)
        except Exception:
            columns = LoggingDatasetModel.get_columns()

            # Init schema with types
            data: Dict[str, List[str]] = {col: [] for col in columns.keys()}
            df = pd.DataFrame(data, columns=list(columns.keys()), dtype=str).astype(columns)
            self.__dataset.write_with_schema(df)

logging_dataset = LoggingDataset()
