import pandas as pd
from tqdm import tqdm
import requests
import io


class SUNTLoader:
    def __init__(self):
        self.__BASE_URL = "https://raw.githubusercontent.com/LabIA-UFBA/SUNT/main/data/"
        self.__ROUTES_URL = "https://raw.githubusercontent.com/LabIA-UFBA/SUNT/main/data/raw/GTFS_INTEGRA_SALVADOR2/"
        self.__HUGFACE_URL_BOARDING = "https://huggingface.co/datasets/labiaufba/PublicTransportationSunt/resolve/main/Boarding/"
        self.__HUGFACE_URL_OD = "https://huggingface.co/datasets/labiaufba/PublicTransportationSunt/resolve/main/OD/"
        self.__HUGFACE_URL_GTFS = "https://huggingface.co/datasets/labiaufba/SSA_StopBusTimeSeries_5/resolve/main/gtfs/"
        self.__HUGFACE_URL_LEADING = "https://huggingface.co/datasets/labiaufba/PublicTransportationSunt/resolve/main/Alighting/"
        self.__DEFAULT_START = "2024-03-01"
        self._DATASETS_TYPE = {
            'od':{"Name": "OD", "Description": "Features estimated full trips by passengers, enabling analysis of bus loading at stops and along routes.", "URL": self.__HUGFACE_URL_OD, "origin": "huggingface"},
            'gtfs-stops':{"Name": "GTFS Stops", "Description": "Information about each bus stop.", "URL": self.__HUGFACE_URL_GTFS, "origin": "huggingface"},
            'gtfs-trips':{"Name": "GTFS Trips", "Description": "Information about trips and their corresponding paths.", "URL": self.__HUGFACE_URL_GTFS, "origin": "huggingface"},
            'gtfs-stop-times':{"Name": "GTFS Stop Times", "Description": "Chronological order of bus stops where each trip paused.", "URL": self.__HUGFACE_URL_GTFS, "origin": "huggingface"},
            'gtfs-routes':{"Name": "GTFS Routes", "Description": " Information about bus lines.", "URL": self.__HUGFACE_URL_GTFS, "origin": "huggingface"},
            'gtfs-shapes':{"Name": "GTFS Shapes", "Description": "GTFS shapes data", "URL": self.__HUGFACE_URL_GTFS, "origin": "huggingface"},
            'gtfs-agency':{"Name": "GTFS Agency", "Description": "Information about the bus companies.", "URL": self.__HUGFACE_URL_GTFS, "origin": "huggingface"},
            'boarding': {"Name": "Boarding", "Description": "A comprehensive table integrating data from AFC, LTI, AVL, and GTFS sources", "URL": self.__HUGFACE_URL_BOARDING, "origin": "huggingface"},
            'alighting': {"Name": "Alighting", "Description": "Features estimated alighting information inferred from passenger behavior, including walking metrics and transfer patterns.", "URL": self.__HUGFACE_URL_LEADING, "origin": "huggingface"},
            'afc': {"Name": "AFC", "Description": "Features information regarding ticket validations by passengers across the public transportation network.", "URL": "AFC/", "origin": "mendeley"},
            'avl-lines': {"Name": "AVL Lines", "Description": "Features information concerning vehicles’ routes and bus schedules", "URL": "AVL/", "origin": "mendeley"},
            'avl-vehicles': {"Name": "AVL Vehicles", "Description": "Features information concerning vehicles’ routes and bus schedules over time.", "URL": "AVL/", "origin": "mendeley"},
            'lti': {"Name": "LTI", "Description": "Maps the start and end of each trip made by a vehicle on a specific route.", "URL": "LTI/", "origin": "mendeley"},
            # 'edge': {"Name": "Edge", "Description": "Features information about connections between stops/stations, including distance, geospatial data, average speed, trip time, and passenger loading.", "URL": "LTI/", "origin": "mendeley"},
            # 'graph-based': {"Name": "Graph-Based", "Description": "Features information about the graph structure of the public transportation network.", "URL": "LTI/", "origin": "mendeley"},
            # | `edge` | Features information about connections between stops/stations, including distance, geospatial data, average speed, trip time, and passenger loading. |
            # | `graph-based` | Describes static features of transportation nodes (e.g., stops and stations), averaged across all data. |   
        }

    def get_datasets(self):
        campos = ["Name", "Description"]
        return {k: {field: dados[field] for field in campos} for k, dados in self._DATASETS_TYPE.items()}

    def _filter_days(self, dates, day_type):
        df = pd.DataFrame({'date': dates})
        df['dow'] = df['date'].dt.dayofweek
        if day_type == 'workdays': return df[df['dow'] < 5]['date']
        if day_type == 'saturdays': return df[df['dow'] == 5]['date']
        if day_type == 'sundays': return df[df['dow'] == 6]['date']
        return df['date']

    def _dataset_sunt_mendeley(self, path_file):
        # dataset fixo
        DATASET_ID = "85fdtx3kr5"
        url = f"https://data.mendeley.com/public-api/datasets/{DATASET_ID}/"
        response = requests.get(url)
        response.raise_for_status()
        data = response.json()
        files = data.get("files", [])
        path_file_lower = path_file.lower()
        file_info = None
        # Localiza o arquivo na lista
        for f in files:
            filename = f['filename'].lower()
            # Verifica se o nome do arquivo corresponde ao solicitado
            if filename == path_file_lower or filename.endswith(path_file_lower):
                file_info = f
                break
        if file_info is None:
            raise Exception(f"File '{path_file}' not found in dataset")
        download_url = file_info["content_details"]["download_url"]
        r = requests.get(download_url)
        r.raise_for_status()
        file_bytes = io.BytesIO(r.content)
        if path_file.endswith(".parquet"):
            df = pd.read_parquet(file_bytes)
        elif path_file.endswith(".csv") or path_file.endswith(".txt"):
            df = pd.read_csv(file_bytes)
        elif path_file.endswith(".json"):
            df = pd.read_json(file_bytes)
        else:
            raise Exception("Incompatible file format")

        return df

    def load_batch(self, dataset_type, start_date=None, periods=1, freq='D', day_type='all'):
        dataset_type = dataset_type.lower()
        if dataset_type not in self._DATASETS_TYPE:
            raise ValueError(f"Dataset type '{dataset_type}' not Found. Options: {list(self._DATASETS_TYPE.keys())}")

        start = start_date if start_date else self.__DEFAULT_START
        dates = pd.date_range(start=start, periods=periods, freq=freq)
        selected_dates = self._filter_days(dates, day_type)
        dataframes = []
        for dt in tqdm(selected_dates, desc=f"Carregando {dataset_type}"):
            y, m, d = dt.year, f"{dt.month:02d}", f"{dt.day:02d}"
            paths = {
                'od': f"{self._DATASETS_TYPE['od']['URL']}od-{y}-{m}-{d}.parquet",                
                'gtfs-stops': f"{self._DATASETS_TYPE['gtfs-stops']['URL']}stops.txt",
                'gtfs-trips': f"{self._DATASETS_TYPE['gtfs-trips']['URL']}trips.txt",
                'gtfs-stop-times': f"{self._DATASETS_TYPE['gtfs-stop-times']['URL']}stop_times.txt",                
                'gtfs-routes': f"{self._DATASETS_TYPE['gtfs-routes']['URL']}routes.txt",                
                'gtfs-shapes': f"{self._DATASETS_TYPE['gtfs-shapes']['URL']}shapes.txt",  
                'gtfs-agency': f"{self._DATASETS_TYPE['gtfs-agency']['URL']}agency.txt",
                "boarding": f"{self._DATASETS_TYPE['boarding']['URL']}boarding-{y}-{m}-{d}.parquet",
                "alighting": f"{self._DATASETS_TYPE['alighting']['URL']}alighting-{y}-{m}-{d}.parquet",
                "afc": f"afc-{y}-{m}-{d}.parquet",
                "avl-lines": f"avl-lines-{y}-{m}-{d}.parquet",
                "avl-vehicles": f"avl-vehicles-{y}-{m}-{d}.parquet",
                "lti": f"lti-{y}-{m}-{d}.parquet",
            }
            try:
                url = paths[dataset_type]
                if self._DATASETS_TYPE[dataset_type]['origin'] == 'mendeley':
                    df = self._dataset_sunt_mendeley(path_file=url)
                else:
                    if url.endswith('.parquet'):
                        df = pd.read_parquet(url)
                    elif url.endswith('.csv'):
                        df = pd.read_csv(url)
                    elif url.endswith('.txt'):
                        df = pd.read_csv(url)
                    else:
                        raise ValueError(f"Error : {url}")
                df['date_ref'] = dt
                dataframes.append(df)
            except Exception as e:
                print(f"Error loading {url}: {e}")
                continue
        return pd.concat(dataframes, ignore_index=True) if dataframes else pd.DataFrame()
