from math import ceil


class Renderer:
    def __init__(
        self, 
        progress_bar_width=50, 
        columns=1, 
        use_space_lines=True, 
        node_names=None, 
        display_power=True,
        display_users=False,
    ) -> None:
        self.progress_bar_width = progress_bar_width
        self.columns = columns
        self.use_space_lines = use_space_lines
        self.node_names = node_names
        self.display_power = display_power
        self.display_users = display_users

    def render_info_dict(self, info_dict):
        line_blocks = []

        for node_dict in info_dict:
            if self.node_names is None or node_dict["name"] in self.node_names:
                lines = self.render_node(node_dict)

                line_blocks.append(lines)

        first_line = "|" + "=" * (len(lines[-1]) - 2) + "|"
        
        final_lines = []

        n_rows = ceil(len(line_blocks) / self.columns)

        # Format rows and columns
        for row in range(n_rows):
            lines = []

            max_size = -1

            # Find max size of the blocks in the same row:
            for col in range(self.columns):
                if col * n_rows + row < len(line_blocks):
                    if len(line_blocks[col * n_rows + row]) > max_size:
                        max_size = len(line_blocks[col * n_rows + row])

            for col in range(self.columns):
                if col * n_rows + row < len(line_blocks):
                    if len(lines) == 0:
                        lines.extend(
                            self.expand_rendered_block_to_size(
                                line_blocks[col * n_rows + row], 
                                max_size
                            )
                        )
                    else:
                        for i, line in enumerate(self.expand_rendered_block_to_size(line_blocks[col * n_rows + row], max_size)):
                            lines[i] += line
            
            final_lines.extend(lines)            

        final_lines.insert(0, first_line * self.columns)

        #lines.append("=" * len(lines[-1]))

        return "\n".join(final_lines)

    def render_node(self, node_dict):
        name = node_dict["name"]
        mem_used = node_dict["latest_info"]["used_memory_mb"]
        mem_total = node_dict["total_memory_mb"]
        utilization = node_dict["latest_info"]["cpu_utilization"]
        temp = node_dict["latest_info"]["temperature"]

        head_line = "|- Node: " + name + " "
        info_line = f"| CPU: {utilization:>4.1f}%   Memory: {mem_used:>6}/{mem_total:<6} MB   Temp: {temp:>3}°C"

        lines = []

        for i, gpu_dict in enumerate(node_dict["gpus"]):
            lines.extend(self.get_rendered_gpu_lines(gpu_dict))

        head_line = head_line + "-" * (len(lines[-1]) - len(head_line) - 1) + "|"
        info_line = info_line + " " * (len(lines[-1]) - len(info_line) - 1) + "|"
        pad_line = "|" + "-" * (len(lines[-1]) - 2) + "|"
        pad_line_empty = "|" + " " * (len(lines[-1]) - 2) + "|"

        lines.append(pad_line_empty)
        lines.append("|" + "=" * (len(lines[-1]) - 2) + "|")

        lines.insert(0, pad_line)
        if self.use_space_lines:
            lines.insert(0, pad_line_empty)
        lines.insert(0, info_line)
        if self.use_space_lines:
            lines.insert(0, pad_line_empty)
        lines.insert(0, head_line)

        return lines

    def get_rendered_gpu_lines(self, gpu_dict):
        gpu_type = gpu_dict["type"]
        index = gpu_dict["index"]
        mem_used = gpu_dict["latest_info"]["used_memory_mb"]
        mem_total = gpu_dict["total_memory_mb"]
        utilization = gpu_dict["latest_info"]["utilization"]
        temp = gpu_dict["latest_info"]["temperature"]
        n_processes = len(gpu_dict["running_processes"])
        power = gpu_dict["latest_info"]["power_draw"]

        mem_used_percent = int(self.progress_bar_width * mem_used / mem_total)
        rest_mem = self.progress_bar_width - mem_used_percent

        line_util = "| [" + \
            "=" * mem_used_percent + \
            " " * rest_mem + "]" + \
            f"{mem_used:>6}/{mem_total:<6} MB,  Util: {int(utilization):>3}% |"

        line_meta = f"| GPU #{index} ({gpu_type}):   #Proc.: {n_processes}   Temp: {temp:>3}°C   " + (f"Pow: {int(power):>3} W" if self.display_power else "")

        line_meta = line_meta + " " * (len(line_util) - len(line_meta) - 1) + "|"
        empty_line = "|" + " " * (len(line_meta) - 2) + "|"

        lines = [line_meta, line_util]

        if self.display_users:
            lines.append(self.get_rendered_users_line(gpu_dict["running_processes"], len(line_meta)))

        if self.use_space_lines:
            lines = [empty_line] + lines

        return lines

    def get_rendered_users_line(self, processes_list, line_len):
        user_names = [p_dict["user_name"] for p_dict in processes_list]
        
        line = "| Users: " + ", ".join(user_names)
        line += " " * (line_len - len(line) - 1) + "|"

        return line

    def expand_rendered_block_to_size(self, block_lines, max_size):
        buffer_line = "|" + " " * (len(block_lines[-1]) - 2) + "|"

        block_lines_new = block_lines[:-1] + [buffer_line] * (max_size - len(block_lines)) + block_lines[-1:]

        return block_lines_new