"""Process tree builder for porTUI htop-style tree view."""

from __future__ import annotations

import ipaddress

import psutil

from portui.models import PortInfo

# Type alias for the rows returned by build()
TreeRow = tuple[PortInfo | None, int | None, str | None, str]


class ProcessTreeBuilder:
    """Builds htop-style tree-ordered rows from a flat list of PortInfo."""

    def __init__(self, sort_column: str | None, sort_ascending: bool) -> None:
        self._sort_column = sort_column
        self._sort_ascending = sort_ascending

    def build(self, filtered: list[PortInfo]) -> list[TreeRow]:
        """Return tree-ordered rows with display prefixes.

        Each element is a tuple of:
          (port, ancestor_pid, ancestor_name, display_prefix)

        For real port-holding processes *port* is set; ancestor fields are None.
        For synthetic ancestor rows *port* is None and *ancestor_pid*/*ancestor_name*
        identify the intermediate process.
        """
        # Group filtered ports by PID
        by_pid: dict[int, list[PortInfo]] = {}
        for p in filtered:
            by_pid.setdefault(p.pid, []).append(p)

        pids_in_set: set[int] = set(by_pid)

        # ancestor_info: pid -> (name, ppid) for intermediate ancestor processes
        ancestor_info: dict[int, tuple[str, int]] = {}

        self._resolve_all_ancestors(by_pid, pids_in_set, ancestor_info)

        # Build parent -> children mapping among all known PIDs
        children_of: dict[int, list[int]] = {}
        parent_of: dict[int, int] = {}

        for pid in pids_in_set:
            ppid = (
                by_pid[pid][0].ppid
                if pid in by_pid
                else ancestor_info.get(pid, ("", 0))[1]
            )
            if ppid in pids_in_set and ppid != pid:
                children_of.setdefault(ppid, []).append(pid)
                parent_of[pid] = ppid

        # Remove ancestors with no port-holding descendants
        stale = [
            pid
            for pid in ancestor_info
            if not self._has_port_descendant(pid, by_pid, children_of)
        ]
        for pid in stale:
            pids_in_set.discard(pid)
            ancestor_info.pop(pid, None)
            children_of.pop(pid, None)
            parent_pid = parent_of.pop(pid, None)
            if parent_pid is not None and parent_pid in children_of:
                children_of[parent_pid] = [
                    c for c in children_of[parent_pid] if c != pid
                ]

        reverse = not self._sort_ascending if self._sort_column else False

        # Sort children by the active sort key
        for ppid in children_of:
            children_of[ppid].sort(
                key=lambda pid: self._pid_sort_key(
                    pid, by_pid, ancestor_info, children_of
                ),
                reverse=reverse,
            )

        # Find root PIDs (no parent in the set)
        roots = sorted(
            (pid for pid in pids_in_set if pid not in parent_of),
            key=lambda pid: self._pid_sort_key(pid, by_pid, ancestor_info, children_of),
            reverse=reverse,
        )

        result: list[TreeRow] = []
        for i, root_pid in enumerate(roots):
            self._walk(
                root_pid,
                "",
                i == len(roots) - 1,
                True,
                by_pid,
                ancestor_info,
                children_of,
                result,
            )
        return result

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _resolve_all_ancestors(
        self,
        by_pid: dict[int, list[PortInfo]],
        pids_in_set: set[int],
        ancestor_info: dict[int, tuple[str, int]],
    ) -> None:
        """Walk up the process tree for every port-PID, adding intermediate ancestors."""
        for pid in list(by_pid):
            ppid = by_pid[pid][0].ppid
            if ppid > 1 and ppid not in pids_in_set:
                self._resolve_ancestors(ppid, pids_in_set, ancestor_info)

    @staticmethod
    def _resolve_ancestors(
        ppid: int,
        pids_in_set: set[int],
        ancestor_info: dict[int, tuple[str, int]],
    ) -> None:
        """Walk up from *ppid*, adding intermediate ancestors to *ancestor_info*."""
        cur = ppid
        visited: set[int] = set()
        chain: list[tuple[int, str, int]] = []
        while cur not in pids_in_set and cur > 1 and cur not in visited:
            visited.add(cur)
            try:
                proc = psutil.Process(cur)
                name = proc.name()
                parent_pid = proc.ppid()
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                break
            chain.append((cur, name, parent_pid))
            cur = parent_pid

        for apid, aname, appid in chain:
            ancestor_info[apid] = (aname, appid)
            pids_in_set.add(apid)

    def _pid_sort_key(
        self,
        pid: int,
        by_pid: dict[int, list[PortInfo]],
        ancestor_info: dict[int, tuple[str, int]],
        children_of: dict[int, list[int]],
    ) -> object:
        """Sort key for a PID — delegates to the first port-holding descendant."""
        is_ancestor = pid not in by_pid
        if not is_ancestor:
            p: PortInfo | None = by_pid[pid][0]
        else:
            p = self._first_port_descendant(pid, by_pid, children_of)
            if p is None:
                return (0,)

        col = self._sort_column
        if col == "process":
            if is_ancestor:
                aname, _ = ancestor_info[pid]
                return aname.lower()
            return p.process_name.lower()
        elif col == "port":
            return p.port
        elif col == "ip":
            try:
                addr = ipaddress.ip_address(p.ip)
                return (0, addr.version, addr)
            except (ValueError, TypeError):
                return (1, 0, p.ip or "")
        elif col == "pid":
            return pid if is_ancestor else p.pid
        elif col == "user":
            return p.username.lower()
        elif col == "state":
            return (p.state or "").lower()
        elif col in ("proto", "protocol"):
            return p.protocol
        else:
            return p.port  # default

    @staticmethod
    def _first_port_descendant(
        pid: int,
        by_pid: dict[int, list[PortInfo]],
        children_of: dict[int, list[int]],
        visited: set[int] | None = None,
    ) -> PortInfo | None:
        """DFS to find the first port-holding descendant.  Cycle-safe."""
        if visited is None:
            visited = set()
        if pid in visited:
            return None
        visited.add(pid)
        if pid in by_pid:
            return by_pid[pid][0]
        for child in children_of.get(pid, []):
            result = ProcessTreeBuilder._first_port_descendant(
                child, by_pid, children_of, visited
            )
            if result is not None:
                return result
        return None

    @staticmethod
    def _has_port_descendant(
        pid: int,
        by_pid: dict[int, list[PortInfo]],
        children_of: dict[int, list[int]],
    ) -> bool:
        if pid in by_pid:
            return True
        return any(
            ProcessTreeBuilder._has_port_descendant(child, by_pid, children_of)
            for child in children_of.get(pid, [])
        )

    @staticmethod
    def _walk(
        pid: int,
        prefix: str,
        is_last: bool,
        is_root: bool,
        by_pid: dict[int, list[PortInfo]],
        ancestor_info: dict[int, tuple[str, int]],
        children_of: dict[int, list[int]],
        result: list[TreeRow],
    ) -> None:
        kids = children_of.get(pid, [])

        if is_root:
            connector = ""
            child_prefix = ""
        else:
            connector = prefix + ("└── " if is_last else "├── ")
            child_prefix = prefix + ("    " if is_last else "│   ")

        if pid in by_pid:
            ports = by_pid[pid]
            for i, port in enumerate(ports):
                result.append((port, None, None, connector if i == 0 else child_prefix))
        else:
            aname, _ = ancestor_info[pid]
            result.append((None, pid, aname, connector))

        for j, child_pid in enumerate(kids):
            ProcessTreeBuilder._walk(
                child_pid,
                child_prefix,
                j == len(kids) - 1,
                False,
                by_pid,
                ancestor_info,
                children_of,
                result,
            )
