from pydantic.v1 import BaseModel

from prowler.lib.logger import logger
from prowler.providers.gcp.config import DEFAULT_RETRY_ATTEMPTS
from prowler.providers.gcp.gcp_provider import GcpProvider
from prowler.providers.gcp.lib.service.service import GCPService


class BigQuery(GCPService):
    def __init__(self, provider: GcpProvider):
        super().__init__(__class__.__name__, provider, api_version="v2")

        self.datasets = []
        self.tables = []
        self._get_datasets()
        self._get_tables()

    def _get_datasets(self):
        for project_id in self.project_ids:
            try:
                request = self.client.datasets().list(projectId=project_id)
                while request is not None:
                    response = request.execute(num_retries=DEFAULT_RETRY_ATTEMPTS)

                    for dataset in response.get("datasets", []):
                        dataset_info = (
                            self.client.datasets()
                            .get(
                                projectId=project_id,
                                datasetId=dataset["datasetReference"]["datasetId"],
                            )
                            .execute(num_retries=DEFAULT_RETRY_ATTEMPTS)
                        )
                        cmk_encryption = False
                        public = False
                        roles = dataset_info.get("access", "")
                        if "allAuthenticatedUsers" in str(roles) or "allUsers" in str(
                            roles
                        ):
                            public = True
                        if dataset_info.get("defaultEncryptionConfiguration"):
                            cmk_encryption = True
                        self.datasets.append(
                            Dataset(
                                name=dataset["datasetReference"]["datasetId"],
                                id=dataset["id"],
                                region=dataset["location"],
                                cmk_encryption=cmk_encryption,
                                public=public,
                                project_id=project_id,
                            )
                        )

                    request = self.client.datasets().list_next(
                        previous_request=request, previous_response=response
                    )
            except Exception as error:
                logger.error(
                    f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
                )

    def _get_tables(self):
        for dataset in self.datasets:
            try:
                request = self.client.tables().list(
                    projectId=dataset.project_id, datasetId=dataset.name
                )
                while request is not None:
                    response = request.execute(num_retries=DEFAULT_RETRY_ATTEMPTS)

                    for table in response.get("tables", []):
                        cmk_encryption = False
                        if (
                            self.client.tables()
                            .get(
                                projectId=dataset.project_id,
                                datasetId=dataset.name,
                                tableId=table["tableReference"]["tableId"],
                            )
                            .execute(num_retries=DEFAULT_RETRY_ATTEMPTS)
                            .get("encryptionConfiguration")
                        ):
                            cmk_encryption = True
                        self.tables.append(
                            Table(
                                name=table["tableReference"]["tableId"],
                                id=table["id"],
                                region=dataset.region,
                                cmk_encryption=cmk_encryption,
                                project_id=dataset.project_id,
                            )
                        )

                    request = self.client.tables().list_next(
                        previous_request=request, previous_response=response
                    )
            except Exception as error:
                logger.error(
                    f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
                )


class Dataset(BaseModel):
    name: str
    id: str
    region: str
    cmk_encryption: bool
    public: bool
    project_id: str


class Table(BaseModel):
    name: str
    id: str
    region: str
    cmk_encryption: bool
    project_id: str
