"""
This module contains classes for tracking machine and action resources.
"""

import os
import subprocess
import logging
from datetime import datetime, timezone
import psutil
import docker
from typing import List, Tuple, Dict, Optional
from matrice_compute.instance_utils import (
    has_gpu,
    get_gpu_info,
    calculate_time_difference,
)
from matrice_compute.scaling import Scaling
from matrice_common.utils import log_errors


class ResourcesTracker:
    """Tracks machine and container resources."""

    def __init__(self) -> None:
        """
        Initialize ResourcesTracker.
        """
        pass

    @log_errors(default_return=(0, 0), raise_exception=False)
    def get_container_cpu_and_memory(self, container: docker.models.containers.Container) -> Tuple[float, float]:
        """
        Get CPU and memory usage for a container.

        Args:
            container (docker.models.containers.Container): Docker container instance.

        Returns:
            Tuple[float, float]: CPU utilization percentage and memory utilization percentage.
        """
        stats = container.stats(stream=False)
        if stats:
            cpu_utilization = 0
            cpu_delta = (
                stats["cpu_stats"]["cpu_usage"]["total_usage"]
                - stats["precpu_stats"]["cpu_usage"]["total_usage"]
            )
            system_delta = stats["cpu_stats"].get("system_cpu_usage", 0) - stats[
                "precpu_stats"
            ].get("system_cpu_usage", 0)
            if system_delta > 0:
                cpu_utilization = cpu_delta / system_delta * 100.0
            memory_usage = stats["memory_stats"].get("usage", 0)
            memory_limit = stats["memory_stats"].get("limit", 1)
            memory_utilization = memory_usage / memory_limit * 100.0
            return cpu_utilization, memory_utilization
        return 0, 0

    @log_errors(default_return=(0, 0), raise_exception=False)
    def get_container_cpu_and_memory_with_container_id(self, container_id: str) -> Tuple[float, float]:
        """
        Get CPU and memory usage for a specific container by its ID.

        Args:
            container_id (str): ID of the Docker container.

        Returns:
            Tuple[float, float]: CPU utilization percentage and memory usage in MB.
        """
        stats_result = subprocess.run(
            [
                "docker",
                "stats",
                "--no-stream",
                "--format",
                "{{.ID}}: {{.CPUPerc}} CPU, {{.MemUsage}} RAM",
                container_id,
            ],
            capture_output=True,
            text=True,
            check=True,
        )
        stats = stats_result.stdout.strip().split(": ")[1].split(", ")
        cpu_usage = float(stats[0].replace("% CPU", "").strip())
        memory_usage = stats[1].split(" / ")[0]
        mem_value, mem_unit = memory_usage[:-3], memory_usage[-3:]
        if mem_unit == "KiB":
            memory_usage_mb = float(mem_value) / 1024
        elif mem_unit == "MiB":
            memory_usage_mb = float(mem_value)
        elif mem_unit == "GiB":
            memory_usage_mb = float(mem_value) * 1024
        else:
            memory_usage_mb = float(mem_value)
        return cpu_usage, memory_usage_mb

    @log_errors(default_return=(0, 0), raise_exception=False, log_error=False)
    def get_container_gpu_info(self, container_id: str) -> Tuple[float, int]:
        """
        Get GPU usage for a specific container.

        Args:
            container_id (str): ID of the Docker container.

        Returns:
            Tuple[float, int]: GPU utilization percentage and GPU memory usage in MB.
        """
        container_pid = self.get_pid_id_by_container_id(container_id)
        gpu_util = self.get_container_gpu_usage(container_pid)
        gpu_mem_used = self.get_container_gpu_memory_usage(container_pid)
        return gpu_util, gpu_mem_used

    @log_errors(default_return="", raise_exception=False)
    def get_pid_id_by_container_id(self, container_id: str) -> str:
        """
        Get PID for a container ID.

        Args:
            container_id (str): ID of the Docker container.

        Returns:
            str: PID of the container.
        """
        pid_result = subprocess.run(
            [
                "docker",
                "inspect",
                "--format",
                "{{.State.Pid}}",
                container_id,
            ],
            capture_output=True,
            text=True,
            check=True,
        )
        container_pid = pid_result.stdout.strip()
        return container_pid

    @log_errors(default_return=0, raise_exception=False, log_error=False)
    def get_container_gpu_usage(self, container_pid: str) -> float:
        """
        Get GPU usage for a container PID.

        Args:
            container_pid (str): PID of the Docker container.

        Returns:
            float: GPU utilization percentage.
        """
        if not has_gpu():
            return 0
        gpu_util = 0
        try:
            result = subprocess.run(
                ["nvidia-smi", "pmon", "-c", "1"],
                capture_output=True,
                text=True,
                check=True,
                timeout=5,
            )
            pmon_output = result.stdout.strip().split("\n")
            for line in pmon_output[2:]:
                parts = line.split()
                if len(parts) >= 8:
                    pid = parts[1]
                    gpu_usage = parts[3]
                    if pid == str(container_pid):
                        gpu_util += float(gpu_usage) if gpu_usage != "-" else 0
        except subprocess.TimeoutExpired:
            logging.warning("nvidia-smi pmon command timed out after 5 seconds in get_container_gpu_usage")
            return 0
        return gpu_util

    @log_errors(default_return=0, raise_exception=False, log_error=False)
    def get_container_gpu_memory_usage(self, container_pid: str) -> int:
        """
        Get GPU memory usage for a container PID.

        Args:
            container_pid (str): PID of the Docker container.

        Returns:
            int: GPU memory usage in MB.
        """
        if not has_gpu():
            return 0
        cmd = [
            "nvidia-smi",
            "--query-compute-apps=pid,used_memory",
            "--format=csv,noheader,nounits",
        ]
        total_memory = 0
        try:
            result = subprocess.run(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                check=True,
                timeout=5,
            )
            for line in result.stdout.splitlines():
                parts = line.strip().split(", ")
                if len(parts) == 2:
                    process_pid, used_memory = parts
                    if process_pid == str(container_pid):
                        total_memory += int(used_memory)
        except subprocess.TimeoutExpired:
            logging.warning("nvidia-smi command timed out after 5 seconds in get_container_gpu_memory_usage")
            return 0
        return total_memory

    @log_errors(default_return=(0, 0, 0, 0), raise_exception=False, log_error=True)
    def get_available_resources(self) -> Tuple[float, float, int, float]:
        """
        Get available machine resources.

        Returns:
            Tuple[float, float, int, float]: Available memory in GB, available CPU percentage,
            free GPU memory in MB, and GPU utilization percentage.
        """
        available_memory = psutil.virtual_memory().available / 1024**3
        available_cpu = 100 - psutil.cpu_percent(1)
        gpu_memory_free, gpu_utilization = self._get_gpu_resources()
        return available_memory, available_cpu, gpu_memory_free, gpu_utilization

    @log_errors(default_return=(0, 0.0), raise_exception=False, log_error=False)
    def _get_gpu_resources(self) -> Tuple[int, float]:
        """
        Get available GPU resources.

        Returns:
            Tuple[int, float]: Free GPU memory in MB and GPU utilization percentage.
        """
        gpu_memory_free = 0
        gpu_utilization = 0.0
        if not has_gpu():
            return gpu_memory_free, gpu_utilization

        try:
            subprocess.check_output("nvidia-smi", timeout=5)
        except subprocess.TimeoutExpired:
            logging.warning("nvidia-smi command timed out after 5 seconds in _get_gpu_resources")
            return 0, 0.0
        
        info_list = get_gpu_info()
        for info in info_list:
            info_split = info.split(", ")
            gpu_memory_free += int(info_split[5])
            gpu_utilization += float(info_split[2])
        gpu_utilization /= len(info_list) if info_list else 1

        return gpu_memory_free, gpu_utilization


