from typing import Dict, Union, Literal, List
import pathlib
import os
import json

import polars as pl
import pandas as pd
import duckdb
import s3fs

from yaspin import yaspin

BLADESIGHT_DATASETS_S3_BUCKET = "bladesight-datasets"

def get_path_to_local_bladesight():
    """ This function returns the path to the local datasets folder.
    If there is no environmental variable called BLADESIGHT_DATASETS_PATH, it 
    will return ~/.bladesight.
    """
    if "BLADESIGHT_DATASETS_PATH" in os.environ:
        return pathlib.Path(os.environ["BLADESIGHT_DATASETS_PATH"]) / ".bladesight"
    else:
        return pathlib.Path.home() / ".bladesight"

def get_local_datasets() -> List[str]:
    """ This function returns a list of the names of the datasets in the local
    datasets folder.
    """
    BLADESIGHT_DATASETS_PATH = get_path_to_local_bladesight()
    if not BLADESIGHT_DATASETS_PATH.exists():
        return []
    else:
        local_datasets = []
        for (path_root, _, files) in os.walk(BLADESIGHT_DATASETS_PATH):
            for file in files:
                if file.endswith('.db'):
                    path_parts = pathlib.Path(path_root).parts
                    path_prefix = None
                    add_parts = False
                    for part in path_parts:
                        if add_parts:
                            if path_prefix is None:
                                path_prefix = part
                            else:
                                path_prefix = path_prefix + "/" + part
                        if part == '.bladesight':
                            add_parts = True
                    local_datasets.append(f'{path_prefix}/{file}'[:-3])
        return local_datasets

def get_bladesight_datasets() -> List[str]:
    """ This function returns a list of all the datasets in 
        the bladesight-datasets bucket.

    Returns:
        List[str]: A list of the names of the datasets in the bucket.
    """
    s3 = s3fs.S3FileSystem(anon=True)
    datasets = []
    with yaspin(text='Getting all datasets from Bladesight Data...'):
        for (bucket_root, _, files) in s3.walk(BLADESIGHT_DATASETS_S3_BUCKET+'/'):
            for file in files:
                if file.endswith('.db'):
                    datasets.append(f'{bucket_root}/{file}'[:-3])
    return datasets

def download_dataset_from_bladesight_data(dataset_path_on_s3 : str) -> None:
    """ This function downloads a dataset from S3 and saves it locally.
    """
    s3 = s3fs.S3FileSystem(anon=True)
    PATH_TO_LOCAL_DB = get_path_to_local_bladesight() / "data"
    for s3_subfolder in dataset_path_on_s3.split('/')[1:]:
        PATH_TO_LOCAL_DB = PATH_TO_LOCAL_DB / s3_subfolder
    
    if not PATH_TO_LOCAL_DB.parent.exists():
        PATH_TO_LOCAL_DB.parent.mkdir(parents=True)

    with yaspin(text=f'Downloading {dataset_path_on_s3} from Bladesight Data...') as spinner:
        s3.download(dataset_path_on_s3+'.db', str(PATH_TO_LOCAL_DB) + '.db')
    spinner.text = f'Done downloading {dataset_path_on_s3} from Bladesight Data... '
    spinner.ok('✅ ')

