import base64
import os
import json
import requests
import pandas as pd
from io import StringIO, BytesIO
import jsonpickle
import gzip
from enum import Enum
import logging
from typing import Optional
from typing import Optional, List, Dict, Any
from datetime import datetime

class ConditionOperator(Enum):
   """
    Enum representing an operator for filtering data in AIDataFilter.

    Attributes:
        Equals: Represents the equality condition.
        NotEquals: Represents the inequality condition.
        GreaterThan: Represents the greater than condition.
        LessThan: Represents the less than condition.
        GreaterThanOrEqualTo: Represents the greater than or equal to condition.
        LessThanOrEqualTo: Represents the less than or equal to condition.
        StartsWith: Represents the starts with condition.
        Like: Represents the like condition.
        NotStartsWith: Represents the not starts with condition.
        NotLike: Represents the not like condition.
        Contains: Represents the contains condition.
        NotContains: Represents the not contains condition.
    """
   Equals = 0
   NotEquals = 1
   GreaterThan = 2
   LessThan = 3
   GreaterThanOrEqualTo = 4
   LessThanOrEqualTo = 5
   StartsWith = 6
   Like = 7
   NotStartsWith = 8
   NotLike = 9
   Contains = 10
   NotContains = 11

class LogicalOperator(Enum):
    """
     Enum representing a logical operator for combining conditions in AIDataFilter.
    
     Attributes:
          And: Represents the logical AND operator.
          Or: Represents the logical OR operator.
          AndNot: Represents the logical AND NOT operator.
     """
    And = 0
    Or = 1
    AndNot = 2

class ChartType(Enum):
    """
    Enum representing different types of charts that can be generated by the inmydata platform.
    
    Attributes:
        Column: Represents a column chart.
        Pie: Represents a pie chart.
        Area: Represents an area chart.
        Bar: Represents a bar chart.
        Scatter: Represents a scatter plot.
        Bubble: Represents a bubble chart.
        Grid: Represents a grid chart.
    """
    Column = 0
    Pie = 1
    Area = 2
    Bar = 3
    Scatter = 4
    Bubble = 5
    Grid = 6

class AIDataFilter:
    """ 
    A class representing a filter for querying data in the inmydata platform.
    This class encapsulates the details of a filter, including the field to filter on, the condition operator,
    the logical operator, the value to filter by, and grouping information.
    
    Attributes:
        LogicalOperator (LogicalOperator): The logical operator to combine this filter with others.
        Field (str): The field to filter on.
        ConditionOperator (ConditionOperator): The condition operator to apply to the filter.   
        Value: The value to filter by.
        StartGroup (int): The starting group index for this filter.
        EndGroup (int): The ending group index for this filter.
        CaseInsensitive (bool): Whether the filter should be case insensitive.
    """
    def __init__(self, Field:str, ConditionOperator:ConditionOperator, LogicalOperator:LogicalOperator, Value, StartGroup:int, EndGroup:int, CaseInsensitive:bool): 
        """
        Initializes the AIDataFilter with the specified field, condition operator, logical operator, value, and grouping information.

        Args:
            Field (str): The field to filter on.
            ConditionOperator (ConditionOperator): The condition operator to apply to the filter.
            LogicalOperator (LogicalOperator): The logical operator to combine this filter with others.
            Value: The value to filter by.
            StartGroup (int): The starting group index for this filter.
            EndGroup (int): The ending group index for this filter.
            CaseInsensitive (bool): Whether the filter should be case insensitive.
        """     
        self.Field = Field
        self.ConditionOperator = ConditionOperator
        self.LogicalOperator = LogicalOperator
        self.Value = Value
        self.StartGroup = StartGroup
        self.EndGroup = EndGroup
        self.CaseInsensitive = CaseInsensitive
    def to_dict(self):
        """
        Converts the AIDataFilter instance to a dictionary representation.

        Returns:
            dict: A dictionary representation of the AIDataFilter instance.
        """

        return {
            "Field": self.Field,
            "ConditionOperator": self.ConditionOperator.value,
            "LogicalOperator": self.LogicalOperator.value,
            "Value": self.Value,
            "StartGroup": self.StartGroup,
            "EndGroup": self.EndGroup,
            "CaseInsensitive": self.CaseInsensitive
        }
    
