depgraph

#!/usr/bin/env python3
"""
Package Dependency Navigator - A tool to explore Gentoo package dependencies
with ncdu-style navigation.
"""

import pickle
import subprocess
import re
import sys
import os
import traceback
from pathlib import Path
from enum import IntEnum, auto
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple
from operator import attrgetter

import curses
import curses.textpad


class ViewType(IntEnum):
    EMPTY = auto()
    OVERLAYS = auto()
    EQUERY_LIST = auto()
    EQUERY_DEPGRAPH = auto()
    EQUERY_DEPENDS = auto()


@dataclass
class Package:
    cpv: str
    dependencies_count: int = 0
    dependencies: List["Dependency"] = field(default_factory=list)


@dataclass
class Dependency:
    cpv: str
    depth: int
    filtered_depend_count: int = 1


@dataclass
class State:
    view_type: ViewType
    selected_index: int = 0
    scroll_offset: int = 0


class Parser:
    EQUERY_LIST_SEPARATOR = "|"
    EQUERY_LIST_PLACEHOLDERS = ["$cpv", "$repo"]
    EQUERY_LIST_SET_PLACEHOLDERS = ["$cp"]

    def __init__(self, atom: str):
        self.atom = atom
        self.equery_list_placeholders_count = len(self.EQUERY_LIST_PLACEHOLDERS)
        self.equery_list_set_placeholders_count = len(self.EQUERY_LIST_SET_PLACEHOLDERS)
        self.equery_list_tmpl = self.EQUERY_LIST_SEPARATOR.join(
            self.EQUERY_LIST_PLACEHOLDERS
        )
        self.equery_list_set_tmpl = self.EQUERY_LIST_SEPARATOR.join(
            self.EQUERY_LIST_SET_PLACEHOLDERS
        )

        self.dependencies = {}
        self.world_packages: Dict[str, bool] = {}
        self.installed_packages: set[str] = set()

    def parse_equery_list(self, output: str) -> Dict[str, List[Package]]:
        """Parse equery list output and return overlays with package dictionaries."""
        overlays = {}

        for line in output.strip().split("\n"):
            if not line.strip():
                continue

            parts = line.split(self.EQUERY_LIST_SEPARATOR)
            if len(parts) != self.equery_list_placeholders_count:
                print(f"Equery list placeholders count mismatch")
                print("Expected: ", self.EQUERY_LIST_PLACEHOLDERS)
                print("Got: ", line)
                sys.exit(1)

            cpv, repo = parts

            pkg = Package(cpv=cpv, dependencies_count=0, dependencies=[])

            if repo not in overlays:
                overlays[repo] = []
            overlays[repo].append(pkg)

        return overlays

    def parse_equery_list_to_set(self, output: str) -> set[str]:
        packages = set()

        for line in output.strip().split("\n"):
            if not line.strip():
                continue

            parts = line.split(self.EQUERY_LIST_SEPARATOR)
            if len(parts) != self.equery_list_set_placeholders_count:
                print(f"Equery list placeholders count mismatch")
                print("Expected: ", self.EQUERY_LIST_SET_PLACEHOLDERS)
                print("Got: ", line)
                sys.exit(1)

            cp = parts[0]

            packages.add(cp)

        return packages

    def parse_dependency_graph(self, output: str) -> List[Dependency]:
        """Parse equery g output and return list of dependencies."""
        lines = output.strip().split("\n")

        deps = []

        # Skip the header line
        for line in lines[1:]:
            if not line.strip():
                continue

            match = re.search(r"\s+\[\s*(\d+)\]\s+([\w\-/]+)-([\-\d\.\w_]+)", line)
            if match:
                depth, cp, fullverion = match.groups()

                # skip self
                if int(depth) == 0:
                    continue

                if cp not in self.installed_packages:
                    continue

                cpv = f"{cp}-{fullverion}"

                dep = Dependency(
                    cpv=cpv,
                    depth=depth,
                )
                deps.append(dep)

                if cpv in self.dependencies:
                    self.dependencies[cpv].filtered_depend_count += 1
                else:
                    self.dependencies[cpv] = dep

        return deps

    def parse_dependents(self, output: str) -> List[str]:
        """Parse equery d output and return list of packages that depend on this package."""
        dependents = []
        for line in output.strip().split("\n"):
            if line.startswith(" * "):
                continue
            match = re.search(r"([^/]+/[^-]+)-([^\s]+)", line)
            if match:
                dependents.append(match.group(1))
        return dependents


