hairball.py

#!/usr/bin/env python3
"""Hairball graph builder.

Merges domain-specific Turtle/RDF files into a single ontology, renders three
pyvis-based "hairball" HTML visualisations with a floating AI chat widget, and
exports the JSON context consumed by the chat Lambda backend.
"""

from __future__ import annotations

import argparse
import json
import shutil
import sys
import uuid
from datetime import datetime
from html import escape as html_escape
from pathlib import Path

ASSISTANT_LOGO_SRC = Path(__file__).parent / "images" / "Liber-AI-logo-300px.png"
ASSISTANT_LOGO_FILENAME = "assistant-logo.png"

from pyvis.network import Network
from rdflib import Graph, Literal, Namespace, URIRef
from rdflib.namespace import OWL, RDF, RDFS, XSD


NETWORK_OPTIONS = {
    "nodes": {"font": {"size": 16, "color": "#e8eaed"}},
    "edges": {
        "font": {"size": 14, "color": "#e8eaed",
                  "background": "rgba(14,17,22,0.8)", "strokeWidth": 3,
                  "strokeColor": "#0e1116"},
        "color": {"color": "#9aa0a6", "highlight": "#e8eaed"},
        "smooth": {"type": "continuous"},
    },
    "interaction": {"hover": True},
    "layout": {"improvedLayout": True, "randomSeed": 2},
    "physics": {
        "enabled": True, "solver": "repulsion",
        "repulsion": {"centralGravity": 0.0, "springLength": 240,
                       "springConstant": 0.02, "nodeDistance": 320, "damping": 0.09},
    },
}

NETWORK_OPTIONS_BY_USE_CASE: dict[str, dict] = {
    "graphrag-toolkit": {
        **NETWORK_OPTIONS,
        "physics": {
            **NETWORK_OPTIONS["physics"],
            "repulsion": {
                **NETWORK_OPTIONS["physics"]["repulsion"],
                "springLength": 140,
                "nodeDistance": 200,
            },
        },
    },
}

DEFAULT_CLASS_COLOR = {"background": "#1f3a5f", "border": "#87b0e6"}
LITERAL_COLOR = {"background": "#2c2c2c", "border": "#9aa0a6"}

DOMAIN_PALETTE = [
    {"background": "#1a5276", "border": "#5dade2"},
    {"background": "#196f3d", "border": "#58d68d"},
    {"background": "#7d3c98", "border": "#c39bd3"},
    {"background": "#b9770e", "border": "#f5b041"},
    {"background": "#922b21", "border": "#ec7063"},
    {"background": "#117a65", "border": "#48c9b0"},
    {"background": "#1f618d", "border": "#85c1e9"},
    {"background": "#6c3483", "border": "#bb8fce"},
    {"background": "#b7950b", "border": "#f7dc6f"},
    {"background": "#a04000", "border": "#e59866"},
    {"background": "#148f77", "border": "#76d7c4"},
    {"background": "#2e4053", "border": "#85929e"},
    {"background": "#78281f", "border": "#d98880"},
    {"background": "#4a235a", "border": "#d2b4de"},
    {"background": "#7e5109", "border": "#f0b27a"},
    {"background": "#0e6655", "border": "#73c6b6"},
    {"background": "#5b2c6f", "border": "#af7ac5"},
]

XSD_BASE = str(XSD)


def fix_graphbuild_issues(ttl_path: Path, apply_fixes: bool) -> Graph:
    g = Graph()
    g.parse(str(ttl_path), format="turtle")
    if not apply_fixes:
        return g

    obj_props = [s for s, _, _ in g.triples((None, RDF.type, OWL.ObjectProperty))]
    for p in obj_props:
        g.remove((None, p, None))
    g.remove((None, OWL.imports, None))
    g.remove((RDFS.label, RDF.type, OWL.DatatypeProperty))
    for prop in [RDFS.comment, RDFS.label]:
        latest = {s: o for s, _, o in g.triples((None, prop, None))}
        g.remove((None, prop, None))
        for s, o in latest.items():
            g.add((s, prop, o))
    return g


