from collections.abc import Iterable, Sequence

from techdocs_graphrag.models import GraphFact, canonical_key

SEED_COLOR = "#C7D2FE"
NODE_COLOR = "#F1F5F9"


def facts_to_dot(facts: Sequence[GraphFact], highlighted_keys: Iterable[str] = ()) -> str:
    highlighted = set(highlighted_keys)
    names = sorted({name for fact in facts for name in (fact.source, fact.target)})
    nodes = [
        f"  {quote(name)} [fillcolor={quote(SEED_COLOR if canonical_key(name) in highlighted else NODE_COLOR)}];"
        for name in names
    ]
    edges = [
        f"  {quote(fact.source)} -> {quote(fact.target)} [label={quote(fact.relation)}];"
        for fact in facts
    ]
    return "\n".join(
        (
            "digraph KnowledgeGraph {",
            "  rankdir=LR;",
            '  node [shape=box, style="rounded,filled", fontname="Helvetica", fontcolor="#0F172A"];',
            '  edge [fontname="Helvetica", fontsize=10, color="#64748B", fontcolor="#64748B"];',
            *nodes,
            *edges,
            "}",
        )
    )


def quote(value: str) -> str:
    escaped = value.replace("\\", "\\\\").replace('"', '\\"').replace("\n", " ")
    return f'"{escaped}"'