class PackageNavigator:
    TITLE_LINE = 0
    HEADER_LINE = 1
    HEADER_SEPARATOR_LINE = 2
    CONTENT_LINE = 3

    # 0 is default
    TITLE_COLOR_PAIR = 1
    HEADER_COLOR_PAIR = 2
    STATUS_BAR_COLOR_PAIR = 3
    SELECTED_COLOR_PAIR = 4
    CONTENT_COLOR_PAIR = 5

    def __init__(self, atom: str, overlays: Dict[str, List[Package]]):
        self.atom = atom
        self.overlays: Dict[str, List[Package]] = {
            name: (sorted(overlay, key=attrgetter("dependencies_count"), reverse=True))
            for name, overlay in overlays.items()
        }
        self.overlay_keys = list(self.overlays.keys())
        self.overlays_length = len(self.overlay_keys)

        self.stack = [State(view_type=ViewType.OVERLAYS)]
        self._breadcrumbs = [self.atom]
        self.overlay: str = self.overlay_keys[0]
        self.package: Optional[Package] = None

        self.height_offset = 1
        self.prev_view_type = ViewType.EMPTY
        self.prev_scroll_offset = 0
        self.items_length = 0

    @property
    def max_content_y(self):
        # minus navbar
        return self.height - self.height_offset - 1

    @property
    def state(self):
        return self.stack[-1]

    def breadcrubms(self, segment: str):
        self._breadcrumbs.append(segment)
        self.title(">".join(self._breadcrumbs))

    def title(self, line):
        line = line.ljust(self.width)[: self.width]
        self.stdscr.addstr(
            self.TITLE_LINE,
            0,
            line,
            curses.color_pair(self.TITLE_COLOR_PAIR) | curses.A_REVERSE,
        )

    def header(self, line):
        line = line.ljust(self.width)
        self.stdscr.addstr(
            self.HEADER_LINE, 0, line, curses.color_pair(self.HEADER_COLOR_PAIR)
        )
        self.stdscr.addstr(
            self.HEADER_SEPARATOR_LINE,
            0,
            "-" * self.width,
            curses.color_pair(self.HEADER_COLOR_PAIR),
        )

    def line(self, y, line, attrs_mask=None):
        if attrs_mask is None:
            attrs_mask = curses.color_pair(self.CONTENT_COLOR_PAIR)

        line = line.ljust(self.width)
        try:
            self.stdscr.addstr(y, 0, line, attrs_mask)
        except curses.error as e:
            max_y = self.height - self.height_offset
            raise Exception(
                f"Cannot draw a line of length {len(line)}/{self.width - 1} on y {y}/{max_y} scroll_offset {self.state.scroll_offset}: '{line}'"
            ) from e

    def selected_line(self, pos, line):
        self.line(
            pos, line, curses.A_BOLD | curses.color_pair(self.SELECTED_COLOR_PAIR)
        )

    def draw_overlays_view(self, first_render: bool, scrolled: bool):
        if first_render:
            self.breadcrubms("Overlays")

        self.draw_table(
            first_render,
            scrolled,
            [60, 20],
            ["Repository", "Packages"],
            [[repo, len(self.overlays[repo])] for repo in self.overlay_keys],
        )

    def draw_equery_list_view(self, first_render: bool, scrolled: bool):
        if first_render:
            self.breadcrubms(f"{self.overlay}>Equery list")

        self.draw_table(
            first_render,
            scrolled,
            [60, 20],
            ["Package", "Installed dependencies"],
            [[p.cpv, p.dependencies_count] for p in self.overlays[self.overlay]],
        )

    def draw_equery_depgraph_view(self, first_render: bool, scrolled: bool):
        if first_render:
            self.breadcrubms(f"{self.package.cpv}>Depgraph")

        self.draw_table(
            first_render,
            scrolled,
            [60, 6, 20],
            ["Package", "Depth", f"Dependants in {self.atom}"],
            [
                [p.cpv, p.depth, p.filtered_depend_count]
                for p in self.package.dependencies
            ],
        )

    def draw_table(
        self,
        first_render: bool,
        scrolled: bool,
        limits: List,
        headers: List,
        items: List,
    ):
        if sum(limits) > 100:
            raise Exception("Limits sum cannot exceed 100%")

        limits = [int(self.width * limit / 100) for limit in limits]

        def _limit(elements: List):
            elements = [str(elements[i]).ljust(limit) for i, limit in enumerate(limits)]
            return "".join(elements)

        items_length = len(items)
        if first_render:
            self.items_length = items_length
            self.header(_limit(headers))

        state = self.state
        max_content_y = self.max_content_y

        if state.selected_index >= items_length:
            state.selected_index = items_length - 1

        if first_render or scrolled:
            for y in range(self.CONTENT_LINE, self.max_content_y + 1):
                i = y + state.scroll_offset - self.CONTENT_LINE
                if i < items_length:
                    if i == state.selected_index:
                        self.selected_line(y, _limit(items[i]))
                    else:
                        self.line(y, _limit(items[i]))
                else:
                    self.line(y, "")
        else:
            i = state.selected_index
            y = self.CONTENT_LINE + i - state.scroll_offset
            self.selected_line(y, _limit(items[i]))
            if i > 0 and y > self.CONTENT_LINE:
                self.line(y - 1, _limit(items[i - 1]))
            if i < items_length - 1 and y + 1 < max_content_y:
                self.line(y + 1, _limit(items[i + 1]))

    def render(self):
        self.height, self.width = self.stdscr.getmaxyx()

        state = self.state
        first_render = state.view_type != self.prev_view_type
        scrolled = state.scroll_offset != self.prev_scroll_offset
        params = [first_render, scrolled]
        if state.view_type == ViewType.OVERLAYS:
            self.draw_overlays_view(*params)
        elif state.view_type == ViewType.EQUERY_LIST:
            self.draw_equery_list_view(*params)
        elif state.view_type == ViewType.EQUERY_DEPGRAPH:
            self.draw_equery_depgraph_view(*params)

        status_line = f"{state.selected_index + 1}/{self.items_length} ^{state.scroll_offset} | h:Back/Up  j/k:Move  l:Enter  q:Quit"
        self.stdscr.addstr(
            self.height - self.height_offset,
            0,
            status_line.ljust(self.width - 1),
            curses.color_pair(self.STATUS_BAR_COLOR_PAIR) | curses.A_REVERSE,
        )

        self.stdscr.refresh()

        self.prev_view_type = state.view_type
        self.prev_scroll_offset = state.scroll_offset

    def handle_input(self):
        key = self.stdscr.getch()

        state = self.state
        if key == ord("q"):
            return False
        elif key == ord("j") or key == curses.KEY_DOWN:
            state.selected_index += 1
            if (
                state.selected_index - state.scroll_offset
                >= self.max_content_y - self.CONTENT_LINE
            ):
                state.scroll_offset += 1
            return True
        elif key == ord("k") or key == curses.KEY_UP:
            if state.selected_index > 0:
                state.selected_index -= 1
                if state.selected_index < state.scroll_offset:
                    state.scroll_offset -= 1
            return True
        elif key == ord("l") or key == ord("\n") or key == curses.KEY_RIGHT:
            view_type = None
            if state.view_type == ViewType.OVERLAYS:
                self.overlay = self.overlay_keys[state.selected_index]
                if len(self.overlays[self.overlay]):
                    view_type = ViewType.EQUERY_LIST
            elif state.view_type == ViewType.EQUERY_LIST:
                self.package = self.overlays[self.overlay][state.selected_index]
                if len(self.package.dependencies):
                    view_type = ViewType.EQUERY_DEPGRAPH
            if view_type:
                self.stack.append(State(view_type=view_type))
                return True
        elif key == ord("h") or key == ord("\x1b") or key == curses.KEY_LEFT:
            if len(self.stack) > 1:
                self.stack.pop()
                # to refill
                self._breadcrumbs = self._breadcrumbs[:-2]
                return True
        elif key == ord("r"):
            return True
        elif key == -1:
            # No key pressed, continue
            pass

    def main_loop(self, stdscr):
        """Main application loop."""
        self.stdscr = stdscr

        if curses.has_colors():
            curses.start_color()
            curses.init_pair(
                self.TITLE_COLOR_PAIR, curses.COLOR_GREEN, curses.COLOR_WHITE
            )
            curses.init_pair(
                self.HEADER_COLOR_PAIR, curses.COLOR_BLUE, curses.COLOR_WHITE
            )
            curses.init_pair(
                self.STATUS_BAR_COLOR_PAIR, curses.COLOR_YELLOW, curses.COLOR_WHITE
            )
            curses.init_pair(
                self.SELECTED_COLOR_PAIR, curses.COLOR_GREEN, curses.COLOR_WHITE
            )
            curses.init_pair(
                self.CONTENT_COLOR_PAIR, curses.COLOR_BLACK, curses.COLOR_WHITE
            )

        curses.curs_set(0)  # Hide cursor
        stdscr.keypad(True)
        # Blocks until keypress. 0% CPU
        stdscr.timeout(-1)

        self.height, self.width = self.stdscr.getmaxyx()
        failed = True
        while failed and self.height_offset < self.height:
            try:
                self.stdscr.addstr(
                    self.height - self.height_offset, 0, "Bottom", curses.A_REVERSE
                )
                failed = False
            except Exception as e:
                print(f"Adjusting bottom {self.height_offset}: {e}")
                self.height_offset += 1

        self.stdscr.clear()
        self.render()

        while True:
            need_repaint = self.handle_input()
            if need_repaint == True:
                try:
                    self.render()
                except curses.error as e:
                    self.line(self.CONTENT_LINE, f"Curses error: {str(e):<20}")
                    raise e
            elif need_repaint == False:
                break
            # None - skip


