from __future__ import annotations

import json
import math
import os
import time
from typing import Any, Callable

from qgis.core import Qgis, QgsApplication, QgsTask
from qgis.PyQt.QtCore import QByteArray, QEventLoop, QFile, QIODevice, QObject, QUrl, pyqtSignal
from qgis.PyQt.QtNetwork import QNetworkAccessManager, QNetworkRequest

from ...compat import NET_CONTENT_LENGTH_HEADER, NET_NO_ERROR, NET_OP_CANCELED_ERROR, QGSTASK_CAN_CANCEL, QGSTASK_SILENT
from ..util import MULTIPART_CHUNK_SIZE

MAX_PROGRESS_INTERVAL = 1  # seconds

# QT5 -> QT6 compatibility for QIODevice.OpenMode type
# Qt5 exposes QIODevice.OpenMode/ReadOnly, Qt6 uses OpenModeFlag.ReadOnly.
if hasattr(QIODevice, "OpenModeFlag"):
    OPEN_MODE_TYPE = QIODevice.OpenModeFlag
    READ_ONLY_MODE = QIODevice.OpenModeFlag.ReadOnly
else:
    OPEN_MODE_TYPE = QIODevice.OpenMode
    READ_ONLY_MODE = QIODevice.ReadOnly


class PartialFile(QFile):
    def __init__(self, filepath: str, start: int, end: int, log_callback: Callable | None = None, progress_callback: Callable | None = None):
        super().__init__(filepath)
        self.log = log_callback
        self.progress_callback = progress_callback
        self.start = start
        self.end = end
        self.current_pos = start
        self.part_size = end - start

    def open(self, mode: Any) -> bool:
        fileopen = super().open(mode)
        self.seek(self.start)
        return fileopen

    def close(self) -> None:
        return super().close()

    def readData(self, maxlen: int) -> bytes:
        actual_len = maxlen

        if self.current_pos + maxlen > self.end:
            actual_len = self.end - self.current_pos

        if actual_len <= 0:
            return QByteArray()

        self.current_pos += actual_len
        self.progress_callback(actual_len)
        return super().readData(actual_len)  # Call read method of superclass


