import time
import subprocess
from collections import Counter
from typing import Optional

from PySide6.QtCore import QObject, QThread, Signal, Slot
from PySide6.QtGui import QIcon
from PySide6.QtWidgets import (
    QProgressDialog, QPushButton, QVBoxLayout, QDialog, QLabel, QComboBox, QHBoxLayout, QMessageBox
)
from PySide6.QtCore import Qt

try:
    import docker
    from docker.errors import DockerException, APIError
    DOCKER_SDK_AVAILABLE = True
except ImportError:
    DOCKER_SDK_AVAILABLE = False
    docker = None
    DockerException = Exception
    APIError = Exception

from src.controller.switch_validate_controller import SwitchValidateController
from src.path_resolver import PathResolver
from src.controller.terminal_controller import TerminalController


class SimpleContainer:
    """
    A lightweight wrapper for Docker containers when using subprocess fallback.
    """
    def __init__(self, container_id):
        self.id = container_id
        self.short_id = container_id[:12]

    def stop(self, timeout=10):
        subprocess.run(
            ["docker", "stop", "-t", str(timeout), self.id],
            check=True,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL
        )


class DockerController:
    """
    Provides static methods to manage Docker containers and images.

    Uses Docker SDK for Python when available for better performance,
    falls back to subprocess if SDK is not installed.
    """
    _docker_check_cache = {
        "last_refresh": 0.0,
        "docker_running": False,
        "installed_repos": set(),
        "running_repos": Counter()
    }
    _cache_min_interval = 10.0
    _client: Optional["docker.DockerClient"] = None
    _connection_failed = False
    _last_connection_attempt = 0.0
    _conn_retry_interval = 60.0  # Wait 60s before retrying if it failed

    # Cache icons to avoid repeated disk I/O
    _check_icon: Optional[QIcon] = None
    _cross_icon: Optional[QIcon] = None
    _working_icon: Optional[QIcon] = None
    _not_working_icon: Optional[QIcon] = None
    _normalized_names_cache: dict[str, str] = {}


    @staticmethod
    def _now() -> float:
        return time.monotonic()

    @staticmethod
    def _get_client() -> Optional["docker.DockerClient"]:
        """
        Get or create a Docker client instance.
        Returns None if Docker is not available.
        """
        if not DOCKER_SDK_AVAILABLE:
            print("Docker SDK not available - install docker package")
            return None

        if DockerController._client is None:
            now = DockerController._now()
            if DockerController._connection_failed and (now - DockerController._last_connection_attempt < DockerController._conn_retry_interval):
                return None

            DockerController._last_connection_attempt = now
            try:
                print("Creating new Docker client...")
                DockerController._client = docker.from_env()
                print("Docker client created, testing connection...")
                DockerController._client.ping()
                print("Docker client connected successfully")
                DockerController._connection_failed = False
            except DockerException:
                # Silently fail to connect to Docker daemon to avoid log spam
                # print(f"Failed to connect to Docker daemon: {type(e).__name__}: {e}")
                # print(f"DOCKER_HOST env: {os.environ.get('DOCKER_HOST', 'not set')}")
                DockerController._client = None
                DockerController._connection_failed = True
            except Exception:
                # print(f"Unexpected error connecting to Docker: {type(e).__name__}: {e}")
                DockerController._client = None

        return DockerController._client

    @staticmethod
    def _get_normalized_repo(image_ref: str) -> str:
        """
        Get normalized repo name from cache or compute it.
        """
        if image_ref not in DockerController._normalized_names_cache:
            DockerController._normalized_names_cache[image_ref] = DockerController._normalize_repo_name(image_ref)
        return DockerController._normalized_names_cache[image_ref]

    @staticmethod
    def _normalize_repo_name(image_ref: str) -> str:
        """
        Normalize a Docker image reference to a repository path without registry and tag.
        """
        ref = image_ref.strip()
        if '@' in ref:
            ref = ref.split('@', 1)[0]
        last_slash = ref.rfind('/')
        last_colon = ref.rfind(':')
        if last_colon > last_slash:
            ref_no_tag = ref[:last_colon]
        else:
            ref_no_tag = ref
        parts = ref_no_tag.split('/')
        if parts and ('.' in parts[0] or ':' in parts[0]):
            parts = parts[1:]
        return '/'.join(parts) if parts else ref_no_tag

    @staticmethod
    def refresh_cache(force: bool = False) -> None:
        """
        Refresh the Docker cache asynchronously in a background thread.
        Non-blocking — starts a QThread and returns immediately.
        UI is updated when the thread finishes via `refresh_mgr.cache_updated` signal.
        """
        now = DockerController._now()
        if (not force and
                (now - DockerController._docker_check_cache["last_refresh"]) < DockerController._cache_min_interval):
            return
        refresh_mgr.start()

    @staticmethod
    def run_container(context):
        """
        Run the Docker container for the selected tool with a specified switch.
        """
        current_item = context.ui.tools_list.currentItem()
        if current_item is not None:

            current_index = context.ui.tools_list.indexOfTopLevelItem(current_item)
            current_tool = context.tools[current_index]
            current_card = context.ui.tool_cards[current_index]

            docker_image = current_tool["docker_image"]
            default_port = current_tool.get("default_port")
            switch = current_card.get_params()
            tool_name = current_tool["name"].lower()

            if not SwitchValidateController.validate_switch(switch, tool_name):

                msg = QMessageBox()
                msg.setIcon(QMessageBox.Icon.Critical)
                msg.setText("Invalid switch")
                msg.setWindowTitle("Error")
                msg.exec()
                return

            import shlex
            command = ["docker", "run", "--rm", "-it"]
            if default_port:
                command.extend(shlex.split(default_port))
            command.append(docker_image)
            if switch:
                command.extend(shlex.split(switch))

            TerminalController.run(command)
        else:
            msg = QMessageBox()
            msg.setIcon(QMessageBox.Icon.Warning)
            msg.setText("Please select a valid tool.")
            msg.setInformativeText(
                "You need to select a tool from the list to run the Docker container.")
            msg.setWindowTitle("Warning")
            msg.setStandardButtons(QMessageBox.StandardButton.Ok)
            msg.exec()

    @staticmethod
    def check_container_status(tool, tool_item, run_button):
        """
        Check if Docker is running and if the Docker image is downloaded.
        """
        docker_image = tool["docker_image"]

        # Cache refresh is controlled by _cache_min_interval
        DockerController.refresh_cache()
        cache = DockerController._docker_check_cache

        target_repo = DockerController._get_normalized_repo(docker_image)

        # Lazy load icons
        if DockerController._check_icon is None:
            DockerController._check_icon = QIcon(PathResolver.resource_path('resources/assets/check.png'))
        if DockerController._cross_icon is None:
            DockerController._cross_icon = QIcon(PathResolver.resource_path('resources/assets/cross.png'))

        installed = target_repo in cache["installed_repos"]
        try:
            icon = DockerController._check_icon if installed else DockerController._cross_icon
            if tool_item.icon(0).cacheKey() != icon.cacheKey():
                tool_item.setIcon(0, icon)

        except Exception as e:
            print(f"Failed to set tool item icon/status for {docker_image}: {e}")

        try:
            running_count = cache["running_repos"].get(target_repo, 0)
            if running_count > 0:
                new_text = f"Running ({running_count})" if running_count > 1 else "Running"
            else:
                new_text = "Run Container"
            
            if run_button.text() != new_text:
                run_button.setText(new_text)
        except Exception as e:
            print(f"Failed to update run button status for {docker_image}: {e}")
            run_button.setText(f"Container status check is not possible\n {e}")

    @staticmethod
    def _get_containers_by_image(docker_image: str) -> list:
        """
        Get the list of container IDs running a specific image.
        Tries Docker SDK first, falls back to subprocess if SDK is unavailable or fails.
        """
        matched = []
        client = DockerController._get_client()
        
        if client is not None:
            try:
                target_norm = DockerController._get_normalized_repo(docker_image)
                all_containers = client.containers.list()
                for container in all_containers:
                    image_tags = container.image.tags
                    if not image_tags:
                        # If no tags, check against image ID
                        if container.image.short_id == docker_image or container.image.id == docker_image:
                            matched.append(container)
                    else:
                        # Check each tag using normalization logic
                        for tag in image_tags:
                            if DockerController._get_normalized_repo(tag) == target_norm:
                                matched.append(container)
                                break
            except (DockerException, APIError):
                client = None # Mark as failed to trigger fallback

        # Subprocess fallback (if SDK failed or unavailable)
        if not matched and (client is None):
             matched = DockerController._get_containers_by_image_subprocess(docker_image)
             
        return matched

    @staticmethod
    def _get_containers_by_image_subprocess(docker_image: str) -> list:
        target_norm = DockerController._get_normalized_repo(docker_image)
        matched = []
        try:
            cmd = ['docker', 'ps', '--no-trunc', '--format', '{{.ID}}|{{.Image}}']
            result = subprocess.run(cmd, capture_output=True, text=True)
            if result.returncode != 0:
                return []
            
            for line in result.stdout.splitlines():
                if '|' not in line: continue
                c_id, c_img = line.split('|', 1)
                c_img = c_img.strip()
                c_id = c_id.strip()
                
                # Normalize the image name from docker ps
                # c_img might be, for example, "parrotsec/set:latest"
                if DockerController._get_normalized_repo(c_img) == target_norm:
                     matched.append(SimpleContainer(c_id))
                     
        except Exception as e:
            print(f"Fallback check failed: {e}")
            
        return matched

    @staticmethod
    def stop_container(tool, parent_widget):
        """
        Stop the Docker container associated with the specified tool.
        """
        docker_image = tool["docker_image"]
        progress_dialog = QProgressDialog("Stopping container...", "Cancel", 0, 100, parent_widget)
        progress_dialog.setWindowModality(Qt.WindowModality.WindowModal)
        progress_dialog.setAutoClose(False)
        progress_dialog.setAutoReset(False)
        progress_dialog.show()

        try:
            containers = DockerController._get_containers_by_image(docker_image)
            progress_dialog.setValue(30)

            if containers:
                container_to_stop = None

                if len(containers) > 1:
                    dialog = QDialog(parent_widget)
                    dialog.setWindowTitle("Select Container")
                    layout = QVBoxLayout()

                    label = QLabel("Choose container ID to stop:")
                    layout.addWidget(label)

                    combo_box = QComboBox()
                    combo_box.addItems([c.short_id for c in containers])
                    layout.addWidget(combo_box)

                    button_layout = QHBoxLayout()
                    stop_all_button = QPushButton("Stop All Containers")
                    ok_button = QPushButton("Ok")
                    cancel_button = QPushButton("Cancel")

                    button_layout.addWidget(stop_all_button)
                    button_layout.addWidget(cancel_button)
                    button_layout.addWidget(ok_button)
                    layout.addLayout(button_layout)

                    def handle_stop_all():
                        dialog.accept()
                        DockerController.stop_all_containers(tool, parent_widget)

                    stop_all_button.clicked.connect(handle_stop_all)
                    ok_button.clicked.connect(dialog.accept)
                    cancel_button.clicked.connect(dialog.reject)

                    dialog.setLayout(layout)
                    dialog.exec()

                    if dialog.result() == QDialog.DialogCode.Accepted:
                        selected_id = combo_box.currentText()
                        container_to_stop = next(
                            (c for c in containers if c.short_id == selected_id), None
                        )
                else:
                    container_to_stop = containers[0]

                if container_to_stop:
                    try:
                        container_to_stop.stop(timeout=10)
                        progress_dialog.setValue(100)
                        print(f"Container {container_to_stop.short_id} stopped successfully.")
                    except Exception as e:
                        print(f"Failed to stop container: {e}")
                else:
                    print("No container selected to stop.")
            else:
                print("No running container found for this tool.")

        except Exception as e:
            print(f"Failed to stop container: {e}")
        finally:
            progress_dialog.close()

    @staticmethod
    def stop_all_containers(tool, parent_widget):
        """
        Stop all Docker containers associated with the specified tool.
        """
        docker_image = tool["docker_image"]
        progress_dialog = QProgressDialog("Stopping all containers...", "Cancel", 0, 100, parent_widget)
        progress_dialog.setWindowModality(Qt.WindowModality.WindowModal)
        progress_dialog.setAutoClose(False)
        progress_dialog.setAutoReset(False)
        progress_dialog.show()

        try:
            containers = DockerController._get_containers_by_image(docker_image)
            progress_dialog.setValue(30)

            if containers:
                for i, container in enumerate(containers):
                    try:
                        container.stop(timeout=10)
                        progress_dialog.setValue(30 + int(70 * (i + 1) / len(containers)))
                        print(f"Container {container.short_id} stopped successfully.")
                    except Exception as e:
                        print(f"Failed to stop container {container.short_id}: {e}")
                progress_dialog.setValue(100)
            else:
                print("No running container found for this tool.")

        except Exception as e:
            print(f"Failed to stop containers: {e}")
        finally:
            progress_dialog.close()


