"""
Shows diagrams from reports generated by Waterfalls.
"""
from __future__ import annotations

import argparse
from collections import defaultdict, namedtuple
from glob import glob
import json
import logging
import os
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import matplotlib.pyplot as plt


TimeUnit = namedtuple("TimeUnit", ["name", "si_symbol", "user_symbol", "ns_multiple"])

TIME_UNITS = [
    TimeUnit("nanoseconds", "ns", "nsec", 1),
    TimeUnit("microseconds", "µs", "usec", 1e3),
    TimeUnit("milliseconds", "ms", "msec", 1e6),
    TimeUnit("seconds", "s", "sec", 1e9),
    TimeUnit("minutes", "m", "min", 1e9 * 60),
    TimeUnit("hours", "h", "hour", 1e9 * 3600),
]
COLOR_LIGHT = "#e5ecf6"
COLOR_DARK = "#20293d"
COLOR_BLUE = "#64c9ec"
COLOR_BORDER = "#44b0d0"

BlocksType = List[dict]
TimersType = Dict[str, BlocksType]


logger = logging.getLogger(__name__)


class Viewer:
    """
    Class for visualizing reports generated by Waterfalls.
    """

    def __init__(
        self,
        directory: str = os.getcwd(),
        user_time_unit: Optional[str] = None,
        show_thread_id: bool = False,
        show_horizontal_lines: bool = False,
        save_image: bool = False,
    ) -> None:
        """
        Constructs an instance of waterfalls.Viewer.

        Args:
            directory: Directory containing report file(s) generated by `waterfalls.Timer`.
            user_time_unit: Time unit defined as `TimeUnit.user_symbol` (e.g., `msec`).
                When `None`, then the most appropriate time unit is determined automatically.
            show_thread_id: When `True`, each timer will show its thread ID.
                When `False`, only timers with overlapping blocks will be differentiated by thread ID.
            show_horizontal_lines: When `True`, then horizontal lines will be rendered in the diagram.
            save_image: When `True`, then diagram will be saved as `waterfalls.svg` image
                in the reports `directory`.
        """
        self.directory = directory
        self.user_time_unit = user_time_unit
        self.show_thread_id = show_thread_id
        self.show_horizontal_lines = show_horizontal_lines
        self.save_image = save_image
        self.directory_path = Path(self.directory)

    def visualize_report(self) -> None:
        """
        Loads Waterfalls report(s) and shows them on the screen or saves them into an image file.
        """
        report_file_paths = self._get_report_file_paths()
        blocks = self._load_blocks_from_reports(report_file_paths)
        timers, time_total, time_min = self._group_blocks_to_timers(blocks)
        timers = self._format_timer_names(timers)
        timers = self._sort_timers(timers)
        time_unit = self._determine_time_unit(time_total)
        self._render_diagram(timers, time_unit, time_total, time_min)

    @staticmethod
    def _add_data_to_diagram(ax: plt.axes.AxesSubplot, data: TimersType, time_unit: TimeUnit, time_min: int) -> None:
        """
        Adds all timers and timing blocks into the diagram.

        Args:
            ax: Matplotlib subplot.
            data: Dictionary mapping timers and a list of all their blocks.
            time_unit: Time unit that should be used for the time axis and all durations.
            time_min: Time of the first block of all timers that will be used as the beginning of the time axis.
        """
        timer_names = []

        # Place blocks on the diagram
        for i, (timer_name, blocks) in enumerate(data.items()):
            timer_names.append(timer_name)
            block_times = []
            thread_durations = []

            for block in blocks:
                start_time = (block["start_time"] - time_min) / time_unit.ns_multiple
                duration = (block["stop_time"] - block["start_time"]) / time_unit.ns_multiple
                cpu_percent = min(block["thread_duration"] / (block["stop_time"] - block["start_time"]) * 100, 100)
                block_times.append([start_time, duration])
                thread_durations.append([start_time, block["thread_duration"] / time_unit.ns_multiple])
                ax.text(
                    x=start_time,
                    y=i * 10 + 5,
                    s=f"time: {duration:.3f} {time_unit.si_symbol}\ncpu: {cpu_percent:.1f} %",
                    color=COLOR_DARK,
                    verticalalignment="bottom",
                )
                ax.text(x=start_time, y=i * 10 + 5, s=block["text"], color=COLOR_DARK, verticalalignment="top")

            # Block duration
            ax.broken_barh(block_times, (i * 10 + 1, 8), facecolors=COLOR_BLUE, edgecolor=COLOR_BORDER)

            # Thread duration
            ax.broken_barh(thread_durations, (i * 10 + 1, 8), facecolors=COLOR_BORDER)

        # Place timer names on the diagram
        y_tick_placement = [i * 10 + 5 for i in range(len(timer_names))]
        ax.set_yticks(y_tick_placement, labels=timer_names)

    @staticmethod
    def _detect_overlap(blocks: BlocksType) -> bool:
        """
        Detects if any two blocks overlap.

        Args:
            blocks: List of timing blocks. Blocks must be sorted by `start_time` from lowest to highest.

        Returns:
            `True` if any two blocks overlap, `False` otherwise.
        """
        previous_stop_time = -1
        for block in blocks:
            if block["start_time"] < previous_stop_time:
                return True
            previous_stop_time = block["stop_time"]
        return False

    def _determine_time_unit(self, time_total: int) -> TimeUnit:
        """
        Determines which time unit should be used.

        Tries to find unit which results in the lowest value of the total time which is not lower than 1
        (it is more human friendly to read 30 seconds than 0.5 minute).

        Args:
            time_total: Number of nanoseconds between the start time of the first block and stop time of the last block.

        Returns:
            A `TimeUnit` object. When `self.user_time_unit` is not `None`, it has the highest priority.
                When it is `None`, the time unit will be automatically determined based on `total_time`.
                As the last resort, nanoseconds time unit is returned.
        """
        if self.user_time_unit is not None:
            for time_unit in TIME_UNITS:
                if self.user_time_unit == time_unit.user_symbol:
                    return time_unit

        for time_unit in reversed(TIME_UNITS):
            if time_total / time_unit.ns_multiple >= 1:
                return time_unit

        return TIME_UNITS[0]

    def _format_diagram(self, fig: plt.Figure, ax: plt.axes.AxesSubplot, time_unit: TimeUnit, time_total: int) -> None:
        """
        Args:
            fig: Matplotlib figure.
            ax: Matplotlib subplot.
            time_unit: Time unit that should be used for the time axis and all durations.
            time_total: Number of nanoseconds between the start time of the first block and stop time of the last block.
        """
        # Show time unit symbols on x axis
        ax.xaxis.set_major_formatter("{x:.0f} " + time_unit.si_symbol)

        ax.set_title("Generated by Waterfalls", color=COLOR_DARK)
        ax.set_xlabel(
            f"{time_unit.name} since start " f"(total: {time_total / time_unit.ns_multiple:.3f} {time_unit.si_symbol})"
        )

        # Hide tick marks (short lines extending from the main diagram area)
        ax.tick_params(bottom=False, left=False)

        # Control the appearance of grid lines
        ax.grid(axis="both" if self.show_horizontal_lines else "x", color="white")

        # Invert y axis to start waterfall from top to bottom
        ax.invert_yaxis()

        # Move grid lines below bars
        ax.set_axisbelow(True)

        # Hide border around the diagram
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)
        ax.spines["bottom"].set_visible(False)
        ax.spines["left"].set_visible(False)

        # Set color of axis labels
        ax.xaxis.label.set_color(COLOR_DARK)

        # Set color of tick text (timer names)
        ax.tick_params(axis="both", colors=COLOR_DARK)

        # Set color of the diagram main area
        ax.set_facecolor(COLOR_LIGHT)

        fig.set_tight_layout(True)

    def _format_timer_names(self, timers: TimersType) -> TimersType:
        """
        Formats timer names by including thread IDs.

        When `self.show_thread_id` is `True`, each timer will show its thread ID. When `False`,
            only timers with overlapping blocks will be differentiated by thread ID.

        Args:
            timers: Dictionary mapping timer names and a list of all their blocks.

        Returns:
            Dictionary mapping timers with formatted names and a list of all their blocks.
        """
        formatted_timers = defaultdict(list)

        for timer_name, blocks in timers.items():
            blocks.sort(key=lambda block: block["start_time"])
            if self.show_thread_id or self._detect_overlap(blocks):
                for block in blocks:
                    timer_and_thread_name = f"{block['name']}\nthread: {block['thread_id']}"
                    formatted_timers[timer_and_thread_name].append(block)
            else:
                formatted_timers[timer_name] = blocks

        return formatted_timers

    def _get_report_file_paths(self) -> List[str]:
        """
        Lists all Waterfalls report files in the reports directory.

        Returns:
            List of all report file paths in the reports directory.

        Raises:
            SystemExit: When no report file is found in the specified directory.
        """
        report_file_paths = glob(str(self.directory_path.joinpath("waterfalls*.json")))
        if not report_file_paths:
            raise SystemExit(
                f"No Waterfalls report file found in directory '{str(self.directory_path.resolve())}'. "
                "Type 'waterfalls --help' to see how to specify the directory"
            )
        return report_file_paths

    @staticmethod
    def _group_blocks_to_timers(blocks: BlocksType) -> Tuple[TimersType, int, int]:
        """
        Groups blocks by the name of the timer and calculates the time duration of the timing.

        Args:
            blocks: List of timing blocks.

        Returns:
            Dictionary mapping the timer name and a list of all its timing blocks,
            total time from the start of the first block to the stop of the last block,
            and star time of the first block.
        """
        timers = defaultdict(list)
        start_times = []
        stop_times = []

        for block in blocks:
            block_name = block["name"]
            timers[block_name].append(block)
            start_times.append(block["start_time"])
            stop_times.append(block["stop_time"])

        time_min = min(start_times)
        time_max = max(stop_times)
        time_total = time_max - time_min

        return timers, time_total, time_min

    @staticmethod
    def _load_blocks_from_reports(report_file_paths: List[str]) -> BlocksType:
        """
        Loads all timing blocks from all report files.

        Args:
            report_file_paths: List of all report file paths in reports directory.

        Returns:
            List of timing blocks.

        Raises:
            SystemExit: When no timing block found in any of the reports.
        """
        blocks = []

        for report_file_path in report_file_paths:
            with open(report_file_path) as report_file:
                blocks.extend(json.load(report_file))

        if not blocks:
            raise SystemExit("No timing block found in the report file(s)")

        return blocks

    def _render_diagram(self, data: TimersType, time_unit: TimeUnit, time_total: int, time_min: int) -> None:
        """
        Renders Matplotlib diagram.

        Args:
            data: Dictionary mapping timers and a list of all their blocks.
            time_unit: Time unit that should be used for the time axis and all durations.
            time_total: Number of nanoseconds between the start time of the first block and stop time of the last block.
            time_min: Time of the first block of all timers that will be used as the beginning of the time axis.
        """
        fig, ax = plt.subplots(figsize=(16, 9))
        self._add_data_to_diagram(ax, data, time_unit, time_min)
        self._format_diagram(fig, ax, time_unit, time_total)

        if self.save_image:
            image_path = self.directory_path.joinpath("waterfalls.svg")
            plt.savefig(image_path)
            logger.info("Diagram image saved into file '%s'", image_path.resolve())
        else:
            fig.canvas.manager.set_window_title("Waterfalls")
            plt.show()

    def _sort_timers(self, timers: TimersType) -> TimersType:
        """
        Sorts timers by the `start_time` of the first block.
        When thread ID showing is enabled, then timers should be primarily sorted by their `thread_id`
        and then by `start_time`.

        Args:
            timers: Dictionary mapping timer names and a list of all their blocks.

        Returns:
            Dictionary mapping sorted timers and a list of all their blocks.
        """
        timers = dict(sorted(timers.items(), key=lambda item: min(block["start_time"] for block in item[1])))

        if self.show_thread_id:
            timers = dict(sorted(timers.items(), key=lambda item: min(block["thread_id"] for block in item[1])))

        return timers


