from sys import path
path.append('C:\\Users\\Chenw\\Documents\\Trae\\PBI-Measure-CallGraph\\MSNET')

import pyadomd
import pandas as pd


class ModelConnector:
    """Class to connect to and retrieve data from Analysis Services models."""
    
    def __init__(self, server, database):
        """
        Initialize the connection to the Analysis Services model.
        
        Args:
            server (str): The server name or IP address
            database (str): The database name
        """
        self.server = server
        self.database = database
        if database == 'null' : 
            self.connection_string = f"Provider=MSOLAP;Data Source={server}"
        else :   
            self.connection_string = f"Provider=MSOLAP;Data Source={server};Initial Catalog={database}"
        self.connection = None
        self.connect()
        
    def connect(self):
        """Establish a connection to the Analysis Services model."""
        try:
            self.connection = pyadomd.Pyadomd(self.connection_string)
            self.connection.open()
            return True
        except Exception as e:
            raise Exception(f"Failed to connect to the model: {str(e)}")
    
    def execute_query(self, query):
        """
        Execute a DAX or MDX query against the model.
        
        Args:
            query (str): The DAX or MDX query to execute
            
        Returns:
            pandas.DataFrame: The query results
        """
        if not self.connection:
            self.connect()
            
        try:
            with self.connection.cursor().execute(query) as conn:
                rows = conn.fetchall()
                df = pd.DataFrame(rows, columns=[col[0] for col in conn._description])
            return df
        except Exception as e:
            raise Exception(f"Failed to execute query: {str(e)}")
    
    def get_measures(self):
        """
        Retrieve all measures from the model.
        
        Returns:
            list: A list of dictionaries containing measure information
        """
        # DMV query to get all measures
        query = """
        SELECT 
            [MEASURE_NAME] as [name],
            [EXPRESSION] as [expression],
            [MEASURE_CAPTION] as [caption],
            [MEASURE_DISPLAY_FOLDER] as [display_folder],
            [DESCRIPTION] as [description]
        FROM $SYSTEM.MDSCHEMA_MEASURES
        WHERE MEASURE_IS_VISIBLE
        """
        
        try:
            df = self.execute_query(query)
            return df if not df.empty else []
        except Exception as e:
            raise Exception(f"Failed to retrieve measures: {str(e)}")
        

if __name__ == '__main__':        
    # Example usage
    server = 'localhost:56195'
    database = 'Exteranl All Channel'
    
    model_connector = ModelConnector(server, database)
    measures = model_connector.get_measures()
    print(measures)  # List of measure dictionaries