"""Cross-Domain Relationship Discovery — Stufe 2: Opus klassifiziert jede Kandidaten-Beziehung
in GENAU EINE Kategorie. Liefert das Rohmaterial der Compliance-Ontologie (insb. SHARED_CAPABILITY
= Capability-Schicht). ANTHROPIC_API_KEY aus ENV (nie hartcodiert). Streaming.

  ANTHROPIC_API_KEY=… python3 classify_relationships.py --pairs /tmp/cd_pairs.json \
      --only-cross-family --out /tmp/cd_classified.json
"""
from __future__ import annotations

import argparse
import json
import os
import re
from collections import Counter

SYS = """Du bist Compliance-Ontologe. Gegeben Paare von Legal Obligations (CRA), bestimme fuer
JEDES Paar GENAU EINE Beziehung. Ziel ist NICHT Aehnlichkeit, sondern die STRUKTURELLE Beziehung.

Kategorien (genau EINE; bei Mehrdeutigkeit gilt diese Prioritaet):
1 SAME_OBLIGATION   — dieselbe rechtliche Pflicht, nur pro Domaene anders formuliert -> MERGE-Kandidat.
2 SUPPORTED_BY      — A ist domaenenspezifische Auspraegung/Teilfall von B ODER A traegt zur Erfuellung von B bei. RICHTUNG angeben.
3 SHARED_CAPABILITY — beide werden durch DIESELBE technische Faehigkeit erfuellt (z.B. MFA, TLS-Verschluesselung, digitale Signatur, Session-Management, Patch-Management, Logging-Pipeline). capability_name (snake_case) angeben.
4 SHARED_PROCEDURE  — beide ueber denselben operativen Prozess erfuellt, ohne gemeinsames technisches Artefakt.
5 SHARED_EVIDENCE   — beide erzeugen/nutzen denselben Nachweis (Audit-Log, SBOM, Release Notes). evidence_name angeben.
6 SHARED_GUIDANCE   — beide berufen sich auf denselben externen Standard (NIST/OWASP/ISO), sonst distinkt.
7 OVERLAP_ONLY      — nur oberflaechliche Wort-/Themenueberlappung, keine echte strukturelle Beziehung.
8 UNRELATED         — Falsch-Positiv der Embedding-Naehe.

Gib AUSSCHLIESSLICH JSON aus:
{"results":[{"i":0,"relation":"SHARED_CAPABILITY","direction":"a->b|b->a|none","capability_name":"","evidence_name":"","reason":"max 18 Woerter"}]}
Regeln: relation = genau eine der 8 Strings. direction nur bei SUPPORTED_BY, sonst "none".
capability_name NUR bei SHARED_CAPABILITY (sonst ""), evidence_name NUR bei SHARED_EVIDENCE (sonst "").
Sei streng: SHARED_GUIDANCE/OVERLAP_ONLY/UNRELATED grosszuegig nutzen; SAME_OBLIGATION nur bei echter Deckungsgleichheit.
Gib fuer JEDES Paar (per Index i) genau ein Ergebnis."""


def build_user(pairs: list[dict]) -> str:
    lines = []
    for i, p in enumerate(pairs):
        lines.append(f'[{i}] A={p["a"]} ({p["fa"]}/{p["ta"]}): {p["da"]}\n'
                     f'     B={p["b"]} ({p["fb"]}/{p["tb"]}): {p["db"]}  [sim={p["sim"]}]')
    return "Paare:\n" + "\n".join(lines)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--pairs", required=True)
    ap.add_argument("--only-cross-family", action="store_true")
    ap.add_argument("--min-sim", type=float, default=0.0)
    ap.add_argument("--model", default="claude-opus-4-8")
    ap.add_argument("--out", required=True)
    a = ap.parse_args()
    d = json.load(open(a.pairs, encoding="utf-8"))
    pairs = [p for p in d["pairs"]
             if (not a.only_cross_family or p["cross_family"]) and p["sim"] >= a.min_sim]

    import anthropic
    client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    with client.messages.stream(model=a.model, max_tokens=24000, system=SYS,
                                messages=[{"role": "user", "content": build_user(pairs)}]) as st:
        msg = st.get_final_message()
    txt = msg.content[0].text
    m = re.search(r"\{.*\}", txt, re.DOTALL)
    data = json.loads(m.group(0) if m else txt)

    res = []
    for r in data.get("results", []):
        i = r.get("i")
        if not isinstance(i, int) or i < 0 or i >= len(pairs):
            continue
        p = pairs[i]
        res.append({"a": p["a"], "fa": p["fa"], "b": p["b"], "fb": p["fb"], "sim": p["sim"],
                    "relation": r.get("relation", "?"), "direction": r.get("direction", "none"),
                    "capability_name": r.get("capability_name", ""),
                    "evidence_name": r.get("evidence_name", ""), "reason": r.get("reason", "")})
    dist = Counter(r["relation"] for r in res)
    out = {"n_pairs": len(pairs), "n_classified": len(res), "distribution": dict(dist),
           "model": a.model, "results": res}
    json.dump(out, open(a.out, "w", encoding="utf-8"), ensure_ascii=False, indent=1)
    print(f"classified {len(res)}/{len(pairs)} | {dict(dist)}")
    print("written:", a.out)


if __name__ == "__main__":
    main()