class Dataset:
    """This object is used to access data from a dataset.
    """
    def __init__(self, path : pathlib.Path):
        self.path = path
        if self._confirm_dataset_valid() is False:
            raise ValueError(f'{self.path} is not a valid dataset.')
        self.tables : List[str] = self._get_tables()
        self.metadata : Dict[str, Dict] = self._get_all_metadata()
        self.dataframe_library : Literal['pd', 'pl'] = "pd"
        self.print_citation()
    
    def set_dataframe_library(self, library : Literal['pd', 'pl']):
        """ This function sets the dataframe library to use when returning data.

        Args:
            library (Literal['pd', 'pl']): The dataframe library to use.
        """
        if library in ['pd', 'pl']:
            self.dataframe_library = library
        else:
            raise ValueError("library must be 'pd' or 'pl'")
    
    def _confirm_dataset_valid(self) -> bool:
        """ This function checks if a dataset exists and has a .db extension in s
        """
        if self.path.exists() and self.path.suffix == '.db':
            return True
        return False
    
    def _get_tables(self) -> List[str]:
        """ This method gets the tables in the dataset. It 
        excludes the metadata table.

        Returns:
            List[str]: The tables in the dataset.
        """
        all_tables = self._read_sql("SHOW TABLES;")["name"].to_list() 
        data_tables = list(set(all_tables) - set(['metadata']))
        return data_tables
    
    def __getitem__(self, key : str):
        table_name = key.replace("table/", "")
        if table_name in self.tables:
            return self._read_sql(f"SELECT * FROM {table_name};", return_mode=self.dataframe_library)
        else:
            raise KeyError(
                f"Table {table_name} not found. These are the tables in the dataset: {self.tables}"
            )
    
    def print_citation(self):
        """ Print the citation provided in the metadata table.
        """
        citation = self.metadata['CITATION']
        if 'repr' not in citation.keys():
            print("No citation provided in metadata table.")
            return
        
        cite_main = f"""If you use this dataset in published work, please use the below citation:\n\n{citation['repr']}"""
        print(cite_main)
        
        if 'url' in citation.keys():
            print(f"""\nLink to paper: {citation['url']}""")

        if 'doi' in citation.keys():
            print(f"""\nDOI: {citation['doi']}""")

    def _ipython_key_completions_(self):
        return ["table/" + i for i in self.tables]

    def _read_sql(
        self, 
        sql : str, 
        return_mode : Literal["pd", 'pl'] = 'pd'
    ) -> Union[pd.DataFrame, pl.DataFrame]:
        """ This function executes a SQL query on the dataset
        and returns its results as a pandas or polars DataFrame.

        Args:
            sql (str): The SQL query to execute.
            return_mode (Literal["pd", 'pl'], optional): The return mode. Defaults to 'pd'.

        Returns:
            Union[pd.DataFrame, pl.DataFrame]: The results of the query.
        """
        with duckdb.connect(str(self.path)) as con:
            if return_mode == 'pd':
                df = con.sql(sql).df()
            elif return_mode == 'pl':
                df = con.sql(sql).pl()
        return df
    
    def _get_all_metadata(self) -> Dict[str, Dict]:
        """ This function returns the value of a metadata field.
        
        Returns:
            Dict[str, Dict]: The metadata.
        """
        df_metadata = self._read_sql("SELECT * FROM metadata;")
        metadata = {}
        for _, row in df_metadata.iterrows():
            metadata[row['metadata_key']] = json.loads(row['metadata_value'])
        return metadata
    
    def __repr__(self):
        table_string = "[\n"
        for table in self.tables:
            table_string += f"\t'table/{table}',\n "
        table_string += "]"
        return f"Dataset({self.path}),\n\n Tables: \n {table_string}" 

class BladesightDatasetDirectory:
    def __init__(self):
        self.path = get_path_to_local_bladesight()
        self.local_datasets = [self.replace_path_prefix(i) for i in get_local_datasets()]
        self.online_datasets : List[str] = ...
        try:
            self._ipython_key_completions_()
        except Exception as e:
            print("Could not read remote datasets. Only listing local datasets")
            self.online_datasets = self.local_datasets
    
    def _getitem_key_correct_format(self, key : str) -> bool:
        """ This function checks if the key is in the correct format. The key
        should be in the format "data/intro_to_btt/intro_to_btt_ch02".

        Args:
            key (str): The key to check.

        Returns:
            bool: True if the key is in the correct format, False otherwise.
        """
        if key.startswith("data/"):
            return True
        return False
    
    def __getitem__(self, key : str) -> Dataset:
        """Get the dataset specified by a key. If the dataset is not found, it
        will be downloaded from the Bladesight Data bucket.

        Args:
            key (str): The name of the dataset.

        Raises:
            KeyError: If the dataset is not found.

        Returns:
            Dataset: The dataset.
        """
        if self._getitem_key_correct_format(key) is False:
            raise KeyError(
                f"Dataset {key} does not start with data/. The key should be in the format 'data/../../etc'."
            )
        
        for local_dataset in self.local_datasets:
            homogenized_local_name = self.replace_path_prefix(local_dataset)
            if key == homogenized_local_name:
                return Dataset( self.path / pathlib.Path(local_dataset + '.db'))
        else:
            # Download the dataset from the online datasets
            for online_set in self.online_datasets:
                homogenized_online_name = self.replace_path_prefix(online_set)
                if key == homogenized_online_name:
                    download_dataset_from_bladesight_data(
                        self.replace_path_prefix(
                            key, 
                            BLADESIGHT_DATASETS_S3_BUCKET
                        )
                    )
                    self.local_datasets = get_local_datasets()
                    return self[key]

            else:
                raise KeyError(f"Dataset {key} not found.")
    
    def replace_path_prefix(
        self, 
        dataset_full_path : str, 
        replace_prefix : str = "data"
    ) -> str:
        """ This function is used to replace the first path prefix with the
            replace_prefix argument. For example, if the dataset path 
            is "bladesight-data/intro_to_btt/intro_to_btt_ch02", and the 
            replace_prefix is "data", and the path is returned as 
            "data/intro_to_btt/intro_to_btt_ch02".

        Args:
            dataset_full_path (str): The full path to the dataset.
            replace_prefix (str, optional): The prefix to replace. Defaults to "data".

        Returns:
            str: The new path.
        """

        new_path = [replace_prefix] + dataset_full_path.split('/')[1:]
        return "/".join(new_path)

    def _ipython_key_completions_(self):
        if self.online_datasets is ...:
            self.online_datasets = get_bladesight_datasets()
        return [self.replace_path_prefix(i) for i in self.online_datasets]
    
Datasets = BladesightDatasetDirectory()