class AIDataSimpleFilter:
    """ 
    A class representing a simple filter for querying data in the inmydata platform.
    This class encapsulates the details of a simple filter, including the field to filter on and the value to filter by.
    
    Attributes:
        Field (str): The field to filter on.    
        Value: The value to filter by.
    """
    def __init__(self, Field:str, Value):
        """
        Initializes the AIDataSimpleFilter with the specified field and value.

        Args:
            Field (str): The field to filter on.
            Value: The value to filter by.
        """
        self.Field = Field
        self.Value = Value
    def to_dict(self):
        """
        Converts the AIDataSimpleFilter instance to a dictionary representation.
        
        Returns:
            dict: A dictionary representation of the AIDataSimpleFilter instance.
        """
        return {
            "Field": self.Field,
            "Value": self.Value
        }
    
class TopNOption:
    """ 
    A class representing a Top N filter for querying data in the inmydata platform.
    This class encapsulates the details of a Top N filter, including the metric field to calculated the value to filter on, and the number of top or bottom items to return.
    
    Attributes:
        MetricField (str): The field to filter on.
        NumberOfResults (int): The number of top or bottom items to return. Positive numbers indicate top N, negative numbers indicate bottom N.
    """
    def __init__(self, MetricField:str, NumberOfResults:int):
        """
        Initializes the TopNUsed with the specified metric field and number of top or bottom items.

        Args:
            MetricField (str): The field to filter on.
            NumberOfResults (int): The number of top items to return.
        """
        self.MetricField = MetricField
        self.NumberOfResults = NumberOfResults
    def to_dict(self):
        """
        Converts the TopNUsed instance to a dictionary representation.

        Returns:
            dict: A dictionary representation of the TopNUsed instance.
        """
        return {
            "MetricField": self.MetricField,
            "NumberOfResults": self.NumberOfResults
        }

