"""Stores supplemental attribute associations in SQLite database"""

import itertools
import sqlite3
from typing import Any, Optional, Sequence
from uuid import UUID

from loguru import logger

from infrasys import Component
from infrasys.supplemental_attribute import SupplementalAttribute
from infrasys.utils.sqlite import execute
from infrasys.exceptions import ISAlreadyAttached

TABLE_NAME = "supplemental_attribute_associations"


class SupplementalAttributeAssociationsStore:
    """Stores supplemental attribute associations in a SQLite database."""

    TABLE_NAME = TABLE_NAME

    def __init__(self, con: sqlite3.Connection, initialize: bool = True):
        self._con = con
        if initialize:
            self._create_association_table()
        self._create_indexes()

    def _create_association_table(self):
        schema = [
            "id INTEGER PRIMARY KEY",
            "attribute_uuid TEXT",
            "attribute_type TEXT",
            "component_uuid TEXT",
            "component_type TEXT",
        ]
        schema_text = ",".join(schema)
        cur = self._con.cursor()
        execute(cur, f"CREATE TABLE {self.TABLE_NAME}({schema_text})")
        self._con.commit()
        logger.debug("Created in-memory time series metadata table")

    def _create_indexes(self) -> None:
        cur = self._con.cursor()
        execute(
            cur,
            f"CREATE INDEX IF NOT EXISTS by_attribute ON {self.TABLE_NAME} "
            f"(attribute_uuid, component_uuid, component_type)",
        )
        execute(
            cur,
            f"CREATE INDEX IF NOT EXISTS by_component ON {self.TABLE_NAME} "
            f"(component_uuid, attribute_uuid, attribute_type)",
        )

    _ADD_ASSOCIATION_QUERY = f"""
        SELECT id FROM {TABLE_NAME}
        WHERE attribute_uuid = ? AND component_uuid = ?
        LIMIT 1
    """

    def add(self, component: Component, attribute: SupplementalAttribute) -> None:
        """Add association to the database.

        Raises
        ------
        ISAlreadyAttached
            Raised if the time series metadata already stored.
        """
        params = (str(attribute.uuid), str(component.uuid))
        cur = self._con.cursor()
        res = execute(cur, self._ADD_ASSOCIATION_QUERY, params=params).fetchone()
        if res:
            msg = f"An association with {component=} {attribute=} is already stored."
            raise ISAlreadyAttached(msg)

        row = (
            None,
            str(attribute.uuid),
            type(attribute).__name__,
            str(component.uuid),
            type(component).__name__,
        )

        placeholder = ",".join(itertools.repeat("?", len(row)))
        query = f"INSERT INTO {self.TABLE_NAME} VALUES ({placeholder})"
        execute(cur, query, params=row)
        self._con.commit()

    _HAS_ASSOCIATION_BY_COMPONENT_AND_ATTRIBUTE_QUERY = f"""
        SELECT id FROM {TABLE_NAME}
        WHERE attribute_uuid = ? AND component_uuid = ?
        LIMIT 1
    """

    def has_association_by_component_and_attribute(
        self,
        component: Component,
        attribute: SupplementalAttribute,
    ) -> bool:
        """Return True if the component and supplemental attribute have an association."""
        params = (str(attribute.uuid), str(component.uuid))
        return self._has_rows(self._HAS_ASSOCIATION_BY_COMPONENT_AND_ATTRIBUTE_QUERY, params)

    _HAS_ASSOCIATION_BY_ATTRIBUTE_QUERY = f"SELECT id FROM {TABLE_NAME} WHERE attribute_uuid = ?"

    def has_association_by_attribute(self, attribute: SupplementalAttribute) -> bool:
        """Return true if there is at least one association matching the inputs."""
        # Note: Unlike the other has_association methods, this is not covered by an index.
        params = (str(attribute.uuid),)
        return self._has_rows(self._HAS_ASSOCIATION_BY_ATTRIBUTE_QUERY, params)

    _HAS_ASSOCIATION_BY_COMPONENT_QUERY = f"SELECT id FROM {TABLE_NAME} WHERE component_uuid = ?"

    def has_association_by_component(self, component: Component) -> bool:
        """Return True if there is at least one association with the component."""
        params = (str(component.uuid),)
        return self._has_rows(self._HAS_ASSOCIATION_BY_COMPONENT_QUERY, params)

    _HAS_ASSOCIATION_BY_COMPONENT_AND_ATTRIBUTE_TYPE_QUERY = f"""
        SELECT attribute_uuid
        FROM {TABLE_NAME}
        WHERE component_uuid = ? AND attribute_type = ?
        LIMIT 1
    """

    def has_association_by_component_and_attribute_type(
        self, component: Component, attribute_type: str
    ) -> bool:
        """Return True if the component has an association with a supplemental attribute of the
        given type.
        """
        params = (str(component.uuid), attribute_type)
        return self._has_rows(self._HAS_ASSOCIATION_BY_COMPONENT_AND_ATTRIBUTE_TYPE_QUERY, params)

    def _has_rows(self, query: str, params: Sequence[Any]) -> bool:
        cur = self._con.cursor()
        res = execute(cur, query, params=params).fetchone()
        return res is not None

    _LIST_ASSOCIATED_COMPONENT_UUIDS_QUERY = f"""
        SELECT component_uuid
        FROM {TABLE_NAME}
        WHERE attribute_uuid = ?
    """

    def list_associated_component_uuids(self, attribute: SupplementalAttribute) -> list[UUID]:
        """Return the component UUIDs associated with the attribute."""
        params = (str(attribute.uuid),)
        cur = self._con.cursor()
        rows = execute(cur, self._LIST_ASSOCIATED_COMPONENT_UUIDS_QUERY, params=params)
        return [UUID(x[0]) for x in rows]

    _LIST_ASSOCIATED_SUPPLEMENTAL_ATTRIBUTE_UUIDS_QUERY1 = f"""
        SELECT attribute_uuid
        FROM {TABLE_NAME}
        WHERE component_uuid = ?
    """
    _LIST_ASSOCIATED_SUPPLEMENTAL_ATTRIBUTE_UUIDS_QUERY2 = f"""
        SELECT attribute_uuid
        FROM {TABLE_NAME}
        WHERE attribute_type = ? AND component_uuid = ?
    """

    def list_associated_supplemental_attribute_uuids(
        self,
        component: Component,
        attribute_type: Optional[str] = None,
    ) -> list[UUID]:
        """Return the supplemental attribute UUIDs associated with the component and attribute
        type.
        """
        if attribute_type is None:
            query = self._LIST_ASSOCIATED_SUPPLEMENTAL_ATTRIBUTE_UUIDS_QUERY1
            params = (str(component.uuid),)
        else:
            query = self._LIST_ASSOCIATED_SUPPLEMENTAL_ATTRIBUTE_UUIDS_QUERY2
            params = (attribute_type, str(component.uuid))  # type: ignore
        cur = self._con.cursor()
        rows = execute(cur, query, params=params)
        return [UUID(x[0]) for x in rows]

    def remove_association_by_attribute(
        self,
        attribute: SupplementalAttribute,
        must_exist=True,
    ) -> None:
        """Remove all associations with the given attribute."""
        where_clause = "WHERE attribute_uuid = ?"
        params = (str(attribute.uuid),)
        num_deleted = self._remove_associations(where_clause, params)
        if must_exist and num_deleted < 1:
            msg = f"Bug: unexpected number of deletions: {num_deleted}. Should have been >= 1."
            raise Exception(msg)

    def remove_association(
        self,
        component: Component,
        attribute: SupplementalAttribute,
    ) -> None:
        """Remove the association between the attribute and component."""
        where_clause = "WHERE attribute_uuid = ? AND component_uuid = ?"
        params = (str(attribute.uuid), str(component.uuid))
        num_deleted = self._remove_associations(where_clause, params)
        if num_deleted != 1:
            msg = f"Bug: unexpected number of deletions: {num_deleted}. Should have been 1."
            raise Exception(msg)

    # This functionality, copied from Sienna, could be added if needed.
    # def remove_associations(self, attribute_type: str) -> None:
    #    """Remove all associations of the given type."""
    #    where_clause = "WHERE attribute_type = ?"
    #    params = (attribute_type,)
    #    num_deleted = self._remove_associations(where_clause, params)
    #    logger.debug("Deleted %s supplemental attribute associations", num_deleted)

    def _remove_associations(self, where_clause: str, params: Sequence[Any]) -> int:
        query = f"DELETE FROM {self.TABLE_NAME} {where_clause}"
        cur = self._con.cursor()
        execute(cur, query, params)
        rows = execute(cur, "SELECT CHANGES() AS changes").fetchall()
        assert len(rows) == 1, rows
        row = rows[0]
        logger.debug("Deleted %s rows from the time series metadata table", row[0])
        self._con.commit()
        return row[0]

    _GET_ATTRIBUTE_COUNTS_BY_TYPE_QUERY = f"""
        SELECT
            attribute_type
            ,count(*) AS count
        FROM {TABLE_NAME}
        GROUP BY
            attribute_type
        ORDER BY
            attribute_type
    """

    def get_attribute_counts_by_type(self) -> list[dict[str, Any]]:
        """Return a list of dicts of stored attribute counts by type."""
        cur = self._con.cursor()
        rows = execute(cur, self._GET_ATTRIBUTE_COUNTS_BY_TYPE_QUERY).fetchall()
        return [{"type": x[0], "count": x[1]} for x in rows]

    # TODO: This could be useful if we want to display a table to users. We don't yet
    # directly depend on Pandas. We could add that dependency or use some other table display.
    # This was copied from InfrastructureSystems.jl.
    # def get_attribute_summary_table(self) -> pd.DataFrame:
    #    """Return a DataFrame with the number of supplemental attributes by type for components."""
    #    query = f"""
    #        SELECT
    #            attribute_type
    #            ,component_type
    #            ,count(*) AS count
    #        FROM {self.TABLE_NAME}
    #        GROUP BY
    #            attribute_type
    #            ,component_type
    #        ORDER BY
    #            attribute_type
    #            ,component_type
    #    """
    #    cur = self._con.cursor()
    #    rows = execute(cur, query).fetchall()
    #    #return DataFrame(_execute(associations, query))

    _GET_NUM_ATTRIBUTES_QUERY = f"""
            SELECT COUNT(DISTINCT attribute_uuid) AS count
            FROM {TABLE_NAME}
        """

    def get_num_attributes(self) -> int:
        """Return the number of supplemental attributes."""
        cur = self._con.cursor()
        return execute(cur, self._GET_NUM_ATTRIBUTES_QUERY).fetchone()[0]

    _GET_NUM_COMPONENTS_WITH_ATTRIBUTES_QUERY = f"""
        SELECT COUNT(DISTINCT component_uuid) AS count
        FROM {TABLE_NAME}
    """

    def get_num_components_with_attributes(self) -> int:
        """Return the number of components with supplemental attributes."""
        cur = self._con.cursor()
        return execute(cur, self._GET_NUM_COMPONENTS_WITH_ATTRIBUTES_QUERY).fetchone()[0]