def upgrade_rdfs_to_owl(g: Graph) -> None:
    """Promote rdfs:Class to owl:Class and rdf:Property to owl:Object/DatatypeProperty.

    The visualisation + context-export queries look for OWL types; this lets
    plain RDFS ontologies (e.g. graphrag-toolkit) flow through unchanged.
    """
    for s in list(g.subjects(RDF.type, RDFS.Class)):
        g.add((s, RDF.type, OWL.Class))
    for p in list(g.subjects(RDF.type, RDF.Property)):
        is_datatype = any(str(r).startswith(XSD_BASE) for r in g.objects(p, RDFS.range))
        g.add((p, RDF.type, OWL.DatatypeProperty if is_datatype else OWL.ObjectProperty))


def domain_label_from_filename(ttl_path: Path, prefixes_to_strip: list[str]) -> str:
    name = ttl_path.stem
    for prefix in prefixes_to_strip:
        if name.startswith(prefix):
            name = name[len(prefix):]
    return name


def merge_domain_models(input_dir: Path, base_uri: str, base_ns: Namespace,
                        ontology_uri: URIRef, prefixes_to_strip: list[str],
                        apply_fixes: bool) -> tuple[Graph, dict[str, list[str]]]:
    ttl_files = sorted(input_dir.glob("*.ttl"))
    if not ttl_files:
        raise FileNotFoundError(f"No .ttl files found in {input_dir}")

    merged = Graph()
    merged.bind("hmo", ontology_uri)
    uri_domains: dict[str, list[str]] = {}

    for ttl in ttl_files:
        print(f"  Loading {ttl.name}")
        sub = fix_graphbuild_issues(ttl, apply_fixes)
        upgrade_rdfs_to_owl(sub)
        domain = domain_label_from_filename(ttl, prefixes_to_strip)
        for uri_type in [OWL.Class, OWL.ObjectProperty, OWL.DatatypeProperty]:
            for s in sub.subjects(RDF.type, uri_type):
                uri_domains.setdefault(str(s), []).append(domain)
        merged += sub

    merged.remove((None, RDF.type, OWL.Ontology))
    merged.add((ontology_uri, RDF.type, OWL.Ontology))

    undeclared = list(merged.query(f"""
        SELECT DISTINCT ?s WHERE {{
            [] rdfs:subClassOf ?s .
            FILTER NOT EXISTS {{ ?s a owl:Class }}
            FILTER (STRSTARTS(STR(?s), "{base_uri}"))
        }}"""))
    for (s,) in undeclared:
        merged.add((s, RDF.type, OWL.Class))
        merged.add((s, RDFS.label, Literal(str(s).split('/')[-1])))

    has_entity_class = any(
        str(o).strip().lower() == "entity"
        for s in merged.subjects(RDF.type, OWL.Class)
        for o in merged.objects(s, RDFS.label)
    )
    if not has_entity_class:
        entity_uri = base_ns['Entity']
        merged.add((entity_uri, RDF.type, OWL.Class))
        merged.add((entity_uri, RDFS.label, Literal("Entity")))
        for (s,) in merged.query("SELECT ?s WHERE { ?s a owl:Class . MINUS { ?s rdfs:subClassOf [] } }"):
            if s != entity_uri:
                merged.add((s, RDFS.subClassOf, entity_uri))

    return merged, uri_domains


def short_label(uri: str, base_uri: str) -> str:
    s = str(uri)
    if s.startswith(base_uri):
        return s[len(base_uri):]
    if '#' in s:
        return s.rsplit('#', 1)[-1]
    tail = s.rstrip('/').rsplit('/', 1)[-1]
    return tail or s


