#! /usr/bin/env python3

from collections import defaultdict
from dataclasses import dataclass
from pprint import pprint
import json
import subprocess
from time import sleep
from typing import Iterable, List, Set, NamedTuple, Dict

delay_seconds = 0.5
call_limit = 4

PHID = str

conduit_uri = "http://localhost.aviv.pw:8080/"
start_from = ["PHID-TASK-pz4qindlie7wqhn7ez6j"]

EDGE_TYPES = {"parent": "task.parent", "child": "task.subtask"}


@dataclass
class VisitedItem:
    parents: Set[PHID]
    children: Set[PHID]

    @staticmethod
    def new():
        return VisitedItem(set(), set())


class EdgeData(NamedTuple):
    src: PHID
    dst: PHID
    type: str

    @staticmethod
    def from_dict(input: dict):
        return EdgeData(
            src=input["sourcePHID"],
            dst=input["destinationPHID"],
            type=input["edgeType"],
        )


def get_edges(phids: Iterable[PHID]) -> List[EdgeData]:
    params = {
        "sourcePHIDs": list(phids),
        "types": list(EDGE_TYPES.values()),
        "limit": 1000,
    }

    print("Making conduit call...")

    params = json.dumps(params)
    response_raw = subprocess.check_output(
        ["arc", "--conduit-uri", conduit_uri, "call-conduit", "--", "edge.search"],
        text=True,
        input=params,
    )

    response = json.loads(response_raw)
    assert response["error"] is None, response_raw
    response = response["response"]
    assert response["cursor"]["after"] is None, response_raw
    return [EdgeData.from_dict(data) for data in response["data"]]


graph_data: Dict[PHID, VisitedItem] = defaultdict(VisitedItem.new)
to_visit: Set[PHID] = set(start_from)
visited: Set[PHID] = set()


while to_visit and call_limit > 0:
    new_edges = get_edges(to_visit)
    call_limit -= 1

    mentioned_phids = {e.src for e in new_edges} | {e.dst for e in new_edges}

    child_edges = {e for e in new_edges if e.type == EDGE_TYPES["child"]}
    parent_edges = {e for e in new_edges if e.type == EDGE_TYPES["parent"]}

    for edge in child_edges:
        graph_data[edge.src].children.add(edge.dst)
        graph_data[edge.dst].parents.add(edge.src)
    for edge in parent_edges:
        graph_data[edge.dst].children.add(edge.src)
        graph_data[edge.src].parents.add(edge.dst)

    visited |= to_visit
    to_visit = mentioned_phids - visited

    if call_limit <= 0 and to_visit:
        print("Reached call limit with unvisited nodes.")
        break

    sleep(delay_seconds)


class JsonEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, VisitedItem):
            return {
                "parents": list(obj.parents),
                "children": list(obj.children),
            }
        print(object.__class__.__name__)
        return json.JSONEncoder.default(self, obj)

with open("task_graph.json", "wt") as storage:
    json.dump(dict(graph_data), storage, indent=2, cls=JsonEncoder)

#! /usr/bin/env python3

import json


with open("task_graph.json") as storage:
    data = json.load(storage)


graph = set()

for phid, edges in data.items():
    graph |= {f"{phid} --> {child}" for child in edges["children"]}
    graph |= {f"{parent} --> {phid}" for parent in edges["parents"]}


graph = "\n".join(graph)
graph = graph.replace("PHID-TASK-", "")

with open("graph-code.txt", "wt") as output:
    output.write('stateDiagram\n')
    output.write(graph)