class _DockerWorker(QObject):
    """
    Performs Docker I/O (images.list, containers.list, ping) on a background thread.
    Emits finished() with results when done.
    """
    finished = Signal(bool, set, object)  # docker_running, installed_repos, running_repos (Counter)

    def __init__(self):
        super().__init__()

    @Slot()
    def refresh(self):
        docker_running = False
        installed_repos = set()
        running_repos = Counter()

        client = DockerController._get_client()

        if client is not None:
            try:
                client.ping()
                docker_running = True

                try:
                    images = client.images.list()
                    for image in images:
                        for tag in image.tags:
                            repo = tag.split(':')[0] if ':' in tag else tag
                            installed_repos.add(DockerController._get_normalized_repo(repo))
                except APIError as e:
                    print(f"Failed to retrieve Docker images via SDK: {e}")

                try:
                    containers = client.containers.list()
                    for container in containers:
                        img = container.image.tags[0] if container.image.tags else container.image.short_id
                        running_repos[DockerController._get_normalized_repo(img)] += 1
                except APIError as e:
                    print(f"Failed to retrieve running containers via SDK: {e}")

            except (DockerException, APIError) as e:
                print(f"Docker connection lost: {e}")
                docker_running = False
                DockerController._client = None
            except Exception as e:
                print(f"Unexpected error checking Docker status: {e}")
                docker_running = False
                DockerController._client = None

        if not installed_repos:
            try:
                result = subprocess.run(
                    ['docker', 'images', '--format', '{{.Repository}}'],
                    capture_output=True, text=True, check=False
                )
                if result.returncode == 0:
                    docker_running = True
                    for line in result.stdout.splitlines():
                        repo = line.strip()
                        if repo and repo != '<none>':
                            installed_repos.add(DockerController._get_normalized_repo(repo))
            except (FileNotFoundError, subprocess.SubprocessError):
                pass

        self.finished.emit(docker_running, installed_repos, running_repos)


class _RefreshManager(QObject):
    """
    Manages background QThread lifecycle for Docker cache refreshes.
    Emits cache_updated when fresh data is available.
    """
    cache_updated = Signal()

    def __init__(self):
        super().__init__()
        self._thread: QThread | None = None
        self._refresh_in_progress = False

    def start(self):
        if self._refresh_in_progress:
            return
        self._refresh_in_progress = True
        self._thread = QThread()
        worker = _DockerWorker()
        worker.moveToThread(self._thread)
        self._thread.started.connect(worker.refresh)
        worker.finished.connect(self._on_result)
        worker.finished.connect(self._thread.quit)
        worker.finished.connect(worker.deleteLater)
        self._thread.finished.connect(self._on_thread_finished)
        self._thread.start()

    def _on_result(self, docker_running, installed_repos, running_repos):
        DockerController._docker_check_cache.update({
            "last_refresh": time.monotonic(),
            "docker_running": docker_running,
            "installed_repos": installed_repos,
            "running_repos": running_repos,
        })
        self.cache_updated.emit()

    def _on_thread_finished(self):
        thread = self._thread
        self._thread = None
        self._refresh_in_progress = False
        thread.deleteLater()


refresh_mgr = _RefreshManager()