import os
import time
import json
import yaml
import logging
import psycopg2
import traceback

from glob import glob
from rich import print, inspect, print_json
from jinja2 import Environment, FileSystemLoader
from pathlib import Path
from decouple import config


class NoConnectionToCITROSDBException(Exception):
    def __init__(self, message="No connection to citros db."):
        super().__init__(message)


class BatchUploader:
    """
    The Uploader class is responsible for uploading data to CITROS db.


    assuming structure:
        .citros/data/{simulation name}/{batch name}/{simulation run id}/bag/{bag file}

    Attributes:
    -----------
    root : str
        The root directory where the data is located. (default to .citros/data)
    simulation : str
        The simulation name.
    batch : str
        The batch name.
    version : str
        The version of the uploader.

    log : Optional
        The log object for logging messages.
    verbose : bool
        Flag indicating whether to enable verbose mode.
    debug : bool
        Flag indicating whether to enable debug mode.

    """

    ###################
    ##### private #####
    ###################

    def _get_citros_bags(self, path: str) -> list:
        """
        Search for `.citros.mcap` and `.citros.db3` bag files in a specified directory.

        This method walks through the directory structure rooted at `path` to locate
        files ending with `.citros.mcap` or `.citros.db3`. It returns a list containing
        the full paths to these files.

        Parameters:
        -----------
        path : str
            The root directory where the search will begin.

        Returns:
        --------
        list
            A list of full file paths to the bag files. An exception is raised if no
            bag files are found.

        Raises:
        -------
        Exception
            If no `.citros.mcap` or `.citros.db3` bag files are found.

        Examples:
        ---------
        >>> _get_citros_bags('/path/to/root')
        ['/path/to/root/folder1/bag1.citros.mcap', '/path/to/root/folder2/bag2.citros.db3']

        """
        bags = []
        for root, dirs, files in os.walk(path):
            for file in files:
                if file.endswith(".citros.mcap"):
                    self.log.debug(f"Found MCAP bag: {file}")
                elif file.endswith(".citros.db3"):
                    self.log.debug(f"Found SQLITE3 bag: {file}")
                else:
                    continue
                bags.append(os.path.join(root, file))
        # if len(bags) < 1:
        #     raise Exception(
        #         f"Didn't find SQLITE3 or MCAP bag in the [{path}] folder ..."
        #     )

        return bags

    def _get_custom_message(self, directory_path: str) -> list:
        """
        Search for and return the paths of folders containing custom message types.

        This method searches through the directory specified by `directory_path` to
        find any folders that contain custom message types. These folders are assumed
        to contain a 'msg/' subdirectory.

        Parameters:
        -----------
        directory_path : str
            The directory path where the method will search for custom message types.

        Returns:
        --------
        list
            A list of strings, each being the path of a folder containing custom message types.

        Raises:
        -------
        Exception
            If any error occurs while accessing the file system.

        Examples:
        ---------
        >>> _get_custom_message('/path/to/directory')
        # ['/path/to/directory/custom_msgs1/msg/', '/path/to/directory/custom_msgs2/msg/']

        """
        try:
            packages = [
                folder
                for folder in os.listdir(directory_path)
                if os.path.isdir(os.path.join(directory_path, folder))
            ]
            msgs = [directory_path + package + "/msg/" for package in packages]
            if len(msgs) > 0:
                self.log.debug(f"Found custom message in {msgs}")
        except Exception:
            msgs = []
        return msgs

    def _get_parameters(self, directory_path: str) -> list:
        """
        Search for and return the paths of parameter files in the given directory.

        This method scans the directory specified by `directory_path` to find parameter files.
        It looks for files ending with ".json", ".yaml", or ".yml".

        Parameters:
        -----------
        directory_path : str
            The directory path where the method will search for parameter files.

        Returns:
        --------
        list
            A list of strings, each being the path of a parameter file.

        Logs:
        -----
        Debug log if no parameters are found in the specified directory.

        Examples:
        ---------
        >>> _get_parameters('/path/to/directory')
        # ['/path/to/directory/params1.json', '/path/to/directory/params2.yaml']

        """
        parameters = []
        for root, dirs, files in os.walk(directory_path):
            for file in files:
                if file.endswith(".json"):
                    self.log.debug(f"Found json parameters file: {file}")
                elif file.endswith(".yaml") or file.endswith(".yml"):
                    self.log.debug(f"Found yaml parameters file: {file}")
                else:
                    continue
                parameters.append(os.path.join(root, file))
        if len(parameters) < 1:
            self.log.debug(
                f"Didn't find parameters in the [{directory_path}] folder ..."
            )

        return parameters

    @staticmethod
    def _read_parameters_file(file_path: str) -> dict:
        """
        Read the content of a parameter file and return it as a JSON-formatted dictionary.

        This method reads a parameter file specified by `file_path`, converts its content to
        a Python dictionary, and then returns the dictionary as a JSON-formatted string.
        It can read both JSON and YAML formatted parameter files.

        Parameters:
        -----------
        file_path : str
            The path to the parameter file to be read. The file should have either a .json, .yaml, or .yml extension.

        Returns:
        --------
        dict
            A dictionary containing the parameters read from the file, serialized to JSON format.

        Raises:
        -------
        Exception if the file cannot be read or parsed.

        Examples:
        ---------
        >>> _read_parameters_file('/path/to/params.json')
        # '{"key": "value"}'

        >>> _read_parameters_file('/path/to/params.yaml')
        # '{"key": "value"}'

        """
        _, ext = os.path.splitext(file_path)
        with open(file_path, "r") as f:
            if ext == ".json":
                return json.dumps(json.load(f))
            elif ext == ".yaml" or ext == ".yml":
                return json.dumps(yaml.safe_load(f))

    def upload_bag_to_pg(
        self,
        connection,
        schema_name: str,
        table_name: str,
        sid: int,
        bag: str,
        msgs: list,
    ):
        """
        Upload a bag file to a PostgreSQL database.

        This method reads messages from a bag file using the `BagReaderCustomMessages` class and uploads them to
        a PostgreSQL database table. The method handles both successful and unsuccessful uploads, logging
        appropriate messages in both cases.

        Parameters:
        ----------
        cursor : psycopg2.cursor
            The cursor object for interacting with the PostgreSQL database.

        connection : psycopg2.connection
            The connection object for the PostgreSQL database.

        schema_name : str

        table_name : str

        sid : str
            The ID of the simulation run to which this bag file belongs.

        bag : str
            The path to the bag file that needs to be uploaded.

        msgs : list
            List of paths to custom message directories (if any).

        Returns:
        -------
        tuple
            A tuple containing:
            - A boolean indicating whether the upload was successful.
            - A status message string.
            - None if successful, or an error message string if unsuccessful.

        Raises:
        ------
        Exception, psycopg2.Error:
            Any exceptions raised during database interactions are caught and logged.

        Examples:
        --------
        >>> upload_bag_to_pg(cursor, connection, 'batch_001', 'sim_001', '/path/to/bag', ['/path/to/custom/msg'])
        # (True, 'Success, uploaded [/path/to/bag] to Postgres. [size: 100 MB]', None)
        """
        from .ros import BagReaderCustomMessages

        cursor = connection.cursor()
        try:
            start_time = time.time()
            self.log.info(f"Uploading {bag} to PG")
            bagReader = BagReaderCustomMessages(
                msgs,
                log=self.log,
                debug=self.debug,
                verbose=self.verbose,
            )
            total_size = 0
            for buffer in bagReader.read_messages(bag, sid):
                size = buffer.seek(0, os.SEEK_END)
                size = buffer.tell()
                total_size = total_size + size
                buffer.seek(0)
                self.log.debug(
                    f" \tInserting buffer size: {((size / 1024 ) / 1024):.2f} MB"
                )
                if size == 0:
                    continue
                cursor.execute(f"SET search_path TO {schema_name}")
                try:
                    cursor.copy_from(
                        buffer,
                        table_name,
                        sep=chr(0x1E),
                        null="",
                        columns=["sid", "rid", "time", "topic", "type", "data"],
                    )
                except (Exception, psycopg2.Error) as error:
                    buffer.seek(0)
                    self.log.error(f"buffer = {buffer.getvalue()}")
                    self.log.error(
                        f" Failed to insert record into table, aborting upload to DB.",
                        error,
                    )
                    self.log.exception(error)
                    return False, "Got exception from pgdb", str(error)
                connection.commit()

            self.log.debug(
                f"Done uploading {bag}, took {(time.time() - start_time):.3f} [sec]"
            )
            return (
                True,
                f"Success, uploaded [{bag}] to Postgres. [size: {(total_size / 1024)/1024} MB]",
                None,
            )
        except (Exception, psycopg2.Error) as error:
            connection.commit()
            self.log.exception(
                f" Failed to insert record into table, aborting upload to DB.", error
            )
            self.log.exception(traceback.format_exc())
            return False, "Got exception from pgdb", str(error)
        # finally:
        #     # closing database connection.
        #     if connection:
        #         cursor.close()
        #         connection.close()
        #         self.log.debug(f"PostgreSQL connection is closed")
        #         logging.shutdown()

    def upload_parameters_to_pg(
        self,
        connection,
        schema_name: str,
        table_name: str,
        sid: str,
        parameter_file: str,
    ):
        """
        Upload simulation parameters to a PostgreSQL database table.

        This function takes a dictionary of parameters and uploads them to a specified table in a PostgreSQL database.
        If the upload is successful, the method will commit the changes and return a success message.

        Parameters:
        ----------

        connection : psycopg2.connection
            The connection object for the PostgreSQL database.

        params_dict : dict
            A dictionary containing key-value pairs of parameters to upload. Keys are parameter names as strings,
            and values are the corresponding parameter values.

        table_name : str, optional
            The name of the database table where parameters should be uploaded. Default is 'parameters'.

        Returns:
        -------
        tuple
            A tuple containing:
            - A boolean indicating whether the upload was successful.
            - A status message string.
            - None if successful, or an error message string if unsuccessful.

        Raises:
        ------
        psycopg2.Error:
            Any exceptions raised during database interactions are caught and logged.

        Examples:
        --------
        >>> params = {'param1': 1, 'param2': 2}
        >>> upload_parameters_to_pg(cursor, connection, params)
        # (True, 'Parameters uploaded successfully to table [parameters].', None)
        """
        try:
            self.log.debug(f"Uploading parameter file: {parameter_file}")
            parameter_dict = self._read_parameters_file(parameter_file)
            self.log.debug(f"Parameter:\n[{json.dumps(parameter_dict, indent=2)}]")
            record_to_insert = {
                "sid": sid,
                "rid": 0,
                "time": 0,
                "topic": "/config",
                "type": parameter_file,
                "data": parameter_dict,
            }

            postgres_insert_query = f""" 
                insert into {schema_name}."{table_name}"
                (sid, rid, time, topic, type, data)
                values (%(sid)s, %(rid)s, %(time)s, %(topic)s, %(type)s, %(data)s);
            """

            cursor = connection.cursor()
            cursor.execute(postgres_insert_query, record_to_insert)
            connection.commit()

        except Exception as e:
            self.log.error(e)

        # finally:
        #     if connection:
        #         cursor.close()
        #         connection.close()
        #         self.log.debug(f"PostgreSQL connection is closed")
        #         logging.shutdown()

    def upload(self):
        self.log.debug(f"{self.__class__.__name__}.upload()")

        # inspect(self)
        # print(f"self.batch_dir: {self.batch_dir}")

        from .database import CitrosDB

        citrosDB = CitrosDB(log=self.log, debug=self.debug, verbose=self.verbose)

        schema_name = f"{self.simulation_name}"
        table_name = f"{self.name}"

        connection = citrosDB.connect()
        if connection is None:
            self.log.error("No creating connection to database. Aborting.")
            raise NoConnectionToCITROSDBException

        citrosDB.create_table(connection, schema_name, table_name, self.version)

        for sid_path in glob(f"{self.batch_dir}/*/"):
            if sid_path.endswith("/"):
                sid_path = sid_path[:-1]

            sid = sid_path.split("/")[-1]

            # upload all parameters
            parameters = self._get_parameters(f"{self.batch_dir}/{sid}/config/")
            # print(parameters)
            for parameter in parameters:
                try:
                    # cursor = connection.cursor()
                    self.upload_parameters_to_pg(
                        connection,
                        schema_name=schema_name,
                        table_name=table_name,
                        sid=sid,
                        parameter_file=parameter,
                    )
                    connection.commit()
                except Exception as e:
                    self.log.error(e)
                    connection = citrosDB.connect()

            # uplaod bags
            bags = self._get_citros_bags(f"{self.batch_dir}/{sid}/bags/")
            if len(bags) < 1:
                self.log.error(
                    f"No bags found in {self.batch_dir}/{sid}/bags/ . Perhaps the run didnt finish properly."
                )
                continue

            # print(bags)
            msgs = self._get_custom_message(f"{self.batch_dir}/{sid}/msgs/")
            # print(msgs)

            for bag in bags:
                try:
                    status, message, error = self.upload_bag_to_pg(
                        connection,
                        schema_name=schema_name,
                        table_name=table_name,
                        sid=sid,
                        bag=bag,
                        msgs=msgs,
                    )
                    self.log.debug(
                        f"status: {status}, message: {message}, error: {error}"
                    )
                except Exception as e:
                    self.log.error(e)
                    connection = citrosDB.connect()

        citrosDB.hot_reload_set_status(
            connection, schema_name, table_name, self.version, "LOADED"
        )
        # cursor.close()
        connection.commit()
        connection.close()

    def unload(self):
        from .database import CitrosDB

        citrosDB = CitrosDB(log=self.log, debug=self.debug, verbose=self.verbose)

        schema_name = f"{self.simulation_name}"
        table_name = f"{self.name}"

        connection = citrosDB.connect()
        if connection is None:
            self.log.error("No creating connection to database. Aborting.")
            raise NoConnectionToCITROSDBException

        citrosDB.drop_table(connection, schema_name, table_name)

        connection.close()