def build_visualisation(ttl_path: Path, version: str, out_dir: Path,
                        uri_domain_map: dict[str, list[str]], base_uri: str,
                        network_options: dict | None = None) -> None:
    network_options = network_options or NETWORK_OPTIONS
    all_domains = sorted(set(d for ds in uri_domain_map.values() for d in ds))
    domain_colors = {d: DOMAIN_PALETTE[i % len(DOMAIN_PALETTE)]
                     for i, d in enumerate(all_domains)}

    g = Graph()
    g.parse(str(ttl_path), format="turtle")
    g.remove((None, RDFS.comment, None))
    g.remove((None, RDF.type, OWL.Class))
    g.remove((None, RDF.type, OWL.ObjectProperty))
    g.remove((None, RDF.type, OWL.DatatypeProperty))
    for (s, o) in g.query("SELECT ?s ?o WHERE { ?s rdfs:subClassOf/rdfs:subClassOf ?o }"):
        g.remove((s, RDFS.subClassOf, o))

    conn_count: dict[str, dict[str, int]] = {}

    def _bump(uri: str) -> None:
        for d in uri_domain_map.get(uri, []):
            conn_count.setdefault(uri, {})
            conn_count[uri][d] = conn_count[uri].get(d, 0) + 1

    for dom_, prop_, rng_ in g.query("SELECT ?dom ?prop ?rng WHERE { ?prop rdfs:domain ?dom ; rdfs:range ?rng }"):
        for u in (str(dom_), str(rng_), str(prop_)):
            _bump(u)
    for s_, o_ in g.query("SELECT ?s ?o WHERE { ?s rdfs:subClassOf ?o }"):
        _bump(str(s_))
        _bump(str(o_))

    def primary_domain(uri: str):
        ds = uri_domain_map.get(uri, [])
        if not ds:
            return None
        if len(ds) == 1:
            return ds[0]
        counts = conn_count.get(uri, {})
        return max(ds, key=lambda d: counts.get(d, 0))

    def color_for(uri: str):
        pd = primary_domain(uri)
        if pd and pd in domain_colors:
            bg, border = domain_colors[pd]["background"], domain_colors[pd]["border"]
            return {"background": bg, "border": border,
                    "highlight": {"background": border, "border": border},
                    "hover": {"background": border, "border": border}}
        return DEFAULT_CLASS_COLOR

    def render_one(add_object_props: bool, add_datatype_props: bool) -> None:
        node_domains: dict[str, list[str]] = {}
        node_primary: dict[str, str] = {}

        net = Network(height="750px", width="100%", directed=True, notebook=False,
                      bgcolor="#0e1116", font_color="#e8eaed")
        net.set_options(json.dumps(network_options))

        def add_node(nid, label, color, uri=None, domain_from=None):
            lookup = uri or domain_from
            if lookup:
                color = color_for(lookup)
                pd = primary_domain(lookup)
                if pd:
                    node_primary[nid] = pd
                node_domains[nid] = uri_domain_map.get(lookup, [])
            else:
                node_domains.setdefault(nid, [])
            tooltip = label + (f"\n\nDomains: {', '.join(node_domains[nid])}" if node_domains.get(nid) else "")
            net.add_node(nid, label=label, shape="box", color=color, title=tooltip)

        for dom_, prop_, rng_ in g.query("SELECT ?dom ?prop ?rng WHERE { ?prop rdfs:domain ?dom ; rdfs:range ?rng }"):
            dom, prop, rng = str(dom_), str(prop_), str(rng_)
            add_node(dom, short_label(dom, base_uri), DEFAULT_CLASS_COLOR, uri=dom)
            if XSD_BASE in rng:
                if not add_datatype_props:
                    continue
                rng_id = rng + str(uuid.uuid4())
                add_node(rng_id, rng.replace(XSD_BASE, "xsd:"), LITERAL_COLOR, domain_from=prop)
            else:
                if not add_object_props:
                    continue
                rng_id = rng
                add_node(rng_id, short_label(rng, base_uri), DEFAULT_CLASS_COLOR, uri=rng)
            prop_domains = uri_domain_map.get(prop, uri_domain_map.get(dom, []))
            net.add_edge(dom, rng_id, label=short_label(prop, base_uri), title=', '.join(prop_domains))

        for s, o in g.query("SELECT ?s ?o WHERE { ?s rdfs:subClassOf ?o }"):
            add_node(str(s), short_label(str(s), base_uri), DEFAULT_CLASS_COLOR, uri=str(s))
            add_node(str(o), short_label(str(o), base_uri), DEFAULT_CLASS_COLOR, uri=str(o))
            net.add_edge(str(s), str(o), label="", dashes=True)

        out = out_dir
        out.mkdir(parents=True, exist_ok=True)
        if add_object_props and add_datatype_props:
            filename = "index.html"
        else:
            filename = f"ontology{'-with-object-properties' if add_object_props else ''}" \
                       f"{'-with-datatype-properties' if add_datatype_props else ''}.html"
        target = str(out / filename)
        html = net.generate_html(target, notebook=False)

        html = html.replace('<script src="lib/bindings/utils.js"></script>\n', '')
        html = html.replace('<meta charset="utf-8">',
                            '<meta charset="utf-8"><title>Ontology Hairball</title>')
        html = html.replace('<center>\n<h1></h1>\n</center>', '')
        html = html.replace('<center>\n          <h1></h1>\n        </center>', '')

        legend_items = ''.join(
            f'<div class="legend-item" data-domain="{html_escape(d, quote=True)}" '
            f'onclick="filterByDomain(domainList[{i}])">'
            f'<span class="legend-swatch" style="background:{domain_colors[d]["background"]};'
            f'border-color:{domain_colors[d]["border"]}"></span>{html_escape(d)}</div>'
            for i, d in enumerate(all_domains)
        )
        header_html = f'''<div id="viz-header">
          <div id="legend" class="legend collapsed">
            <div class="legend-toggle" onclick="toggleLegend()">&#9654; Domains</div>
            <div class="legend-body">
              <div class="legend-item legend-all" onclick="filterByDomain('')">
                <span class="legend-swatch" style="background:#1f3a5f;border-color:#87b0e6"></span>All domains
              </div>
              {legend_items}
            </div>
          </div>
          <span style="color: #cccccc;">Version {version}</span>
          <a href="ontology.html" style="color: #ffffff;">Taxonomy only</a>
          <a href="ontology-with-object-properties.html" style="color: #ffffff;">Classes only</a>
          <a href="." style="color: #ffffff;">Full ontology</a>
        </div>
        <div id="spinner" style="display:none; position:fixed; inset:0; z-index:9999; background:rgba(14,17,22,0.6); align-items:center; justify-content:center;">
          <div class="spinner-ring"></div>
        </div>'''

        filter_script = f'''<script>
var nodeDomains = {json.dumps(node_domains)};
var nodePrimaryDomain = {json.dumps(node_primary)};
var domainColors = {json.dumps(domain_colors)};
var domainList = {json.dumps(all_domains)};
var activeDomain = '';
function toggleLegend() {{
  var el = document.getElementById('legend');
  el.classList.toggle('collapsed');
  var toggle = el.querySelector('.legend-toggle');
  toggle.innerHTML = el.classList.contains('collapsed') ? '&#9654; Domains' : '&#9660; Domains';
}}
function showSpinner() {{ document.getElementById('spinner').style.display = 'flex'; }}
function hideSpinner() {{ document.getElementById('spinner').style.display = 'none'; }}
function filterByDomain(domain) {{
  activeDomain = domain;
  document.querySelectorAll('.legend-item').forEach(function(el) {{
    var d = el.getAttribute('data-domain') || '';
    el.classList.toggle('active', d === domain);
  }});
  if (!domain) document.querySelector('.legend-all').classList.add('active');
  showSpinner();
  setTimeout(function() {{ _applyFilter(domain); hideSpinner(); }}, 30);
}}
function _applyFilter(domain) {{
  if (!domain) {{
    network.body.data.nodes.forEach(function(n) {{
      var pd = nodePrimaryDomain[n.id];
      var color = pd && domainColors[pd]
        ? {{background: domainColors[pd].background, border: domainColors[pd].border,
            highlight: {{background: domainColors[pd].border, border: domainColors[pd].border}},
            hover: {{background: domainColors[pd].border, border: domainColors[pd].border}}}}
        : {{background: '#1f3a5f', border: '#87b0e6'}};
      network.body.data.nodes.update({{id: n.id, hidden: false, color: color}});
    }});
    network.body.data.edges.forEach(function(e) {{
      network.body.data.edges.update({{id: e.id, hidden: false}});
    }});
    return;
  }}
  var visible = new Set();
  for (var nid in nodeDomains) {{
    if (nodeDomains[nid].indexOf(domain) >= 0) visible.add(nid);
  }}
  network.body.data.nodes.forEach(function(n) {{
    network.body.data.nodes.update({{id: n.id, hidden: !visible.has(n.id)}});
  }});
  network.body.data.edges.forEach(function(e) {{
    var show = visible.has(e.from) && visible.has(e.to);
    network.body.data.edges.update({{id: e.id, hidden: !show}});
  }});
}}
</script>'''

        chat_widget_html = f'''<div id="chat-widget">
  <button id="chat-toggle" onclick="toggleChat()" title="Ask AI about this ontology">
    <img src="{ASSISTANT_LOGO_FILENAME}" alt="Ontology Assistant">
  </button>
  <div id="chat-panel">
    <div id="chat-header">
      <span>Ontology Assistant</span>
      <button onclick="toggleChat()" style="background:none;border:none;color:#999;font-size:18px;cursor:pointer;padding:0 4px;">&times;</button>
    </div>
    <div id="chat-messages">
      <div class="chat-row assistant">
        <img class="chat-avatar" src="{ASSISTANT_LOGO_FILENAME}" alt="">
        <div class="chat-msg assistant">Ask me anything about this ontology &mdash; classes, properties, relationships, or domains.</div>
      </div>
    </div>
    <div id="chat-input-wrap">
      <input id="chat-input" type="text" placeholder="Ask about the ontology..." autocomplete="off">
      <button id="chat-send" onclick="sendChat()">
        <svg width="16" height="16" viewBox="0 0 24 24" fill="currentColor"><path d="M2.01 21L23 12 2.01 3 2 10l15 2-15 2z"/></svg>
      </button>
    </div>
  </div>
</div>'''

        chat_widget_script = '''<script>
var chatOpen = false;
var chatMessages = [];
var ontologyVersion = "__ONTOLOGY_VERSION__";
function toggleChat() {
  chatOpen = !chatOpen;
  document.getElementById('chat-panel').style.display = chatOpen ? 'flex' : 'none';
  document.getElementById('chat-toggle').style.display = chatOpen ? 'none' : 'flex';
  if (chatOpen) document.getElementById('chat-input').focus();
}
document.addEventListener('keydown', function(e) { if (e.key === 'Escape' && chatOpen) toggleChat(); });
var chatInputEl = document.getElementById('chat-input');
chatInputEl.addEventListener('keydown', function(e) {
  if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); sendChat(); }
});
function appendMsg(role, text, raw) {
  var el = document.getElementById('chat-messages');
  var bubble = document.createElement('div');
  bubble.className = 'chat-msg ' + role;
  if (raw) { bubble.innerHTML = text; } else { bubble.textContent = text; }
  if (role === 'assistant') {
    var row = document.createElement('div');
    row.className = 'chat-row assistant';
    var avatar = document.createElement('img');
    avatar.className = 'chat-avatar';
    avatar.src = 'assistant-logo.png';
    avatar.alt = '';
    row.appendChild(avatar);
    row.appendChild(bubble);
    el.appendChild(row);
    el.scrollTop = el.scrollHeight;
    return row;
  }
  el.appendChild(bubble);
  el.scrollTop = el.scrollHeight;
  return bubble;
}
function sendChat() {
  var input = document.getElementById('chat-input');
  var q = input.value.trim();
  if (!q) return;
  input.value = '';
  chatMessages.push({role: 'user', content: q});
  appendMsg('user', q);
  var typing = appendMsg('assistant', '<span class="chat-typing"><span></span><span></span><span></span></span>', true);
  fetch('api/chat', {
    method: 'POST',
    headers: {'Content-Type': 'application/json'},
    body: JSON.stringify({messages: chatMessages, version: ontologyVersion})
  })
  .then(function(r) { return r.json(); })
  .then(function(d) {
    typing.remove();
    var reply = d.content || 'Sorry, I could not generate a response.';
    chatMessages.push({role: 'assistant', content: reply});
    appendMsg('assistant', reply);
  })
  .catch(function(err) {
    typing.remove();
    appendMsg('assistant', 'Error: ' + err.message);
  });
}
</script>'''
        chat_widget_script = chat_widget_script.replace('__ONTOLOGY_VERSION__', version)

        extra_css = '''
body { background-color: #0e1116; overflow: hidden; }
.spinner-ring { width: 48px; height: 48px; border: 4px solid rgba(255,255,255,.15); border-top-color: #5dade2; border-radius: 50%; animation: spin .7s linear infinite; }
@keyframes spin { to { transform: rotate(360deg); } }
.card { border: 0; }
#viz-header { padding: 5px 16px; text-align: center; font-size: 24px; display: flex; align-items: center; justify-content: center; gap: 16px; flex-wrap: wrap; position: relative; }
.legend { position: relative; z-index: 100; font-family: sans-serif; font-size: 12px; color: #ccc; margin-right: auto; }
.legend-toggle { padding: 6px 12px; cursor: pointer; font-weight: 600; font-size: 14px; user-select: none; letter-spacing: .3px; white-space: nowrap; background: rgba(14,17,22,0.92); border: 1px solid #444; border-radius: 8px; }
.legend-toggle:hover { color: #fff; }
.legend.collapsed .legend-body { display: none; }
.legend:not(.collapsed) .legend-toggle { border-radius: 8px 8px 0 0; border-bottom: 1px solid #333; }
.legend-body { position: absolute; left: 0; top: 100%; background: rgba(14,17,22,0.95); border: 1px solid #444; border-top: none; border-radius: 0 0 8px 8px; min-width: 200px; box-shadow: 0 4px 16px rgba(0,0,0,.6); padding: 6px 0; max-height: 70vh; overflow-y: auto; text-align: left; }
.legend-item { padding: 5px 14px; display: flex; align-items: center; gap: 8px; cursor: pointer; transition: background .1s; }
.legend-item:hover { background: rgba(255,255,255,.08); }
.legend-item.active { background: rgba(255,255,255,.12); }
.legend-swatch { width: 14px; height: 14px; border-radius: 3px; border: 2px solid; display: inline-block; flex-shrink: 0; }
#chat-widget { position: fixed; bottom: 24px; right: 24px; z-index: 10000; font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif; }
#chat-toggle { width: 60px; height: 60px; border-radius: 50%; border: none; background: #fff; padding: 6px; cursor: pointer; display: flex; align-items: center; justify-content: center; box-shadow: 0 4px 16px rgba(0,0,0,.45); transition: transform .15s, box-shadow .15s; }
#chat-toggle img { width: 100%; height: 100%; object-fit: contain; display: block; }
#chat-toggle:hover { transform: scale(1.08); box-shadow: 0 6px 20px rgba(0,0,0,.55); }
#chat-panel { display: none; flex-direction: column; width: 800px; max-width: 90vw; height: 75vh; background: #1a1a2e; border: 1px solid #333; border-radius: 12px; box-shadow: 0 8px 32px rgba(0,0,0,.5); overflow: hidden; }
#chat-header { padding: 12px 16px; background: #16162a; border-bottom: 1px solid #333; font-size: 14px; font-weight: 600; color: #e8eaed; display: flex; align-items: center; justify-content: space-between; }
#chat-messages { flex: 1; overflow-y: auto; padding: 12px 16px; display: flex; flex-direction: column; gap: 10px; }
.chat-msg { font-size: 13px; line-height: 1.5; padding: 8px 12px; border-radius: 8px; max-width: 90%; word-wrap: break-word; font-family: 'SF Mono', 'Cascadia Code', 'Fira Code', 'Consolas', monospace; white-space: pre-wrap; }
.chat-msg.user { background: #4a6cf7; color: white; align-self: flex-end; border-bottom-right-radius: 2px; }
.chat-msg.assistant { background: #252540; color: #ddd; border-bottom-left-radius: 2px; flex: 1 1 auto; min-width: 0; }
.chat-row.assistant { display: flex; align-items: flex-end; gap: 8px; align-self: flex-start; max-width: 95%; }
.chat-avatar { width: 28px; height: 28px; border-radius: 50%; background: #fff; padding: 3px; box-sizing: border-box; flex-shrink: 0; object-fit: contain; box-shadow: 0 1px 4px rgba(0,0,0,.3); }
#chat-input-wrap { padding: 10px 12px; border-top: 1px solid #333; background: #16162a; display: flex; gap: 8px; align-items: center; }
#chat-input { flex: 1; padding: 8px 12px; border: 1px solid #444; border-radius: 6px; background: #1a1a2e; color: #e8eaed; font-size: 13px; outline: none; font-family: inherit; }
#chat-input:focus { border-color: #4a6cf7; }
#chat-input::placeholder { color: #666; }
#chat-send { width: 34px; height: 34px; border-radius: 6px; border: none; background: #4a6cf7; color: white; cursor: pointer; display: flex; align-items: center; justify-content: center; }
#chat-send:hover { background: #3a5ce5; }
.chat-typing { display: inline-flex; gap: 4px; padding: 4px 0; }
.chat-typing span { width: 6px; height: 6px; background: #666; border-radius: 50%; animation: chatBounce .6s infinite alternate; }
.chat-typing span:nth-child(2) { animation-delay: .2s; }
.chat-typing span:nth-child(3) { animation-delay: .4s; }
@keyframes chatBounce { to { opacity: .3; transform: translateY(-4px); } }
'''
        html = html.replace('</style>', extra_css + '</style>')
        html = html.replace('<body>', f'<body>\n{header_html}')
        html = html.replace('</body>',
                            f'{filter_script}\n{chat_widget_html}\n{chat_widget_script}\n</body>')
        html = html.replace('height: 750px;', 'height: 100vh;')
        html = html.replace('border: 1px solid lightgray;', 'padding: 0; border: 0;')

        with open(target, 'w') as f:
            f.write(html)
        print(f"  Wrote {target}")

    render_one(add_object_props=True, add_datatype_props=True)
    render_one(add_object_props=True, add_datatype_props=False)
    render_one(add_object_props=False, add_datatype_props=False)


