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
from IPython.display import display, HTML

BLADESIGHT_DATASETS_S3_BUCKET = "bladesight-datasets"

def is_notebook() -> bool:
    try:
        shell = get_ipython().__class__.__name__
        if shell == 'ZMQInteractiveShell':
            return True   # Jupyter notebook or qtconsole
        elif shell == 'TerminalInteractiveShell':
            return False  # Terminal running IPython
        else:
            return False  # Other type (?)
    except NameError:
        return False      # Probably standard Python interpreter

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() / "datasets"
    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']
        cite_html = f"""<p>If you use this dataset in published work, please use the below citation:</p>
        <p><b>{citation['repr']}</b></p>
        """
        if 'url' in citation.keys():
            cite_html +=  f""" <p><a href="{citation['url']}">Link to paper</a></p>"""

        if 'doi' in citation.keys():
            cite_html +=  f""" <p>DOI: { citation['doi'] } </p>"""
        if is_notebook():
            display(HTML(cite_html))
        else:
            print(cite_html)
    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

class BladesightDatasetDirectory:
    def __init__(self):
        self.path = get_path_to_local_bladesight()
        self.local_datasets = 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__(self, key : str):
        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
            download_dataset_from_bladesight_data(
                self.replace_path_prefix(
                    key, 
                    BLADESIGHT_DATASETS_S3_BUCKET
                )
            )
            self.local_datasets = get_local_datasets()
            return self[key]
        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()