def _parse_arguments() -> argparse.Namespace:
    """
    Parses command line arguments.

    Returns:
        A `Namespace` object holding parsed argument values.
    """
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "directory",
        nargs="?",
        default=os.getcwd(),
        type=str,
        help="Directory containing report file(s) generated by waterfalls.Timer. The path should be quoted. "
        "By default, the current working directory is used.",
    )
    parser.add_argument(
        "-u",
        "--unit",
        choices=[time_unit.user_symbol for time_unit in TIME_UNITS],
        help="Specifies time unit. By default, the time unit is determined automatically.",
    )
    parser.add_argument(
        "-t",
        "--thread_id",
        action="store_true",
        help="Shows thread ID next to each timer. "
        "By default, thread ID is hidden and only shown to distinguish overlapping blocks belonging to the same timer.",
    )
    parser.add_argument("-l", "--lines", action="store_true", help="Shows horizontal lines in the graph.")
    parser.add_argument(
        "-i",
        "--image",
        action="store_true",
        help="Saves diagram as an image (waterfalls.svg) to the reports directory "
        "instead of launching interactive window.",
    )
    return parser.parse_args()


def main() -> None:
    """
    Parses arguments from the command line and visualizes report(s).
    """
    logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO)

    args = _parse_arguments()

    Viewer(
        directory=args.directory,
        user_time_unit=args.unit,
        show_thread_id=args.thread_id,
        show_horizontal_lines=args.lines,
        save_image=args.image,
    ).visualize_report()


if __name__ == "__main__":
    main()