def export_ontology_context(ttl_path: Path, version: str, out_dir: Path,
                            uri_domain_map: dict[str, list[str]], base_uri: str) -> None:
    g = Graph()
    g.parse(str(ttl_path), format="turtle")

    def _label(uri):
        return str(uri).replace(base_uri, '')

    def _domains(uri):
        return uri_domain_map.get(str(uri), [])

    classes = []
    for s, label, comment in g.query("""
        SELECT ?s ?label ?comment WHERE {
            ?s a owl:Class ; rdfs:label ?label .
            OPTIONAL { ?s rdfs:comment ?comment }
        }"""):
        parent = None
        for (p,) in g.query(f"SELECT ?p WHERE {{ <{s}> rdfs:subClassOf ?p }}"):
            parent = _label(p)
            break
        classes.append({"uri": str(s), "label": str(label),
                        "comment": str(comment) if comment else "",
                        "domains": _domains(s), "subClassOf": parent})

    obj_props = []
    for s, label, dom, rng in g.query("""
        SELECT ?s ?label ?dom ?rng WHERE {
            ?s a owl:ObjectProperty ; rdfs:label ?label .
            OPTIONAL { ?s rdfs:domain ?dom }
            OPTIONAL { ?s rdfs:range ?rng }
        }"""):
        obj_props.append({"uri": str(s), "label": str(label),
                          "domain": _label(dom) if dom else None,
                          "range": _label(rng) if rng else None,
                          "domains": _domains(s)})

    dt_props = []
    for s, label, dom, rng in g.query("""
        SELECT ?s ?label ?dom ?rng WHERE {
            ?s a owl:DatatypeProperty ; rdfs:label ?label .
            OPTIONAL { ?s rdfs:domain ?dom }
            OPTIONAL { ?s rdfs:range ?rng }
        }"""):
        dt_props.append({"uri": str(s), "label": str(label),
                         "domain": _label(dom) if dom else None,
                         "range": str(rng).split('#')[-1] if rng else None,
                         "domains": _domains(s)})

    hierarchy = [{"child": _label(s), "parent": _label(o)}
                 for s, o in g.query("SELECT ?s ?o WHERE { ?s rdfs:subClassOf ?o }")]

    domain_stats: dict[str, dict[str, int]] = {}
    for uri, domains in uri_domain_map.items():
        for d in domains:
            stats = domain_stats.setdefault(d, {"classes": 0, "objectProperties": 0, "datatypeProperties": 0})
            if (URIRef(uri), RDF.type, OWL.Class) in g:
                stats["classes"] += 1
            elif (URIRef(uri), RDF.type, OWL.ObjectProperty) in g:
                stats["objectProperties"] += 1
            elif (URIRef(uri), RDF.type, OWL.DatatypeProperty) in g:
                stats["datatypeProperties"] += 1

    context = {"version": version, "classes": classes,
               "objectProperties": obj_props, "datatypeProperties": dt_props,
               "hierarchy": hierarchy, "domainStats": domain_stats}

    out_path = out_dir / "ontology-context.json"
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with open(out_path, 'w') as f:
        json.dump(context, f, indent=2)
    print(f"  Wrote {out_path}")
    print(f"  Classes: {len(classes)}")
    print(f"  Object properties: {len(obj_props)}")
    print(f"  Datatype properties: {len(dt_props)}")
    print(f"  Hierarchy edges: {len(hierarchy)}")
    print(f"  Domains: {len(domain_stats)}")


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Merge domain TTL files into a hairball visualisation + chat context JSON.",
    )
    p.add_argument("--input-dir", type=Path, default=Path("./sample-models"),
                   help="Directory containing one *.ttl per domain (default: ./sample-models)")
    p.add_argument("--output-dir", type=Path, default=Path("./output"),
                   help="Directory where the version folder gets written (default: ./output)")
    p.add_argument("--base-uri", default="https://example.com/ontology/",
                   help="Base URI of the merged ontology, used as a label-prefix hint")
    p.add_argument("--version", default=None,
                   help="Output version tag; defaults to today's date (YYYYMMDD)")
    p.add_argument("--strip-prefix", action="append", default=None,
                   metavar="PREFIX",
                   help="Domain-label prefix to strip from TTL filenames. Repeatable. "
                        "Defaults to the H&M Semantic Model prefixes.")
    p.add_argument("--no-graphbuild-fixes", action="store_true",
                   help="Skip graph.build-style cleanup (object-property re-declarations, "
                        "owl:imports, duplicate labels)")
    return p.parse_args(argv)