class UploadMultiPartFileTask(QgsTask):
    cancelled = pyqtSignal()

    def __init__(self, rel_path: str, abs_path: str, urls: list[str], ext_prog_callback: Callable[[int], None] | None = None, log_callback: Callable[[int], None] | None = None, retries=5):
        super().__init__(f"Upload {rel_path}", QGSTASK_CAN_CANCEL | QGSTASK_SILENT)
        self.rel_path = rel_path
        self.file_path = abs_path
        self.log = log_callback
        self.ext_prog_callback = ext_prog_callback
        self.urls = urls
        self.uploaded_size = 0
        self.allowed_retries = retries
        self.retry_count = 0
        self.nam = QNetworkAccessManager()
        self.total_size = os.path.getsize(abs_path)
        self.chunk_size = MULTIPART_CHUNK_SIZE
        self.chunks = math.ceil(self.total_size / self.chunk_size)
        self.error = None

    def cancel(self) -> None:
        """Implements a really simple cancel method that just emits a signal when the task is cancelled"""
        self.file_upload_log(f"Cancelling task: {self.description()}", Qgis.Info)
        super().cancel()
        self.cancelled.emit()

    def file_upload_log(self, message: str, level: int, context_obj=None) -> None:
        if self.log:
            log_str = "UploadMultiPartFileTask: " + message
            # indent everuything by 6 spaces
            spacer = " " * 6
            log_str = spacer + log_str.replace(os.linesep, "\n" + spacer)
            self.log(log_str, level, context_obj)

    def debug_log(self) -> str:
        """Useful helper function for printing task state to a log file

        Returns:
            str: _description_
        """
        debug_obj = {
            "urls": self.urls,
            "uploaded_size": self.uploaded_size,
            "total_size": self.total_size,
            "chunk_size": self.chunk_size,
            "chunks": self.chunks,
            "retry_count": self.retry_count,
            "allowed_retries": self.allowed_retries,
            "errors": str(self.error),
        }
        json_str = json.dumps(debug_obj, indent=4, sort_keys=True)
        # Replace all \n line breaks with a newline character
        json_str = json_str.replace(os.linesep, "\n                ")
        return json_str

    def _progress_callback(self, bytesSent: int) -> None:
        """When we get a progress update from the task, we update the total uploaded size"""
        # self.file_upload_log(f"      Progress callback: {bytesSent} of {self.file_path}", Qgis.Info)
        self.uploaded_size += bytesSent
        if self.ext_prog_callback:
            self.ext_prog_callback()

    def upload_file_part(self, url: str, start: int, end: int) -> bool:
        """This is the actual upload code that will physically send all or part of a local file to a single remote URL

        Args:
            url (str): _description_
            start (int): 0 for single file or the first chunk of a multipart upload
            end (int): The last byte of the chunk to upload. End of file for single file or the last byte of the chunk for multipart

        Returns:
            bool: _description_
        """

        # Keep track of the original size before upload so we can reset if the upload fails
        original_size = self.uploaded_size

        for _ in range(self.allowed_retries):
            self.error = None
            partial_file = None
            handle_cancelled_fn = None
            try:
                loop = QEventLoop()
                part_size = end - start
                self.file_upload_log(f"Uploading chunk bytes {start:,} - {end:,} for file: {self.rel_path} Retry: {self.retry_count}", Qgis.Info)

                request = QNetworkRequest(QUrl(url))
                request.setHeader(NET_CONTENT_LENGTH_HEADER, part_size)

                partial_file = PartialFile(self.file_path, start, end, self.file_upload_log, self._progress_callback)
                partial_file.open(READ_ONLY_MODE)

                # Start the actual Call
                self.reply = self.nam.put(request, partial_file)

                def handle_cancelled(pf=partial_file, lp=loop):
                    """Make sure we clean up and close file handles if the cancel signal has been called"""
                    self.error = NET_OP_CANCELED_ERROR
                    self.reply.abort()
                    # Quit the loop directly as a safety measure — abort() triggering
                    # finished → handle_done is not guaranteed if the reply is already
                    # done (abort() on a finished reply is a no-op in Qt).
                    lp.quit()

                # Handle the network request completion
                def handle_done(pf=partial_file, lp=loop):
                    # Detect if the reply was cancelled
                    if self.reply.error() == NET_OP_CANCELED_ERROR:
                        self.error = NET_OP_CANCELED_ERROR
                        lp.quit()
                    elif self.reply.error() != NET_NO_ERROR:
                        # Capture the error and quit the loop cleanly.  Do NOT raise
                        # an exception here: in PyQt6, exceptions raised inside a Qt
                        # slot called from C++ (i.e. from within loop.exec()) are
                        # caught by PyQt's trampoline, NOT propagated through exec().
                        # They are dispatched to sys.excepthook instead, which in QGIS
                        # shows an error dialog — and because lp.quit() is never
                        # reached the event loop keeps spinning, causing repeat dialogs.
                        self.error = self.reply.errorString()
                        lp.quit()
                    else:
                        self.reply.close()
                        pf.close()
                        lp.quit()

                # If the task is cancelled, abort the upload
                handle_cancelled_fn = handle_cancelled
                self.cancelled.connect(handle_cancelled)
                self.reply.finished.connect(handle_done)

                # Wait for the upload to complete or for the task to be cancelled
                loop.exec()

            except Exception as e:
                self.uploaded_size = original_size
                self.error = f"Error while uploading chunk {start}-{end} to {url}: {e!s}"
                self.file_upload_log(f"{self.error}", Qgis.Info)
            finally:
                # Always disconnect handle_cancelled to prevent connections accumulating
                # across retries and chunks.  Without this, every retry adds another
                # closure to self.cancelled, so a cancel signal would fire N×retries
                # handler calls, each calling self.reply.abort() on the live reply.
                if handle_cancelled_fn is not None:
                    try:
                        self.cancelled.disconnect(handle_cancelled_fn)
                    except Exception:
                        pass
                # Ensure the partial file is closed even if handle_done didn't run
                if partial_file is not None and partial_file.isOpen():
                    partial_file.close()

            # Evaluate outcome after the loop (and after finally cleanup)
            if self.error == NET_OP_CANCELED_ERROR:
                # Cancelled — don't retry
                return False
            elif self.error:
                self.uploaded_size = original_size
                self.file_upload_log(f"ERROR: uploading chunk {start}-{end} to {url}: {self.error}", Qgis.Critical)
                self.retry_count += 1
                time.sleep(1)  # Brief pause before retrying
                # fall through to next iteration
            else:
                self.file_upload_log(f"SUCCESS: Finished uploading chunk {start:,}-{end:,}", Qgis.Info)
                return True

        return False

    def run(self) -> bool:
        """Implements the QgsTask run method. This is where the actual work is done

        Returns:
            _type_: _description_
        """

        # Quick check to make sure our chunk math is right
        if self.chunks != len(self.urls):
            self.error = f"Number of URLs ({len(self.urls)}) does not match number of chunks ({self.chunks})"
            return False

        # For each URL (1 for single file, many for multipart) upload the chunk
        for idx, url in enumerate(self.urls):
            # self.file_upload_log(f"START: Uploading chunk {idx + 1}:{self.chunks} for file: {self.rel_path}", Qgis.Info)
            start = idx * self.chunk_size
            end = (idx + 1) * self.chunk_size if idx < len(self.urls) - 1 else self.total_size
            if not self.upload_file_part(url, start, end):
                return False  # Stop if the upload fails after retries (or is cancelled)
            # self.file_upload_log(f"DONE: Uploaded chunk {idx + 1}:{self.chunks} for file: {self.rel_path}", Qgis.Info)
        return True