class ActionsResourcesTracker:
    """Tracks Docker container action resources"""

    def __init__(self, scaling: Scaling):
        """Initialize ActionsResourcesTracker"""
        self.scaling = scaling
        self.max_actions_usage = {}
        self.resources_tracker = ResourcesTracker()
        self.client = docker.from_env()
        self.logged_stopped_containers = []

    @log_errors(raise_exception=False, log_error=True)
    def update_actions_resources(self) -> None:
        """Process both running and exited containers.
        
        Note: Does not remove containers to keep logs. Only tracks resource usage.
        """
        exited_containers = self.client.containers.list(
            filters={"status": "exited"},
            all=True,
        )
        running_containers = self.client.containers.list(filters={"status": "running"})
        if exited_containers:
            for container in exited_containers:
                try:
                    if container.id in self.logged_stopped_containers:
                        continue
                    self._update_container_action_status(container, "completed")
                    self.logged_stopped_containers.append(container.id)
                    # COMMENTED OUT: Do not remove containers to keep logs
                    # container.remove()
                except Exception as err:
                    logging.error(
                        "Error processing exited container %s: %s",
                        container.id,
                        str(err),
                    )
        if running_containers:
            for container in running_containers:
                try:
                    self._update_container_action_status(container, "running")
                except Exception as err:
                    logging.error(
                        "Error processing running container %s: %s",
                        container.id,
                        str(err),
                    )

    @log_errors(default_return=[], raise_exception=False)
    def get_sub_containers_by_label(self, label_key: str, label_value: str) -> list:
        """Get running containers with specified label key and value"""
        containers = self.client.containers.list(
            filters={
                "label": [f"{label_key}={label_value}"],
                "status": "running",
            }
        )
        return containers

    @log_errors(raise_exception=False, log_error=True)
    def _update_container_action_status(self, container, status: str) -> None:
        """Update action status for a specific container"""
        inspect_data = self.client.api.inspect_container(container.id)
        start_time = inspect_data["State"]["StartedAt"]
        finish_time = (
            inspect_data["State"]["FinishedAt"]
            if status == "completed"
            else datetime.now(timezone.utc).isoformat()
        )

        def remove_quotation_marks(args):
            """Remove quotes from container args"""
            new_args = []
            for arg in args:
                new_args.extend(x.replace('"', "").replace("'", "") for x in arg.split(" "))
            return new_args

        args_24 = [arg for arg in remove_quotation_marks(inspect_data["Args"]) if len(arg) == 24 and "pypi" not in arg]
        action_record_id = args_24[-1] if args_24 else None
        if not action_record_id:
            logging.debug("No valid action_id found for the container. Container ID: %s, Args: %s", container.id, inspect_data["Args"])
        duration = calculate_time_difference(start_time, finish_time)
        (
            current_gpu_utilization,
            current_gpu_memory,
            current_cpu_utilization,
            current_memory_utilization,
        ) = self.get_current_action_usage(container, status)
        sub_containers = self.get_sub_containers_by_label("action_id", action_record_id)
        for sub_container in sub_containers:
            if sub_container.id in self.logged_stopped_containers:
                continue
            (
                sub_container_gpu_utilization,
                sub_container_gpu_memory,
                sub_container_cpu_utilization,
                sub_container_memory_utilization,
            ) = self.get_current_action_usage(sub_container, status)
            current_gpu_utilization += sub_container_gpu_utilization
            current_gpu_memory += sub_container_gpu_memory
            current_cpu_utilization += sub_container_cpu_utilization
            current_memory_utilization += sub_container_memory_utilization
            # COMMENTED OUT: Do not stop/remove sub-containers to keep logs
            if status == "completed":
                try:
                    sub_container.stop()
                    self.logged_stopped_containers.append(sub_container.id)
            #         sub_container.remove(force=True)
                except Exception as err:
                    logging.error(
                        "Error removing sub-container %s: %s",
                        sub_container.id,
                        str(err),
                    )
        (
            max_gpu_utilization,
            max_gpu_memory,
            max_cpu_utilization,
            max_memory_utilization,
        ) = self.update_max_action_usage(
            action_record_id,
            current_gpu_utilization,
            current_gpu_memory,
            current_cpu_utilization,
            current_memory_utilization,
        )
        logging.info(
            "Updating action status: service_provider=%s, action_id=%s, running=%s, status=%s, duration=%s, start=%s, gpu_util=%.2f%%, cpu_util=%.2f%%, gpu_mem=%dMB, mem_util=%.2f%%, created=%s, updated=%s",
            os.environ["SERVICE_PROVIDER"],
            action_record_id,
            status == "running",
            status,
            duration,
            start_time,
            max_gpu_utilization,
            max_cpu_utilization,
            max_gpu_memory,
            max_memory_utilization,
            start_time,
            finish_time,
        )
        self.scaling.update_action_status(
            service_provider=os.environ["SERVICE_PROVIDER"],
            action_record_id=action_record_id,
            isRunning=status == "running",
            status=status,
            action_duration=duration,
            docker_start_time=start_time,
            gpuUtilisation=max_gpu_utilization,
            cpuUtilisation=max_cpu_utilization,
            gpuMemoryUsed=max_gpu_memory,
            memoryUtilisation=max_memory_utilization,
            createdAt=start_time,
            updatedAt=finish_time,
        )

    @log_errors(default_return=(0, 0, 0, 0), raise_exception=False)
    def get_current_action_usage(self, container, status: str) -> Tuple[float, int, float, float]:
        """Get current resource usage for a container"""
        current_gpu_utilization = 0
        current_gpu_memory = 0
        current_cpu_utilization = 0
        current_memory_utilization = 0
        if status == "running":
            try:
                (
                    current_cpu_utilization,
                    current_memory_utilization,
                ) = self.resources_tracker.get_container_cpu_and_memory(container)
                (
                    current_gpu_utilization,
                    current_gpu_memory,
                ) = self.resources_tracker.get_container_gpu_info(container_id=container.id)
            except Exception as err:
                logging.error(
                    "Error getting container usage metrics: %s",
                    str(err),
                )
        return (
            current_gpu_utilization,
            current_gpu_memory,
            current_cpu_utilization,
            current_memory_utilization,
        )

    @log_errors(default_return=(0, 0, 0, 0), raise_exception=False, log_error=True)
    def update_max_action_usage(
        self,
        action_record_id: str,
        current_gpu_utilization: float,
        current_gpu_memory: int,
        current_cpu_utilization: float,
        current_memory_utilization: float,
    ) -> Tuple[float, int, float, float]:
        
        """Update and return maximum resource usage values for an action"""
        if action_record_id not in self.max_actions_usage:
            self.max_actions_usage[action_record_id] = {
                "gpu_utilization": 0,
                "gpu_memory": 0,
                "cpu_utilization": 0,
                "memory_utilization": 0,
            }
        current_values = {
            "gpu_utilization": current_gpu_utilization or 0,
            "gpu_memory": current_gpu_memory or 0,
            "cpu_utilization": current_cpu_utilization or 0,
            "memory_utilization": current_memory_utilization or 0,
        }
        for key in current_values:
            self.max_actions_usage[action_record_id][key] = max(
                current_values[key],
                self.max_actions_usage[action_record_id][key],
            )
        return (
            self.max_actions_usage[action_record_id]["gpu_utilization"],
            self.max_actions_usage[action_record_id]["gpu_memory"],
            self.max_actions_usage[action_record_id]["cpu_utilization"],
            self.max_actions_usage[action_record_id]["memory_utilization"],
        )


class MachineResourcesTracker:
    """Tracks machine-level resources like CPU, memory and GPU"""

    def __init__(self, scaling: Scaling):
        """Initialize MachineResourcesTracker"""
        self.scaling = scaling
        self.resources_tracker = ResourcesTracker()

    @log_errors(raise_exception=False, log_error=True)
    def update_available_resources(self):
        """Update available machine resources"""
        (
            available_memory,
            available_cpu,
            gpu_memory_free,
            gpu_utilization,
        ) = self.resources_tracker.get_available_resources()
        _, err, _ = self.scaling.update_available_resources(
            availableCPU=available_cpu,
            availableMemory=available_memory,
            availableGPU=100 - gpu_utilization,
            availableGPUMemory=gpu_memory_free,
        )
        if err is not None:
            logging.error(
                "Error in updating available resources: %s",
                err,
            )