class StructuredDataDriver:
    """ 
    A driver class for interacting with structured data in the inmydata platform.
    This class provides methods to retrieve data based on subjects, fields, and filters.
    It uses the inmydata API to fetch data and returns it as a pandas DataFrame.
    
    Attributes:
        tenant (str): 
            The tenant identifier for the inmydata platform.
        server (str): 
            The server address for the inmydata platform, default is "inmydata.com".
        user (Optional[str]): 
            The user for whom the driver is initialized, if None, no user is set. Useful to identify the user when generating a chart (see https://developer.inmydata.com/a/solutions/articles/36000577995?portalId=36000061664).
        session_id (Optional[str]): 
            The session ID for the driver, if None, no session ID is set. Useful to identify the session when generating a chart (see https://developer.inmydata.com/a/solutions/articles/36000577995?portalId=36000061664).
        api_key (Optional[str]): 
                The API key for authenticating with the inmydata platform. If None, it will attempt to read from the environment variable 'INMYDATA_API_KEY'.
        type (Optional[str]): 
                The type of the data source to use. When blank this will indicate the inmydata platform. If None, it will attempt to read from the environment variable 'INMYDATA_TYPE'.
        logging_level (Optional[int]): 
            The logging level for the logger, default is logging.INFO.
        log_file (Optional[str]): 
            The file to log messages to, if None, logs will be printed to the console.
    """
    class _AIDataAPIRequest:
        def __init__(self, Subject: str, Fields: list[str], Filters: list['AIDataFilter'], TopNUsed: dict['str', 'TopNOption'], SummaryRequest: bool, System: str):
            self.Subject = Subject
            self.Fields = Fields
            self.Filters = Filters  # List of AIDataFilterUsed
            self.TopNUsed = TopNUsed
            self.SummaryRequest = SummaryRequest
            self.System = System
        def to_dict(self):
            return {
                "Subject": self.Subject,
                "Fields": self.Fields,
                "Filters": [f.to_dict() for f in self.Filters],
                "TopNUsed": {k: v.to_dict() for k, v in self.TopNUsed.items()},
                "SummaryRequest": self.SummaryRequest,
                "System": self.System
            } 

    class _AIDataAPIResponse:
        def __init__(self,noRows,fileSize,csvDataString,columnNamesandTypes):      
          self.noRows = noRows
          self.fileSize = fileSize
          self.csvDataString = csvDataString
          self.columnNamesandTypes = columnNamesandTypes
        def toJSON(self):
          return json.dumps(self, default=lambda o: o.__dict__, sort_keys=True, indent=4)
        
    class _AIChartAPIRequest:
        def __init__(self,Subject,RowFields,ColumnsFields,MetricFields,Filters,ChartType,Caption,User,SessionID,TopNUsed):
          self.Subject = Subject
          self.RowFields = RowFields
          self.ColumnFields = ColumnsFields
          self.MetricFields = MetricFields
          self.Filters = Filters
          self.ChartType = ChartType
          self.User = User
          self.Caption = Caption
          self.SessionID = SessionID
          self.TopNUsed = TopNUsed
        def to_dict(self):
            return {
                "Subject": self.Subject,
                "RowFields": self.RowFields,
                "ColumnFields": self.ColumnFields,
                "MetricFields": self.MetricFields,
                "Filters": [f.to_dict() for f in self.Filters], 
                "ChartType": self.ChartType.name.lower(),  
                "User": self.User,
                "Caption": self.Caption,
                "SessionID": self.SessionID,
                "TopNUsed": {k: v.to_dict() for k, v in self.TopNUsed.items()} 
            }
        
    class _AIChartAPIResponse:
        def __init__(self,visualisationID):      
          self.visualisationID = visualisationID
        def toJSON(self):
          return json.dumps(self, default=lambda o: o.__dict__, sort_keys=True, indent=4)

    def __init__(self, tenant: str, server:str ="inmydata.com", user: Optional[str] = None, session_id: Optional[str] = None,  api_key: Optional[str] = None, type: Optional[str] = None , logging_level: Optional[int] = logging.INFO, log_file: Optional[str] = None):
        """
        Initializes the StructuredDataDriver with the specified tenant, server, logging level, and log file.
        
        Args:
            tenant (str): 
                The tenant identifier for the inmydata platform.      
            server (str): 
                The server address for the inmydata platform, default is "inmydata.com".  
            user (Optional[str]): 
                The user for whom the driver is initialized, if None, no user is set. Useful to identify the user when generating a chart (see https://developer.inmydata.com/a/solutions/articles/36000577995?portalId=36000061664).
            session_id (Optional[str]): 
                The session ID for the driver, if None, no session ID is set. Useful to identify the session when generating a chart (see https://developer.inmydata.com/a/solutions/articles/36000577995?portalId=36000061664).
            api_key (Optional[str]): 
                The API key for authenticating with the inmydata platform. If None, it will attempt to read from the environment variable 'INMYDATA_API_KEY'.
            type (str): 
                The type of data source the MCP should use. This will be "" by default which will be inmydata data platform. The alternative is OpenEdge. If None, it will attempt to read from the environment variable 'INMYDATA_TYPE'.               
            logging_level (int): 
                The logging level for the logger, default is logging.INFO.
            log_file (Optional[str]): 
                The file to log messages to, if None, logs will be printed to the console.
        """
        self.server = server
        self.tenant = tenant
        self.user = user
        self.session_id = session_id
        self.type = type

        # Create a logger specific to this class/instance
        self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}.{tenant}")
        if logging_level is None:
            logging_level = logging.INFO
        self.logger.setLevel(logging_level)
        
        # Avoid adding multiple handlers if this gets called multiple times
        if not self.logger.handlers:
            formatter = logging.Formatter('[%(asctime)s] %(levelname)s - %(message)s')

            if log_file:
                handler = logging.FileHandler(log_file)
            else:
                handler = logging.StreamHandler()

            handler.setFormatter(formatter)
            self.logger.addHandler(handler)

        self.logger.propagate = False  # Prevent propagation to the root logger
        
        if api_key:
            self.api_key = api_key
        else:
            try:
               self.api_key = os.environ['INMYDATA_API_KEY']
            except KeyError:
               self.api_key = ""
               self.logger.warning("Environment variable INMYDATA_API_KEY not set. API requests to the inmydata platform will fail.")

        self.logger.info("StructuredDataDriver initialized.")

        pass

    def get_user(self):
        """ 
        Returns the user for whom the driver is initialized.
        
        Returns:
            Optional[str]: The user for whom the driver is initialized, or None if no user is set.
        """
        return self.user
    
    def set_user(self, user: str):
        """ 
        Sets the user for whom the driver is initialized.
        
        Args:
            user (str): The user to set for the driver.
        """
        self.user = user
        self.logger.info(f"User set to {user}")
    
    def get_session_id(self):
        """ 
        Returns the session ID for the driver.
        
        Returns:
            Optional[str]: The session ID for the driver, or None if no session ID is set.
        """
        return self.session_id

    def set_session_id(self, session_id: str):
        """ 
        Sets the session ID for the driver.
        
        Args:
            session_id (str): The session ID to set for the driver.
        """
        self.session_id = session_id
        self.logger.info(f"Session ID set to {session_id}")

    def get_type(self):
        """ 
        Returns the type of data source for the driver.
        
        Returns:
            Optional[str]: The type of data source for the driver, or blank/empty.
        """
        return self.type
    def set_type(self, type: str):
        """ 
        Sets the type of data source for the driver.
        
        Args:
            type (str): The type of data source to set for the driver.
        """
        self.type = type
        self.logger.info(f"Type set to {type}")

    def get_schema(self, source: Optional[str] = None):
        """ Retrieves the schema of the structured data available in the inmydata platform.

        Returns:
            dict: A dictionary representing the schema of the structured data.
        """
        try:
            result = None
            headers = {'Authorization': 'Bearer ' + self.api_key,
                    'Content-Type': 'application/json'}
            url = 'https://' + self.tenant + '.' + self.server + '/api/developer/v1/ai/getapisubjectlistinfo'
            req_body = {"subject": None, "type": self.type}
            x = requests.post(url,headers=headers, json=req_body)
            if x.status_code == 200:    
                try:
                    parsed = json.loads(x.text)
                except json.JSONDecodeError as e:
                    raise RuntimeError(f"Invalid JSON from backend: {e}") from e

                payload = parsed.get("value", parsed) if isinstance(parsed, dict) else parsed

                # Ensure structure we expect
                subjects: List[Dict[str, Any]] = []
                if isinstance(payload, dict) and isinstance(payload.get("subjects"), list):
                    subjects = payload["subjects"]
                elif isinstance(payload, list):
                    # Some backends might directly return a list of subjects
                    subjects = payload
                else:
                    # Be generous but explicit
                    subjects = []

                # Enrich with counts
                for subj in subjects:
                    if not isinstance(subj, dict):
                        continue
                    fact_field_types = subj.get("factFieldTypes") or {}
                    metric_field_types = subj.get("metricFieldTypes") or {}
                    system = subj.get("system") or ""
                    subj["numDimensions"] = len(fact_field_types) if isinstance(fact_field_types, dict) else 0
                    subj["numMetrics"] = len(metric_field_types) if isinstance(metric_field_types, dict) else 0

                result = {
                    "schemaVersion": 1,
                    "generatedAt": datetime.now().isoformat(timespec="seconds") + "Z",
                    "source": source,
                    "subjectsCount": len(subjects),
                    "subjects": subjects,
                }

                # Compact JSON, stable ordering for easier diffs
                return json.dumps(result, ensure_ascii=False, separators=(",", ":"), sort_keys=False)

        except Exception as e:
            # Mirror your C# error string style
            return f"Error retrieving schema: {e}"

    def get_data(self, subject: str, fields: list[str], filters: list[AIDataFilter],SummaryRequest: bool,System: str, TopNUsed: Optional[dict['str', 'TopNOption']] = None) -> pd.DataFrame | None:
            """ Retrieves data from the inmydata platform based on the specified subject, fields, and filters.

            Args:
                subject (str): The subject to query data from.
                fields (list[str]): The list of fields to retrieve.
                filters (list[AIDataFilter]): The list of filters to apply to the query.
                SummaryRequest (bool): A flag to indicate if the data request should use a summary query which will summarize the data based on the fields specified. This is useful when datasets are large and is the default. If SummaryRequest is set to false then it allows data to be read without being summarized.
                System (str): The system to use. This will come from the System property of the subject selected.
                TopNUsed (Optional[dict['str', 'TopNOption']]): Optional dictionary of Top N filters to apply to the query. Defaults to None. They key of the dictionary is the name of the column the top N filter will be applied to.
            Returns:
                pd.DataFrame: A pandas DataFrame containing the retrieved data.
            """
            result = None
            if TopNUsed is None:
                TopNUsed = {}
            aidatareq = self._AIDataAPIRequest(subject,fields,filters,TopNUsed,SummaryRequest,System)
            input_json_string  = jsonpickle.encode(aidatareq.to_dict(), unpicklable=False)
            self.logger.info("Executing " + str(input_json_string))
            if input_json_string is None:
                raise ValueError("input_json_string is None and cannot be loaded as JSON")
            myobj = json.loads(input_json_string)
            headers = {'Authorization': 'Bearer ' + self.api_key,
                    'Content-Type': 'application/json'}
            url = 'https://' + self.tenant + '.' + self.server + '/api/developer/v1/ai/data'
            x = requests.post(url, json=myobj,headers=headers)
            if x.status_code == 200:    
                decoded_response = jsonpickle.decode(x.text)
                if isinstance(decoded_response, dict):
                    value = decoded_response.get("value")
                else:
                    raise ValueError("Decoded response is not a dictionary. Actual type: {}".format(type(decoded_response)))
                if value is None:
                    raise ValueError("Response does not contain 'value' or it is None")
                value_json = jsonpickle.encode(value)
                if value_json is None:
                    raise ValueError("value_json is None and cannot be loaded as JSON")
                aidataresp = self._AIDataAPIResponse(**json.loads(value_json))
                if aidataresp.noRows > 0:            
                  buff = BytesIO(base64.standard_b64decode(aidataresp.csvDataString))
                  with gzip.GzipFile(fileobj=buff) as gz:
                    decompressed_data = gz.read()    
                    data = StringIO(decompressed_data.decode('utf-8'))
                    result = pd.read_csv(filepath_or_buffer = data)
            return result

    def get_data_simple(
            self,
            subject:str,
            fields:list[str],
            simplefilters:list[AIDataSimpleFilter], 
            SummaryRequest: bool,
            System: str,            
            caseSensitive: Optional[bool] = True, 
            TopNUsed: Optional[dict['str', 'TopNOption']] = None) -> pd.DataFrame | None:
            """ 
            Retrieves data from the inmydata platform based on the specified subject, fields, and simple filters.
            
            Args:
                subject (str): The subject to query data from.
                fields (list[str]): The list of fields to retrieve.
                simplefilters (list[AIDataSimpleFilter]): The list of simple filters to apply to the query.
                SummaryRequest (bool): A flag to indicate if the data request should use a summary query which will summarize the data based on the fields specified. This is useful when datasets are large and is the default. If SummaryRequest is set to false then it allows data to be read without being summarized.
                System (str): The system to use. This will come from the System property of the subject selected.                
                caseSensitive (Optional[bool]): Whether the filter should be case sensitive. Defaults to True.
                TopNUsed (Optional[dict['str', 'TopNOption']]): Optional dictionary of Top N filters to apply to the query. Defaults to None. They key of the dictionary is the name of the column the top N filter will be applied to.
            
            Returns:
                pd.DataFrame: A pandas DataFrame containing the retrieved data.
            """
            filters = []
            # Ensure caseSensitive is always a bool
            case_insensitive = bool(caseSensitive) if caseSensitive is not None else True
            for simpleFilter in simplefilters:           
              filter = AIDataFilter(Field=simpleFilter.Field,ConditionOperator=ConditionOperator.Equals,LogicalOperator=LogicalOperator.And,Value=simpleFilter.Value,StartGroup=0,EndGroup=0, CaseInsensitive=case_insensitive)
              filters.append(filter)
            return self.get_data(subject,fields,filters,TopNUsed,SummaryRequest,System)