def process_use_case(use_case: str, input_dir: Path, out_dir: Path, version: str,
                     base_uri: str, prefixes_to_strip: list[str], apply_fixes: bool) -> None:
    base_ns = Namespace(base_uri)
    ontology_uri = URIRef(base_uri)

    print(f"\n=== Use-case: {use_case} ===")
    print(f"Input  : {input_dir.resolve()}")
    print(f"Output : {out_dir.resolve()}")

    print("\nStep 1 — Merging domain models...")
    merged_graph, uri_domain_map = merge_domain_models(
        input_dir, base_uri, base_ns, ontology_uri, prefixes_to_strip, apply_fixes,
    )
    out_dir.mkdir(parents=True, exist_ok=True)
    shutil.copyfile(ASSISTANT_LOGO_SRC, out_dir / ASSISTANT_LOGO_FILENAME)
    ttl_out = out_dir / "merged-ontology.ttl"
    merged_graph.serialize(destination=str(ttl_out), format="turtle", encoding="utf-8")
    print(f"\nMerged ontology saved to {ttl_out}")
    print(f"  Triples: {len(merged_graph)}")
    print(f"  URIs with domain info: {len(uri_domain_map)}")

    print("\nStep 2 — Building visualisations...")
    build_visualisation(ttl_out, version, out_dir, uri_domain_map, base_uri,
                        network_options=NETWORK_OPTIONS_BY_USE_CASE.get(use_case))

    print("\nStep 3 — Exporting ontology context JSON...")
    export_ontology_context(ttl_out, version, out_dir, uri_domain_map, base_uri)


