"""Class to execute Snowflake queries used to validate adjustments."""

from string import Template
from typing import cast

from snowflake_connector.etl_connector import SnowflakeSQLExecutor

from abacus_common_logic.adjustments_validation.constants import (
    CONTRACT_LABEL_TERM,
    CONTRACT_PRODUCT_TERM,
)
from abacus_common_logic.adjustments_validation.snowflake.queries import (
    GET_ACCOUNT_CONTRACTS,
    GET_ACCOUNT_UPCS,
    GET_ACCOUNTS,
    GET_ADJUSTMENT_TYPES,
    GET_CLOSE_BALANCE_STATUSES,
    GET_CONTRACT_TERMS,
    GET_DISPLAY_UPCS,
    GET_PAYMENT_ENTITIES,
    GET_STATEMENT_PERIODS,
)
from abacus_common_logic.adjustments_validation.types_definitions import (
    SnowflakeConfig,
)
from abacus_common_logic.adjustments_validation.utils import load_json


class AdjustmentsValidationSnowflakeExecutor(SnowflakeSQLExecutor):
    """Execute Snowflake queries to validate adjustments."""

    def __init__(self, sf_config: SnowflakeConfig):
        """Initialize an instance."""
        super().__init__(sf_config)

        self._base_params = {
            'db': sf_config['db'],
            'schema': sf_config['schema'],
        }

    def fetch_close_balance_statuses(self, statement_period_id: int) -> dict[str, str]:
        """Fetch the payment entities `close_balance` status from Snowflake."""
        params = {**self._base_params, 'statement_period_id': statement_period_id}
        query = Template(GET_CLOSE_BALANCE_STATUSES).substitute(params)
        rows = cast(list[dict], self.fetchall(query, dict_cursor=True))

        results = {}
        for row in rows:
            payment_entity_id = str(row['PAYMENT_ENTITY_ID'])
            status = row['ACTION_STATUS']

            results[payment_entity_id] = status

        return results

    def fetch_accounts(self, account_ids: set[str]) -> set[str]:
        """Fetch a list of accounts by IDs from Snowflake."""
        if not account_ids:
            return set()

        params = {**self._base_params, 'account_ids': ', '.join(account_ids)}
        query = Template(GET_ACCOUNTS).substitute(params)
        rows = cast(list[dict], self.fetchall(query, dict_cursor=True))

        return set([str(row['ACCOUNT_ID']) for row in rows])

    def fetch_account_contracts(
        self, account_contract_map: dict[str, set[str]]
    ) -> dict[str, set[str]]:
        """Fetch a list of account contracts from Snowflake."""
        where_clause = []

        for account_id, contract_ids in account_contract_map.items():
            for contract_id in contract_ids:
                where_clause.append(
                    f'(account_id = {account_id} AND contract_id = {contract_id})'
                )

        if not where_clause:
            return {}

        params = {**self._base_params, 'where_clause': ' OR '.join(where_clause)}
        query = Template(GET_ACCOUNT_CONTRACTS).substitute(params)
        rows = cast(list[dict], self.fetchall(query, dict_cursor=True))

        results = {}
        for row in rows:
            account_id = str(row['ACCOUNT_ID'])
            contract_id = str(row['CONTRACT_ID'])

            if account_id not in results.keys():
                results[account_id] = set()

            results[account_id].add(contract_id)

        return results

    def fetch_payment_entities(self, account_ids: set[str]) -> dict[str, str]:
        """Fetch a list of payment entities by account IDs from Snowflake."""
        if not account_ids:
            return {}

        params = {**self._base_params, 'account_ids': ', '.join(account_ids)}
        query = Template(GET_PAYMENT_ENTITIES).substitute(params)
        rows = cast(list[dict], self.fetchall(query, dict_cursor=True))

        results = {}
        for row in rows:
            account_id = str(row['ACCOUNT_ID'])
            payment_entity_id = str(row['PAYMENT_ENTITY_ID'])

            if payment_entity_id:
                results[account_id] = payment_entity_id

        return results

    def fetch_display_upcs(self, display_upcs: set[str]) -> dict[str, set[str]]:
        """Fetch a list of display UPCs and their associated UPCs from Snowflake."""
        if not display_upcs:
            return {}

        formatted_upcs = [f"'{upc}'" for upc in display_upcs]
        params = {**self._base_params, 'display_upcs': ', '.join(formatted_upcs)}
        query = Template(GET_DISPLAY_UPCS).substitute(params)
        rows = cast(list[dict], self.fetchall(query, dict_cursor=True))

        results = {}
        for row in rows:
            display_upc = str(row['DISPLAY_UPC'])
            upcs = row['UPCS'].split(',')

            if display_upc not in results.keys():
                results[display_upc] = set()

            results[display_upc].update(upcs)

        return results

    def fetch_product_terms(self, contract_ids: set[str]) -> dict[str, set[str]]:
        """Fetch a list of contract product terms from Snowflake."""
        if not contract_ids:
            return {}

        params = {
            **self._base_params,
            'term_type': CONTRACT_PRODUCT_TERM,
            'contract_ids': ', '.join(contract_ids),
        }

        query = Template(GET_CONTRACT_TERMS).substitute(params)
        rows = cast(list[dict], self.fetchall(query, dict_cursor=True))

        results = {}
        for row in rows:
            contract_id = str(row['CONTRACT_ID'])
            attachments = row['ATTACHMENTS']
            upcs = load_json(attachments)

            if upcs:
                if contract_id not in results.keys():
                    results[contract_id] = set()

                results[contract_id].update(upcs)

        return results

    def fetch_label_terms(self, contract_ids: set[str]) -> dict[str, set[str]]:
        """Fetch a list of contract product terms from Snowflake."""
        if not contract_ids:
            return {}

        params = {
            **self._base_params,
            'term_type': CONTRACT_LABEL_TERM,
            'contract_ids': ', '.join(contract_ids),
        }

        query = Template(GET_CONTRACT_TERMS).substitute(params)
        rows = cast(list[dict], self.fetchall(query, dict_cursor=True))

        results = {}
        for row in rows:
            contract_id = str(row['CONTRACT_ID'])
            attachments = row['ATTACHMENTS']
            label_ids = load_json(attachments)

            if label_ids:
                if contract_id not in results.keys():
                    results[contract_id] = set()

                results[contract_id].update(label_ids)

        return results

    def fetch_account_upcs(
        self, account_upc_map: dict[str, set[str]]
    ) -> dict[str, set[str]]:
        """Fetch a list of UPCs by Account from Snowflake."""
        where_clause = []

        for account_id, upcs in account_upc_map.items():
            upc_list = ', '.join(f"'{upc}'" for upc in upcs)

            where_clause.append(
                f'(v.vendor_id = {account_id} AND (r.upc IN ({upc_list}) OR r.display_upc IN ({upc_list})))'
            )

        if not where_clause:
            return {}

        params = {**self._base_params, 'where_clause': ' OR '.join(where_clause)}
        query = Template(GET_ACCOUNT_UPCS).substitute(params)
        rows = cast(list[dict], self.fetchall(query, dict_cursor=True))

        results = {}
        for row in rows:
            account_id = str(row['VENDOR_ID'])
            upc = row['UPC']
            display_upc = row['DISPLAY_UPC']

            if account_id not in results.keys():
                results[account_id] = set()

            results[account_id].add(upc)
            results[account_id].add(display_upc)

        return results

    def fetch_statement_periods(self, statement_years: set[str]) -> dict[str, str]:
        """Fetch a list of statement periods by years from Snowflake."""
        if not statement_years:
            return {}

        params = {
            **self._base_params,
            'statement_years': ', '.join(statement_years),
        }

        query = Template(GET_STATEMENT_PERIODS).substitute(params)
        rows = cast(list[dict], self.fetchall(query, dict_cursor=True))

        results = {}
        for row in rows:
            key = f'{str(row["STATEMENT_MONTH"])}/{str(row["STATEMENT_YEAR"])}'

            results[key] = row['STATEMENT_PERIOD_STATUS']

        return results

    def fetch_adjustment_types(self) -> set[str]:
        """Fetch a list of adjustment types from Snowflake."""
        params = {**self._base_params}
        query = Template(GET_ADJUSTMENT_TYPES).substitute(params)
        rows = cast(list[dict], self.fetchall(query, dict_cursor=True))

        return set([row['TYPE_NAME'].lower() for row in rows])