import requests
import pandas as pd
import ast


class PowerBIClient:
    # Power BI Client Class

    # Init and get token
    def __init__(self, tenant_id, client_id, client_secret):
        self.tenant_id = tenant_id
        self.client_id = client_id
        self.client_secret = client_secret
        self.access_token = self.get_token()
        self.base_url = "https://api.powerbi.com/v1.0/myorg/groups/"

    def get_token(self):
        # Grab token
        header = {
            "Content-Type": "application/x-www-form-urlencoded",
            "Host": "login.microsoftonline.com:443",
        }

        data = {
            "grant_type": "client_credentials",
            "scope": "https://analysis.windows.net/powerbi/api/.default",
            "client_id": self.client_id,
            "client_secret": self.client_secret,
        }

        result = requests.post(
            f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token",
            headers=header,
            data=data,
        )
        self.access_token = result.json()["access_token"]
        return self.access_token

    def get_header(self):
        # Define header
        header = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.access_token}",
        }
        return header

    # Other functions
    def extract_connection_details(self, x):
        try:
            if isinstance(x, str):
                details = ast.literal_eval(x)
            elif isinstance(x, dict):
                details = x
            else:
                return pd.Series(
                    [None] * 5,
                    index=["server", "database", "connectionString", "url", "path"],
                )

            return pd.Series(
                {
                    "server": details.get("server"),
                    "database": details.get("database"),
                    "connectionString": details.get("connectionString"),
                    "url": details.get("url"),
                    "path": details.get("path"),
                }
            )
        except Exception as e:
            print(f"Error parsing connectionDetails: {e}")
            return pd.Series(
                [None] * 5,
                index=["server", "database", "connectionString", "url", "path"],
            )

    # Trigger actions
    def refresh_dataflow(self, workspace_id, dataflow_id):
        # Define URL

        url = f"{self.base_url}/{workspace_id}/dataflows/{dataflow_id}/refreshes"
        # Send POST request
        result = requests.post(url, headers=self.get_header())
        if result.status_code == 200:
            print(f"Start refreshing dataflow {dataflow_id}")
        else:
            # If failed to refresh dataflow, print error message
            print(
                f"Failed to refresh dataflow {dataflow_id}. Status code: {result.status_code}"
            )

    def refresh_dataset(self, workspace_id, dataset_id):
        # Define URL
        url = f"{self.base_url}/{workspace_id}/datasets/{dataset_id}/refreshes"
        # Send POST request
        result = requests.post(url, headers=self.get_header())
        if result.status_code == 202:
            print(f"Start refreshing dataset {dataset_id}")
        else:
            # If failed to refresh dataset, print error message
            print(
                f"Failed to refresh dataset {dataset_id}. Status code: {result.status_code}"
            )

    # Get data functions
    def get_report_by_workspace(self, workspace_id):
        # Define URL endpoint
        get_report_url = f"{self.base_url}/{workspace_id}/reports"
        # Send API call to get data
        result = requests.get(url=get_report_url, headers=self.get_header())
        return result

    def get_dataset_by_workspace(self, workspace_id):
        # Define URL endpoint
        get_dataset_url = f"{self.base_url}/{workspace_id}/datasets"
        # Send API call to get data
        result = requests.get(url=get_dataset_url, headers=self.get_header())
        return result

    def get_dataflow_by_workspace(self, workspace_id):
        # Define URL endpoint
        get_dataflow_url = f"{self.base_url}/{workspace_id}/dataflows"
        # Send API call to get data
        result = requests.get(url=get_dataflow_url, headers=self.get_header())
        return result

    def get_dataset_refresh_history_by_id(self, workspace_id, dataset_id, top_n=10):
        """
        Get dataset refresh history by dataset id.
        Args:
            workspace_id (str): The ID of the Power BI workspace.
            dataset_id (str): The ID of the Power BI dataset.
            top_n (int): The number of most recent refreshes to retrieve. Default is 10.
        Returns:
            result (requests.Response): The response object containing the dataset refresh history.
        """
        # Define URL endpoint
        get_dataset_refresh_history_url = f"{self.base_url}/{workspace_id}/datasets/{dataset_id}/refreshes?$top={top_n}"
        # Send API to get data
        result = requests.get(
            url=get_dataset_refresh_history_url, headers=self.get_header()
        )
        return result

    def get_dataflow_refresh_history_by_id(self, workspace_id, dataflow_id):
        """
        Get dataflow refresh history by dataflow id.
        Args:
            workspace_id (str): The ID of the Power BI workspace.
            dataflow_id (str): The ID of the Power BI dataflow.
        Returns:
            result (requests.Response): The response object containing the dataflow refresh history.
        """
        get_dataflow_refresh_history_url = (
            f"{self.base_url}/{workspace_id}/dataflows/{dataflow_id}/transactions"
        )
        # Send API to get data
        result = requests.get(
            url=get_dataflow_refresh_history_url, headers=self.get_header()
        )

        return result

    def get_dataset_sources_by_id(self, workspace_id, dataset_id):

        """
        Get dataset sources by dataset id.

        Args:
            workspace_id (str): The ID of the Power BI workspace.
            dataset_id (str): The ID of the Power BI dataset.
        Returns:
            result (requests.Response): The response object containing the dataset sources.
        """
        # Define URL endpoint
        get_dataset_source_url = (
            f"{self.base_url}/{workspace_id}/datasets/{dataset_id}/datasources"
        )
        # Send API to get data
        result = requests.get(url=get_dataset_source_url, headers=self.get_header())
        return result

    def get_dataflow_sources_by_id(self, workspace_id, dataflow_id):
        """
        Get dataflow sources by dataflow id.

        Args:
            workspace_id (str): The ID of the Power BI workspace.
            dataflow_id (str): The ID of the Power BI dataflow.
        Returns:
            result (requests.Response): The response object containing the dataflow sources.
        """
        # Define URL endpoint
        get_dataflow_source_url = (
            f"{self.base_url}/{workspace_id}/dataflows/{dataflow_id}/datasources"
        )
        # Send API to get data
        result = requests.get(url=get_dataflow_source_url, headers=self.get_header())
        return result

    def get_dataset_users_by_id(self, workspace_id, dataset_id):
        """
        Get dataset users by dataset id.

        Args:
            workspace_id (str): The ID of the Power BI workspace.
            dataset_id (str): The ID of the Power BI dataset.
        Returns:
            result (requests.Response): The response object containing the dataset users.
        """
        # Define URL endpoint
        get_dataset_users_url = (
            f"{self.base_url}/{workspace_id}/datasets/{dataset_id}/users"
        )
        # Send API to get data
        result = requests.get(url=get_dataset_users_url, headers=self.get_header())
        return result

    # Get data in bulk
    def get_all_workspaces(self):
        """
        Get all workspaces from Power BI.
        """
        # Send API Request
        result = requests.get(url=self.base_url, headers=self.get_header())
        # Convert to dataframe
        df_get_all_workspaces = pd.DataFrame.from_dict(
            result.json()["value"], orient="columns"
        )
        return df_get_all_workspaces

    def get_all_dataflows(self):
        """
        Get all dataflows from Power BI.
        """
        # Set workspace list
        df_get_all_workspaces = self.get_all_workspaces()
        workspace_id_list = df_get_all_workspaces["id"]
        # Define an empty dataframe
        df_get_all_dataflows = pd.DataFrame()
        # Loop through workspace
        for workspace_id in workspace_id_list:
            try:
                workspace_name = df_get_all_workspaces.query(
                    'id == "{0}"'.format(workspace_id)
                )["name"].iloc[0]
                # Send API call to get data
                result = self.get_dataflow_by_workspace(workspace_id)
                # If result success then proceed:
                if result.status_code == 200:
                    # Create dataframe to store data
                    df = pd.DataFrame.from_dict(
                        result.json()["value"], orient="columns"
                    )
                    # Add column
                    df["workspaceId"] = workspace_id
                    # Add workspace name column
                    df["workspaceName"] = workspace_name
                    # Convert all columns to string type (optional)
                    df = df.astype("str")
                    # Append data
                    df_get_all_dataflows = pd.concat([df_get_all_dataflows, df])
            except Exception as e:
                print(f"Error processing workspace {workspace_id}: {e}")
                continue
        return df_get_all_dataflows

    def get_all_reports(self):
        """
        Get all reports from Power BI.
        """
        # Set workspace list
        df_get_all_workspaces = self.get_all_workspaces()
        workspace_id_list = df_get_all_workspaces["id"]
        # Define an empty dataframe
        df_get_all_reports = pd.DataFrame()
        # Loop through workspace
        for workspace_id in workspace_id_list:
            try:
                workspace_name = df_get_all_workspaces.query(
                    'id == "{0}"'.format(workspace_id)
                )["name"].iloc[0]
                # Send API call to get data
                result = self.get_report_by_workspace(workspace_id)
                # If result success then proceed:
                if result.status_code == 200:
                    # Create dataframe to store data
                    df = pd.DataFrame.from_dict(
                        result.json()["value"], orient="columns"
                    )
                    # Add workspace name column
                    df["workspaceName"] = workspace_name
                    # Convert all columns to string type (optional)
                    df = df.astype("str")
                    # Append data
                    df_get_all_reports = pd.concat([df_get_all_reports, df])
            except Exception as e:
                print(f"Error processing workspace {workspace_id}: {e}")
                continue
        return df_get_all_reports

    def get_all_dataset_refresh_history(self):
        """
        Get all dataset refresh history from Power BI.
        """
        # Get all datasets
        df_get_all_datasets = self.get_all_datasets()
        # Get dataset refresh history
        df_get_all_datasets_refresh_history = pd.DataFrame()
        list_of_ds = df_get_all_datasets.query('isRefreshable == "True"')["id"]
        # Loop through dataset
        for dataset_id in list_of_ds:
            try:
                # Get workspace id
                workspace_id = df_get_all_datasets.query(
                    'id == "{0}"'.format(dataset_id)
                )["workspaceId"].iloc[0]
                # Get workspace name
                workspace_name = df_get_all_datasets.query(
                    'id == "{0}"'.format(dataset_id)
                )["workspaceName"].iloc[0]
                # Get dataset name
                dataset_name = df_get_all_datasets.query(
                    'id == "{0}"'.format(dataset_id)
                )["name"].iloc[0]
                # Send API to get data
                result = self.get_dataset_refresh_history_by_id(
                    workspace_id, dataset_id
                )
                # If result success then proceed:
                if result.status_code == 200:
                    # Parse data from json output
                    df = pd.DataFrame.from_dict(
                        result.json()["value"], orient="columns"
                    )
                    # Add column
                    df["datasetId"] = dataset_id
                    # Add column
                    df["datasetName"] = dataset_name
                    # Add column
                    df["workspaceId"] = workspace_id
                    # Add column
                    df["workspaceName"] = workspace_name
                    # Convert all columns to string type (optional)
                    df = df.astype("str")
                    # Append data
                    df_get_all_datasets_refresh_history = pd.concat(
                        [df_get_all_datasets_refresh_history, df]
                    )
            except Exception as e:
                print(f"Error processing dataset log {dataset_id}: {e}")
                continue

        return df_get_all_datasets_refresh_history

    def get_all_dataflow_refresh_history(self):
        """
        Get all dataflow refresh history from Power BI.
        """
        # Get all dataflows
        df_get_all_dataflows = self.get_all_dataflows()
        # Get dataflow refresh history
        df_get_all_dataflows_refresh_history = pd.DataFrame()
        list_of_dataflows = df_get_all_dataflows["objectId"]

        # Loop through dataflow
        for dataflow_id in list_of_dataflows:
            try:
                # Get workspace id
                workspace_id = df_get_all_dataflows.query(
                    'objectId == "{0}"'.format(dataflow_id)
                )["workspaceId"].iloc[0]
                # Get workspace name
                workspace_name = df_get_all_dataflows.query(
                    'objectId == "{0}"'.format(dataflow_id)
                )["workspaceName"].iloc[0]
                # Get dataflow name
                dataflow_name = df_get_all_dataflows.query(
                    'objectId == "{0}"'.format(dataflow_id)
                )["name"].iloc[0]
                # Send API to get data
                result = self.get_dataflow_refresh_history_by_id(
                    workspace_id, dataflow_id
                )
                # If api_call success then proceed:
                if result.status_code == 200:
                    # Parse data from json output
                    df = pd.DataFrame.from_dict(
                        result.json()["value"], orient="columns"
                    )
                    # Add column
                    df["dataflowId"] = dataflow_id
                    # Add column
                    df["dataflowName"] = dataflow_name
                    # Add column
                    df["workspaceId"] = workspace_id
                    # Add column
                    df["workspaceName"] = workspace_name
                    # Convert all columns to string type (optional)
                    df = df.astype("str")
                    # Append data
                    df_get_all_dataflows_refresh_history = pd.concat(
                        [df_get_all_dataflows_refresh_history, df]
                    )
            except Exception as e:
                print(f"Error processing dataflow {dataflow_id}: {e}")
                continue

        return df_get_all_dataflows_refresh_history

    def get_all_dataset_users(self):
        """
        Get all dataset users from Power BI.
        """
        # Get all datasets
        df_get_all_datasets = self.get_all_datasets()
        # Filter Usage Report dataset
        df_get_all_datasets = df_get_all_datasets[
            ~df_get_all_datasets["name"].str.contains("Usage Metrics")
        ]
        # Get report list
        dataset_id_list = df_get_all_datasets["id"]
        # Define an empty dataframe
        df_get_all_dataset_users = pd.DataFrame()
        # Loop through dataset
        for dataset_id in dataset_id_list:
            try:
                workspace_id = df_get_all_datasets.query(f'id == "{dataset_id}"')[
                    "workspaceId"
                ].iloc[0]
                workspace_name = df_get_all_datasets.query(f'id == "{dataset_id}"')[
                    "workspaceName"
                ].iloc[0]
                dataset_name = df_get_all_datasets.query(f'id == "{dataset_id}"')[
                    "name"
                ].iloc[0]
                # Send API call to get data
                result = self.get_dataset_users_by_id(workspace_id, dataset_id)
                # If result success then proceed:
                if result.status_code == 200:
                    # Create dataframe to store data
                    df = pd.DataFrame.from_dict(
                        result.json()["value"], orient="columns"
                    )
                    # Add workspace name column
                    df["workspaceId"] = workspace_id
                    df["workspaceName"] = workspace_name
                    df["datasetId"] = dataset_id
                    df["datasetName"] = dataset_name
                    # Convert all columns to string type (optional)
                    df = df.astype("str")
                    # Append data
                    df_get_all_dataset_users = pd.concat([df_get_all_dataset_users, df])
            except Exception as e:
                print(f"Error processing dataset {dataset_id}: {e}")
                continue

        return df_get_all_dataset_users

    def get_all_dataset_sources(self):
        """
        Get all dataset data sources from Power BI.
        """
        # Filter Usage Report dataset
        df_get_all_datasets = self.get_all_datasets()
        df_get_all_datasets = df_get_all_datasets[
            ~df_get_all_datasets["name"].str.contains("Usage Metrics")
        ]
        # Get report list
        dataset_id_list = df_get_all_datasets["id"]
        # Define an empty dataframe
        df_get_all_dataset_sources = pd.DataFrame()

        # Loop through dataset
        for dataset_id in dataset_id_list:
            try:
                workspace_id = df_get_all_datasets.query(f'id == "{dataset_id}"')[
                    "workspaceId"
                ].iloc[0]
                workspace_name = df_get_all_datasets.query(f'id == "{dataset_id}"')[
                    "workspaceName"
                ].iloc[0]
                dataset_name = df_get_all_datasets.query(f'id == "{dataset_id}"')[
                    "name"
                ].iloc[0]
                # Send API call to get data
                result = self.get_dataset_sources_by_id(workspace_id, dataset_id)
                # If result success then proceed:
                if result.status_code == 200:
                    # Create dataframe to store data
                    df = pd.DataFrame(result.json()["value"])
                    # Add workspace name column
                    df["workspaceId"] = workspace_id
                    df["workspaceName"] = workspace_name
                    df["datasetId"] = dataset_id
                    df["datasetName"] = dataset_name
                    # Extract more useful columns
                    df[["server", "database", "connectionString", "url", "path"]] = df[
                        "connectionDetails"
                    ].apply(self.extract_connection_details)
                    # Convert all columns to string type (optional)
                    df = df.astype("str")
                    # Append data
                    df_get_all_dataset_sources = pd.concat(
                        [df_get_all_dataset_sources, df]
                    )
            except Exception as e:
                print(f"Error processing dataset {dataset_id}: {e}")
                continue

        return df_get_all_dataset_sources

    def get_all_dataflow_sources(self):
        """
        Get all dataflow data sources from Power BI.
        """
        # Get all dataflows
        df_get_all_dataflows = self.get_all_dataflows()
        # Get report list
        dataflow_id_list = df_get_all_dataflows["objectId"]
        # Define an empty dataframe
        df_get_all_dataflow_sources = pd.DataFrame()
        # Loop through dataset
        for dataflow_id in dataflow_id_list:
            workspace_id = df_get_all_dataflows.query(f'objectId == "{dataflow_id}"')[
                "workspaceId"
            ].iloc[0]
            workspace_name = df_get_all_dataflows.query(f'objectId == "{dataflow_id}"')[
                "workspaceName"
            ].iloc[0]
            dataflow_name = df_get_all_dataflows.query(f'objectId == "{dataflow_id}"')[
                "name"
            ].iloc[0]

            result = self.get_dataflow_sources_by_id(workspace_id, dataflow_id)
            # If result success then proceed:
            if result.status_code == 200:
                try:
                    # Create dataframe to store data
                    df = pd.DataFrame(result.json()["value"])
                    # Add workspace name column
                    df["workspaceId"] = workspace_id
                    df["workspaceName"] = workspace_name
                    df["dataflowId"] = dataflow_id
                    df["dataflowName"] = dataflow_name
                    # Extract more useful columns
                    df[["server", "database", "connectionString", "url", "path"]] = df[
                        "connectionDetails"
                    ].apply(self.extract_connection_details)
                    # Convert all columns to string type (optional)
                    df = df.astype("str")
                    # Append data
                    df_get_all_dataflow_sources = pd.concat(
                        [df_get_all_dataflow_sources, df]
                    )
                except Exception as e:
                    print(f"Error processing dataflow {dataflow_id}: {e}")
                    continue

        return df_get_all_dataflow_sources