def discover_use_cases(input_dir: Path) -> list[tuple[str, Path]]:
    """Return [(name, dir), ...] for each subdirectory containing *.ttl files.

    Falls back to a single anonymous use-case named after `input_dir` when the
    input contains *.ttl files directly (legacy single-use-case layout).
    """
    subdirs = sorted(d for d in input_dir.iterdir() if d.is_dir() and any(d.glob("*.ttl")))
    if subdirs:
        return [(d.name, d) for d in subdirs]
    if any(input_dir.glob("*.ttl")):
        return [(input_dir.name, input_dir)]
    raise FileNotFoundError(f"No .ttl files (or use-case subdirs) found in {input_dir}")


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)

    input_dir: Path = args.input_dir
    output_dir: Path = args.output_dir
    base_uri: str = args.base_uri
    version: str = args.version or datetime.now().strftime("%Y%m%d")
    prefixes_to_strip = args.strip_prefix if args.strip_prefix is not None else [
        "H&M Semantic Model - ", "HM Semantic Model - ",
    ]
    apply_fixes = not args.no_graphbuild_fixes

    use_cases = discover_use_cases(input_dir)
    print(f"Version: {version}")
    print(f"Use-cases: {', '.join(name for name, _ in use_cases)}")

    for name, uc_input_dir in use_cases:
        process_use_case(
            use_case=name,
            input_dir=uc_input_dir,
            out_dir=output_dir / version / name,
            version=version,
            base_uri=base_uri,
            prefixes_to_strip=prefixes_to_strip,
            apply_fixes=apply_fixes,
        )

    return 0


if __name__ == "__main__":
    sys.exit(main())