class UploadQueue(QObject):
    """_summary_

    Returns:
        _type_: _description_
    """

    progress_signal = pyqtSignal(str, int, int, int)
    complete_signal = pyqtSignal()
    cancelled_signal = pyqtSignal()

    MAX_CONCURRENT_UPLOADS = 4

    def __init__(self, log_callback: Callable | None = None):
        super().__init__()
        self.log_callback = log_callback
        self.active = True
        self.last_progress = None
        self.cancelled = False

        self.queue: list[UploadMultiPartFileTask] = []
        self.active_tasks: list[UploadMultiPartFileTask] = []
        self.completed_tasks: list[UploadMultiPartFileTask] = []
        self.cancelled_tasks: list[UploadMultiPartFileTask] = []
        self._handled_task_ids: set[int] = set()  # guard against double task_finished calls
        self._cancelling: bool = False  # guard against re-entrant cancel_all calls

    def reset(self) -> None:
        self.queue = []
        self.active_tasks = []
        self.completed_tasks = []
        self.cancelled_tasks = []
        self.last_progress = None
        self.cancelled = False
        self.active = True
        self._handled_task_ids = set()
        self._cancelling = False

    def queue_logger(self, message: str, level: int, context_obj=None) -> None:
        if self.log_callback:
            log_str = "UploadQueue: " + message
            # indent everuything by 4 spaces
            spacer = " " * 4
            log_str = spacer + log_str.replace(os.linesep, "\n" + spacer)
            self.log_callback(log_str, level, context_obj)

    def enqueue(self, rel_path: str, abs_path: str, upload_urls: list[str], retries: int = 5) -> None:
        """Push a file onto the queue for upload

        Args:
            rel_path(_type_): Relpath is used as an id of sorts here. Not for any other reason
            abs_path(_type_): _description_
            upload_url(_type_): _description_
            retries(int, optional): _description_. Defaults to 5.
        """
        self.queue_logger(f"Enqueued {rel_path} for upload", Qgis.Info)
        # Hook into the task's finished signal
        task = UploadMultiPartFileTask(rel_path, abs_path, upload_urls, self.get_overall_status, self.log_callback, retries)
        task.taskCompleted.connect(lambda: self.task_finished(task))
        task.taskTerminated.connect(lambda: self.task_finished(task))

        self.queue.append(task)
        self.active = True
        self.process_queue()

    def get_overall_status(self, force=False) -> tuple[float, UploadMultiPartFileTask, int, int]:
        """
        Returns the overall status of the queue
        Status is the percent of the total bytes sent
        """
        if not force and self.last_progress is not None and (time.time() - self.last_progress) < MAX_PROGRESS_INTERVAL:
            return
        self.last_progress = time.time()

        total_size = 0
        uploaded_size = 0

        # First add all the tasks in the queue. These are all assumed to be 0% uplaoded
        for task in self.queue:
            total_size += task.total_size

        for task in self.completed_tasks:
            total_size += task.total_size
            uploaded_size += task.total_size

        # Now add all the active tasks
        for task in self.active_tasks:
            total_size += task.total_size
            uploaded_size += task.uploaded_size

        big_task = self.get_biggest_active_task()
        progress = int((uploaded_size / total_size) * 100) if total_size > 0 else 0

        self.progress_signal.emit(big_task.rel_path if big_task else "...", progress, uploaded_size, total_size)

        return (progress, big_task, uploaded_size, total_size)

    def get_biggest_active_task(self) -> UploadMultiPartFileTask | None:
        """Pull task with the biggest file size off the stack and return it

        Returns:
            _type_: _description_
        """
        if self.active_tasks and self.active_tasks[0].total_size > 0:
            return self.active_tasks[0]
        else:
            return None

    def process_queue(self) -> None:
        """Here we process the queue and start tasks as slots become available"""
        while self.active and len(self.active_tasks) < self.MAX_CONCURRENT_UPLOADS:
            # There are things to queue and slots open
            if len(self.queue) > 0:
                task = self.queue.pop(0)
                self.queue_logger(f"Starting upload of {task.rel_path}", Qgis.Info)
                self.active_tasks.append(task)
                QgsApplication.taskManager().addTask(task)

            # If the queue is empty and there are no active tasks, we're done
            elif len(self.queue) == 0 and len(self.active_tasks) == 0:
                self.queue_logger("Queue is empty, stopping", Qgis.Info)
                self.active = False
                self.complete_signal.emit()
            else:
                # There are some active tasks but the queue is empty
                # Nothing to do but wait until the next task finishes and this function is called again
                break

    def task_finished(self, task: UploadMultiPartFileTask):
        """When a single task finishes we need to clean up and process the queue

        Args:
            task(_type_): _description_
        """
        # Both taskCompleted and taskTerminated are connected to this slot; QGIS can
        # fire both for the same task in edge cases (e.g. task finishes just as a
        # cancel is issued).  Guard with an id-based set so each task is handled once.
        task_id = id(task)
        if task_id in self._handled_task_ids:
            self.queue_logger(f"Skipping duplicate task_finished for: {task.rel_path}", Qgis.Warning)
            return
        self._handled_task_ids.add(task_id)

        if task.error == NET_OP_CANCELED_ERROR:
            self.queue_logger(f"Cancelled upload of {task.rel_path}", Qgis.Warning)
            # Still remove from active_tasks so the queue can progress / complete correctly
            try:
                self.active_tasks.remove(task)
                self.cancelled_tasks.append(task)
            except ValueError:
                pass  # already removed (e.g. queue was reset)
        elif task.error:
            self.queue_logger(f"Error uploading {task.rel_path}: {task.error}", Qgis.Critical, task)
            try:
                self.active_tasks.remove(task)
                self.cancelled_tasks.append(task)
            except ValueError:
                pass
        else:
            self.queue_logger(f"Finished uploading: {task.rel_path}", Qgis.Info)
            try:
                self.active_tasks.remove(task)
                self.completed_tasks.append(task)
            except ValueError:
                pass

        # Kick off queue processing again
        self.process_queue()

    def cancel_all(self) -> None:
        # Guard against re-entrant calls (e.g. user double-clicking Stop while the
        # nested QEventLoop is already running).  A second call would create another
        # event loop with its own tasks_need_cancel counter that might never be reached,
        # leaving QGIS processing events indefinitely (the "infinite dialogs" scenario).
        if self._cancelling:
            self.queue_logger("cancel_all called while already cancelling – ignoring", Qgis.Warning)
            return
        self._cancelling = True

        self.queue_logger(f"Canceling all {len(self.active_tasks)} active + {len(self.queue)} queued uploads", Qgis.Info)
        # Shut down the queue processor
        self.active = False
        self.cancelled = True

        # Tasks sitting in self.queue have never been handed to the QGIS task manager,
        # so taskCompleted / taskTerminated will NEVER fire for them.  Just drop them
        # from the queue directly — they never ran and need no async waiting.
        for task in list(self.queue):
            self.queue_logger(f"Discarding unstarted queued task: {task.rel_path}", Qgis.Info)
            self.cancelled_tasks.append(task)
        self.queue.clear()

        # Drain the queue
        finished_tasks = 0
        tasks_need_cancel = 0
        loop_quit_called = False

        # Now QLoop wait for all tasks to be finished
        loop = QEventLoop()

        def on_task_finished():
            nonlocal finished_tasks, loop_quit_called
            finished_tasks += 1
            self.queue_logger(f"Finished cancelling {finished_tasks} of {tasks_need_cancel} uploads", Qgis.Info)
            # Use >= (not ==) so a double-signal firing the count past the target value
            # doesn't leave the event loop hanging forever.
            if not loop_quit_called and finished_tasks >= tasks_need_cancel:
                loop_quit_called = True
                self.queue_logger("All uploads cancelled", Qgis.Info)
                loop.quit()

        # Only wait for active tasks — they are registered with the task manager and
        # will emit taskCompleted / taskTerminated when done.
        for task in list(self.active_tasks):
            if not task.isCanceled():
                self.queue_logger(f"Cancelling upload of {task.rel_path}", Qgis.Info)
                # Connect to both signals so we catch whichever one fires
                task.taskCompleted.connect(on_task_finished)
                task.taskTerminated.connect(on_task_finished)
                task.cancel()
                tasks_need_cancel += 1

        if tasks_need_cancel == 0:
            self._cancelling = False
            self.cancelled_signal.emit()
            return
        loop.exec()

        # Now signal that we're done
        self._cancelling = False
        self.queue_logger("Uploads cancelled", Qgis.Info)
        self.cancelled_signal.emit()