def get_overlays(atom: str):
    parser = Parser(atom)

    # with open(Path("/var/lib/portage/world"), 'r') as f:
    #     parser.world_packages = {line.strip(): True for line in f if line.strip()}

    overlays: dict = None

    try:
        result = subprocess.run(
            ["equery", "list", "--format", parser.equery_list_tmpl, atom],
            capture_output=True,
            text=True,
            check=True,
        )
        overlays = parser.parse_equery_list(result.stdout)
    except subprocess.CalledProcessError as e:
        print(f"Error executing equery: {e}", file=sys.stderr)
        sys.exit(1)

    try:
        result = subprocess.run(
            ["equery", "list", "--format", parser.equery_list_set_tmpl, "@installed"],
            capture_output=True,
            text=True,
            check=True,
        )
        parser.installed_packages = parser.parse_equery_list_to_set(result.stdout)
    except subprocess.CalledProcessError as e:
        print(f"Error executing equery: {e}", file=sys.stderr)
        sys.exit(1)

    for overlay_name, overlay in overlays.items():
        print(
            f"Fetching dependencies of {len(overlay)} packages from {overlay_name}..."
        )
        for pkg in overlay:
            try:
                dep_result = subprocess.run(
                    ["equery", "g", "--linear", pkg.cpv],
                    capture_output=True,
                    text=True,
                    check=True,
                )
                deps = parser.parse_dependency_graph(dep_result.stdout)
                pkg.dependencies = deps
                pkg.dependencies_count = len(deps)
            except subprocess.CalledProcessError as e:
                print(
                    f"Warning: Could not get dependencies for {pkg.cpv}: {e}",
                    file=sys.stderr,
                )

    for overlay in overlays.values():
        for pkg in overlay:
            for dep in pkg.dependencies:
                dep.filtered_depend_count = parser.dependencies[
                    dep.cpv
                ].filtered_depend_count

    return overlays


def main():
    if len(sys.argv) != 2:
        print("Usage: depgraph <atom>")
        print("Example: depgraph @selected")
        sys.exit(1)

    atom = sys.argv[1]
    overlays: dict = None

    cache_dir = Path.home() / ".cache" / "nceql" / "overlays"
    cache_file = cache_dir / atom
    if os.path.exists(cache_file):
        try:
            with open(cache_file, "rb") as f:
                overlays = pickle.load(f)
                print(f"Loaded overlays from {cache_file}")
        except Exception as e:
            print(f"Error loading cache: {e}", file=sys.stderr)
            traceback.print_tb(e.__traceback__)

    if overlays is None:
        overlays = get_overlays(atom)

        try:
            cache_dir.mkdir(parents=True, exist_ok=True)
            with open(cache_file, "wb") as f:
                pickle.dump(overlays, f, protocol=pickle.HIGHEST_PROTOCOL)
        except Exception as e:
            print(f"Error saving cache: {e}", file=sys.stderr)
            traceback.print_tb(exception.__traceback__)

    print("Sorting...")
    navigator = PackageNavigator(atom, overlays)
    curses.wrapper(navigator.main_loop)


if __name__ == "__main__":
    main()