from enum import Enum


class DatasetType(Enum):
    """Dataset type enumeration matching Dataframer backend."""
    SINGLE_FILE = "SINGLE_FILE"
    MULTI_FILE = "MULTI_FILE"
    MULTI_FOLDER = "MULTI_FOLDER"


class FileType(Enum):
    """File type enumeration matching Dataframer backend."""
    MD = "md"
    TXT = "txt"
    CSV = "csv"
    PDF = "pdf"
    JSON = "json"
    JSONL = "jsonl"
     

class DatabricksConnector:
    """
    Databricks connector for Dataframer workflows.
    
    This class provides methods to interact with Databricks SQL, fetch sample data,
    and load generated data into Databricks tables.
    
    Parameters
    ----------
    dbutils : DBUtils
        The dbutils object from your Databricks notebook context.
        This is automatically available in Databricks notebooks.
    
    Examples
    --------
    >>> databricks_connector = DatabricksConnector(dbutils)
    >>> df = databricks_connector.fetch_sample_data(
    ...     num_items_to_select=25,
    ...     table_name="samples.bakehouse.media_customer_reviews"
    ... )
    >>> df.head()
    """
    
    def __init__(self, dbutils):
        """
        Initialize the Databricks connector.
        
        Parameters
        ----------
        dbutils : DBUtils
            The dbutils object from your Databricks notebook context.
        """
        self.dbutils = dbutils
    
    def get_connection(self):
        """
        Return an authenticated Databricks SQL connection.
        
        Returns
        -------
        Connection
            A Databricks SQL connection object.
        """
        from databricks import sql
        from databricks.sdk.core import Config, oauth_service_principal

        server_hostname = self.dbutils.secrets.get("dataframer", "DATABRICKS_SERVER_HOSTNAME")
        http_path = self.dbutils.secrets.get("dataframer", "DATABRICKS_HTTP_PATH")

        def credential_provider():
            config = Config(
                host=f"https://{server_hostname}",
                client_id=self.dbutils.secrets.get("dataframer", "DATABRICKS_CLIENT_ID"),
                client_secret=self.dbutils.secrets.get("dataframer", "DATABRICKS_CLIENT_SECRET"),
            )
            return oauth_service_principal(config)

        return sql.connect(
            server_hostname=server_hostname,
            http_path=http_path,
            credentials_provider=credential_provider,
            user_agent_entry="dataframer_user_agent",
        )
    
    def fetch_sample_data(self, num_items_to_select, table_name):
        """
        Fetch sample data from a Databricks table and return it as a Pandas DataFrame.
        
        Parameters
        ----------
        num_items_to_select : int
            Number of rows to fetch from the table.
        table_name : str
            Fully qualified table name (e.g., "catalog.schema.table").
        
        Returns
        -------
        pd.DataFrame
            A Pandas DataFrame containing the sample data.
        
        Examples
        --------
        >>> databricks_connector = DatabricksConnector(dbutils)
        >>> df = databricks_connector.fetch_sample_data(
        ...     num_items_to_select=25,
        ...     table_name="samples.bakehouse.media_customer_reviews"
        ... )
        >>> df.head()
        """
        import pandas as pd

        query = f"""
            SELECT *
            FROM {table_name}
            LIMIT {num_items_to_select}
        """

        try:
            with self.get_connection() as connection:
                with connection.cursor() as cursor:
                    cursor.execute(query)
                    rows = cursor.fetchall()
                    columns = [desc[0] for desc in cursor.description]
        except Exception as e:
            error_msg = f"Failed to fetch data from table `{table_name}`"
            print(f"{error_msg}: {str(e)}")
            print("Verify table exists, is accessible, and you have SELECT permissions")
            raise RuntimeError(f"{error_msg}: {str(e)}") from e

        return pd.DataFrame(rows, columns=columns)
    
    def load_generated_data(self, table_name, downloaded_zip, dataset_type, file_type):
        """
        Load generated samples from a ZIP file into a Databricks table using Databricks SQL.

        Parameters
        ----------
        table_name : str
            Target table name (e.g., "catalog.schema.table")
        downloaded_zip : file-like
            ZIP file object containing the generated data file
        dataset_type : DatasetType
            Type of dataset structure (DatasetType.SINGLE_FILE, DatasetType.MULTI_FILE, or DatasetType.MULTI_FOLDER)
        file_type : FileType
            Type of file in the ZIP (FileType.CSV, FileType.JSON, FileType.JSONL, etc.)
        
        Examples
        --------
        >>> databricks_connector = DatabricksConnector(dbutils)
        >>> with open("samples.zip", "rb") as f:
        ...     databricks_connector.load_generated_data(
        ...         table_name="my_catalog.my_schema.my_table",
        ...         downloaded_zip=f,
        ...         dataset_type=DatasetType.SINGLE_FILE,
        ...         file_type=FileType.CSV
        ...     )
        """
        import zipfile
        import pandas as pd
        from io import BytesIO
        
        zip_buffer = BytesIO(downloaded_zip.read())

        if dataset_type == DatasetType.SINGLE_FILE:
            try:
                with zipfile.ZipFile(zip_buffer) as z:
                    file_list = z.namelist()
                    
                    generated_data_files = [f for f in file_list if f.lower().endswith(f'.{file_type.value}')]
                    
                    if len(generated_data_files) != 1:
                        error_msg = f"Expected exactly one .{file_type.value} file in ZIP"
                        print(f"{error_msg}. Available files: {file_list}")
                        raise ValueError(error_msg)
                    
                    data_filename = generated_data_files[0]
                    data_bytes = z.read(data_filename)
                    print(f"Found {file_type.value} file: {data_filename}")
            
            except zipfile.BadZipFile as e:
                error_msg = "Invalid or corrupted ZIP file"
                print(f"{error_msg}: {str(e)}")
                raise ValueError(f"{error_msg}: {str(e)}") from e
            except ValueError:
                raise
            except Exception as e:
                error_msg = "Failed to extract file from ZIP"
                print(f"{error_msg}: {str(e)}")
                raise RuntimeError(f"{error_msg}: {str(e)}") from e

            if file_type == FileType.CSV:
                pandas_df = pd.read_csv(BytesIO(data_bytes))
            elif file_type == FileType.JSON:
                # TODO: Implement JSON file handling
                pass
            elif file_type == FileType.JSONL:
                # TODO: Implement JSONL file handling
                pass
            else:
                raise ValueError(f"Unsupported file_type: {file_type}. Supported: CSV, JSON, JSONL for SINGLE_FILE datasets")
            
            with self.get_connection() as connection:
                cursor = connection.cursor()

                columns_sql = ", ".join(
                    f"`{col}` STRING" for col in pandas_df.columns
                )

                try:
                    cursor.execute(f"""
                        CREATE OR REPLACE TABLE {table_name} (
                            {columns_sql}
                        )
                    """)
                except Exception as e:
                    error_msg = f"Failed to create table `{table_name}`"
                    print(f"{error_msg}: {str(e)}")
                    print("Verify table name format (catalog.schema.table), permissions, and warehouse is running")
                    cursor.close()
                    raise RuntimeError(f"{error_msg}: {str(e)}") from e

                insert_sql = f"""
                    INSERT INTO {table_name}
                    VALUES ({", ".join(["?"] * len(pandas_df.columns))})
                """

                try:
                    cursor.executemany(
                        insert_sql,
                        pandas_df.values.tolist()
                    )
                except Exception as e:
                    error_msg = f"Failed to insert data into table `{table_name}`"
                    print(f"{error_msg}: {str(e)} | Rows attempted: {len(pandas_df)}")
                    cursor.close()
                    raise RuntimeError(f"{error_msg}: {str(e)}") from e

                cursor.close()

            print(f"✅ Table `{table_name}` saved successfully using Databricks SQL")
        
        elif dataset_type == DatasetType.MULTI_FILE:
            # TODO: Implement MULTI_FILE handling
            pass
        
        elif dataset_type == DatasetType.MULTI_FOLDER:
            # TODO: Implement MULTI_FOLDER handling
            pass
        
        else:
            raise ValueError(f"Invalid dataset_type: {dataset_type}. Expected DatasetType enum")
