from __future__ import annotations

from typing import TYPE_CHECKING, Any, Dict, Union

from sarus_query_builder.core.core import OptimizableQueryBuilder, QueryBuilder
from sarus_query_builder.core.typing import Task
from sarus_query_builder.protobuf.query_pb2 import Query

if TYPE_CHECKING:
    from sarus_data_spec.typing import Dataset

from sarus_differential_privacy.query import (  # type: ignore
    ComposedQuery,
    EpsilonDeltaQuery,
    EpsilonQuery,
    LaplaceQuery,
    PrivateQuery,
)
from sarus_sql.ast_utils import (  # type: ignore
    load_symbols_in_query,
    parse_query,
)
from sarus_sql.convert_dataspec_to_metadata import load_metadata
from sarus_sql.estimate_privacy_spending import (  # type: ignore
    estimate_privacy_spending,
)
from sarus_sql.protobuf.sql_pb2 import SQL, Budget  # type: ignore
from sarus_sql.validation import validate  # type: ignore


class SQLBuilder(QueryBuilder):
    """Generate SQL hyperparameters"""

    def __init__(self, dataset: Dataset):
        self._dataset = dataset
        self._metadata = load_metadata(dataset)
        self._is_validated = False
        self._valid = False

    def build_query(self, input_parameter: Query.SQL) -> Task:
        sql_string_query = input_parameter.sql_query
        target_epsilon = input_parameter.epsilon
        target_delta = input_parameter.delta

        if target_epsilon == 0.0:
            epsilon_multiplier = 0.0
            target_epsilon = 0.0
        else:
            epsilon_multiplier = target_epsilon

        try:
            if not self._is_validated:
                self._is_validated = True
                parsed_query = parse_query(
                    sql_string_query,
                )
                load_symbols_in_query(parsed_query, self._metadata)
                if input_parameter.sql_unlimited:
                    parsed_query.row_privacy = True
                validate(
                    parsed_query,
                    metadata=self._metadata,
                    sql_unlimited=input_parameter.sql_unlimited,
                )
                self._is_valid = True

            if self._is_valid:
                budget = budget_repartition(epsilon_multiplier, target_delta)
            else:
                raise NotImplementedError

        except NotImplementedError:
            self._is_valid = False
            target_epsilon = 0.0
            budget = Budget(
                false_positive_proba=0.99999,
                max_ids_approx_bounds=1,
                epsilon_approx_bounds_histogram=0.0,
                epsilon_approx_bounds_threshold=0.0,
                epsilon_agg=0.0,
                epsilon_tau_thresholding=0.0,
                delta_tau_thresholding=0.0,
            )

        return SQL(
            query=sql_string_query,
            budget=budget,
            sql_unlimited=input_parameter.sql_unlimited,
            epsilon=target_epsilon,
        )

    def private_query(self, out: Task) -> PrivateQuery:
        try:
            spendings = estimate_privacy_spending(
                self.dataset, out, self._metadata
            )(out.budget)
        except NotImplementedError:
            spendings = []
        return ComposedQuery([convert_mech(mech) for mech in spendings])


def convert_mech(
    mech: Dict[str, Any]
) -> Union[LaplaceQuery, EpsilonQuery, EpsilonDeltaQuery]:
    """Convert Sarus SQL mechanism to the Accountant queries"""
    if mech["mechanism"] == "laplace":
        return LaplaceQuery(1.0 / mech["epsilon"])
    if mech["delta"] == 0:
        EpsilonQuery(mech["epsilon"])
    return EpsilonDeltaQuery(mech["epsilon"], mech["delta"])


class OptimizableSQLBuilder(OptimizableQueryBuilder):
    def __init__(self, dataset: Dataset, query: Query):
        self._dataset = dataset
        self.query = query
        self._builders = [SQLBuilder(dataset)]

    def build_query(self, input_parameter: float) -> Task:
        query = self.query
        query.sql.epsilon = input_parameter
        return self.builders[0].build_query(query.sql)


def sql_builder(dataset: Dataset, query: Query) -> OptimizableSQLBuilder:
    return OptimizableSQLBuilder(dataset, query)


def budget_repartition(epsilon_multiplier: float, delta: float) -> Budget:
    return Budget(
        false_positive_proba=0.99999,
        max_ids_approx_bounds=1,
        epsilon_approx_bounds_histogram=epsilon_multiplier / 8.0,
        epsilon_approx_bounds_threshold=epsilon_multiplier / 8.0,
        epsilon_agg=epsilon_multiplier / 2,
        epsilon_tau_thresholding=epsilon_multiplier / 4,
        delta_tau_thresholding=delta / 1e3,
    )