import glob
import os
import queue
import threading
from functools import partial
from typing_extensions import Dict
import boto3

import tqdm
from boto3.s3.transfer import TransferConfig
from botocore.client import BaseClient
from rich import print

from kleinkram.api_client import AuthenticatedClient


class TransferCallback:
    """
    Handle callbacks from the transfer manager.

    The transfer manager periodically calls the __call__ method throughout
    the upload process so that it can take action, such as displaying progress
    to the user and collecting data about the transfer.
    """

    def __init__(self):
        """
        Initialize the TransferCallback.

        This initializes an empty dictionary to hold progress bars for each file.
        """
        self._lock = threading.Lock()
        self.file_progress = {}

    def add_file(self, file_id, target_size):
        """
        Add a new file to track.

        :param file_id: A unique identifier for the file (e.g., file name or ID).
        :param target_size: The total size of the file being transferred.
        """
        with self._lock:
            tqdm_instance = tqdm.tqdm(
                total=target_size,
                unit="B",
                unit_scale=True,
                desc=f"Uploading {file_id}",
            )
            self.file_progress[file_id] = {
                "tqdm": tqdm_instance,
                "total_transferred": 0,
            }

    def __call__(self, file_id, bytes_transferred):
        """
        The callback method that is called by the transfer manager.

        Display progress during file transfer and collect per-thread transfer
        data. This method can be called by multiple threads, so shared instance
        data is protected by a thread lock.

        :param file_id: The identifier of the file being transferred.
        :param bytes_transferred: The number of bytes transferred in this call.
        """
        with self._lock:
            if file_id in self.file_progress:
                progress = self.file_progress[file_id]
                progress["total_transferred"] += bytes_transferred

                # Update tqdm progress bar
                progress["tqdm"].update(bytes_transferred)

    def close(self):
        """Close all tqdm progress bars."""
        with self._lock:
            for progress in self.file_progress.values():
                progress["tqdm"].close()


def create_transfer_callback(callback_instance, file_id):
    """
    Factory function to create a partial function for TransferCallback.
    :param callback_instance: Instance of TransferCallback.
    :param file_id: The unique identifier for the file.
    :return: A callable that can be passed as a callback to boto3's upload_file method.
    """
    return partial(callback_instance.__call__, file_id)


def expand_and_match(path_pattern):
    expanded_path = os.path.expanduser(path_pattern)
    expanded_path = os.path.expandvars(expanded_path)

    normalized_path = os.path.normpath(expanded_path)

    if "**" in normalized_path:
        file_list = glob.glob(normalized_path, recursive=True)
    else:
        file_list = glob.glob(normalized_path)

    return file_list


def uploadFiles(files: Dict[str, str], credentials: Dict[str, str], nrThreads: int):
    client = AuthenticatedClient()

    session = boto3.Session(
        aws_access_key_id=credentials["accessKey"],
        aws_secret_access_key=credentials["secretKey"],
        aws_session_token=credentials["sessionToken"],
    )
    api_endpoint = client.tokenfile.endpoint
    if api_endpoint == "http://localhost:3000":
        minio_endpoint = "http://localhost:9000"
    else:
        minio_endpoint = api_endpoint.replace("api", "minio")
    s3 = session.resource("s3", endpoint_url=minio_endpoint)

    _queue = queue.Queue()
    for file in files.items():
        _queue.put(file)
    threads = []
    transferCallback = TransferCallback()

    for i in range(nrThreads):
        thread = threading.Thread(
            target=uploadFile, args=(_queue, s3, transferCallback)
        )
        thread.start()
        threads.append(thread)
    for thread in threads:
        thread.join()


def uploadFile(_queue: queue.Queue, s3: BaseClient, transferCallback: TransferCallback):
    while True:
        try:
            filename, _file = _queue.get(timeout=3)
            queueUUID = _file["queueUUID"]
            filepath = _file["filepath"]
            bucket = _file["bucket"]
            target_location = _file["location"]
            config = TransferConfig(
                multipart_chunksize=10 * 1024 * 1024, max_concurrency=5
            )
            with open(filepath, "rb") as f:
                size = os.path.getsize(filepath)
                transferCallback.add_file(filename, size)
                callback_function = create_transfer_callback(transferCallback, filename)
                s3.Bucket(bucket).upload_file(
                    filepath, target_location, Config=config, Callback=callback_function
                )

                client = AuthenticatedClient()
                res = client.post("/queue/confirmUpload", json={"uuid": queueUUID})
                res.raise_for_status()
            _queue.task_done()
        except queue.Empty:
            break
        except Exception as e:
            print(f"Error uploading {filename}: {e}")
            _queue.task_done()


if __name__ == "__main__":
    res = expand_and_match(
        "~/Downloads/dodo_mission_2024_02_08-20240408T074313Z-003/**.bag"
    )
    print(res)
