#!/usr/bin/env python3
"""Shared logic for the staged generator-verifier pipeline:
stage prompts, spec assembly from stages, and the VERIFIER."""
import json, copy
from layout_engine import build_hier
from drawio_engine import validate
from arch_rules import check_rules

SYS = ("You are an SAP solution architect that designs BTP architecture diagrams in "
       "STAGES. Stage 1: plan the BLOCKS (the surrounding containers/zones and their "
       "nesting). Stage 2: fill each block with its SAP COMPONENTS. Stage 3: add the "
       "CONNECTIONS between components. A deterministic engine renders nested blocks "
       "with real SAP icons, so you provide structure only — no coordinates. "
       "Follow SAP best practices: consume GenAI/LLMs via SAP Generative AI Hub / SAP AI Core "
       "(with SAP AI Launchpad as management interface); integrate external identity providers "
       "via SAP Cloud Identity Services; reach S/4HANA / on-premise via Cloud Connector or "
       "Private Link; use SAP Joule for AI copilots; use an SAP Event service (Advanced Event "
       "Mesh) for event-driven flows. A verifier checks these and may ask you to revise. "
       "Always reply with ONLY a JSON object for the requested stage.")
U2 = ("Now populate each block with its SAP components. Reply JSON: "
      '{"components": {"<block_id>": [{"id":"c0","label":"SAP service name"}, ...], ...}}')
U3 = ("Now add the connections between components (use component ids). Reply JSON: "
      '{"connections": [{"source":"c0","target":"c1","label":"optional"}, ...]}')

def strip_components(blocks):
    out = []
    for b in blocks:
        out.append({"id": b["id"], "label": b["label"], "blocks": strip_components(b.get("blocks", []))})
    return out

def comp_map(blocks, m=None):
    if m is None: m = {}
    for b in blocks:
        if b.get("components"):
            m[b["id"]] = [{"id": c["id"], "label": c["label"]} for c in b["components"]]
        comp_map(b.get("blocks", []), m)
    return m

def assemble(title, skeleton, components, connections):
    blocks = copy.deepcopy(skeleton)
    def fill(bs):
        for b in bs:
            b["components"] = components.get(b["id"], [])
            fill(b.get("blocks", []))
    fill(blocks)
    return {"title": title, "blocks": blocks, "connections": connections}

def all_comp_ids(blocks, s=None):
    if s is None: s = set()
    for b in blocks:
        for c in b.get("components", []): s.add(c["id"])
        all_comp_ids(b.get("blocks", []), s)
    return s

def parse_json(text):
    s, e = text.find("{"), text.rfind("}")
    try: return json.loads(text[s:e+1]) if s >= 0 and e > s else {}
    except Exception: return {}

def run_staged(gen, prompt, max_revise=2, log=lambda *a: None, extra_validators=None):
    """Generator-verifier loop. `gen(messages)->assistant_text`. Returns (spec, issues, rounds).
    extra_validators: optional list of fn(spec, prompt)->[issue str] (custom validators; the default
    judge is the deterministic SAP rules engine — no LLM is used)."""
    extra_validators = extra_validators or []
    def full_verify(spec):
        iss = verify(spec)
        for v in extra_validators:
            try: iss = iss + list(v(spec, prompt))
            except Exception as e: log("validator error:", repr(e)[:80])
        return iss
    msgs = [{"role": "system", "content": SYS}, {"role": "user", "content": prompt}]
    def step(user=None):
        if user: msgs.append({"role": "user", "content": user})
        out = gen(msgs); msgs.append({"role": "assistant", "content": out}); return parse_json(out)
    s1 = step(); log("stage1 blocks:", len(s1.get("blocks", [])))
    s2 = step(U2);  log("stage2 components")
    s3 = step(U3);  log("stage3 connections:", len(s3.get("connections", [])))
    title = s1.get("title", prompt[:40])
    spec = assemble(title, s1.get("blocks", []), s2.get("components", {}), s3.get("connections", []))
    issues = full_verify(spec); rounds = 0
    while issues and rounds < max_revise:
        rounds += 1; log(f"verify round {rounds} issues:", issues[:3])
        rr = step(f"Verifier found problems: {'; '.join(issues[:4])}. Fix by WIRING EXISTING components "
                  f"(do NOT add duplicate services; do NOT invent non-SAP components; use official SAP "
                  f"service names). Reply the FULL architecture JSON: "
                  f'{{"blocks":[{{"id","label","blocks","components":[{{"id","label"}}]}}],"connections":[...]}}')
        if isinstance(rr.get("blocks"), list):     # full revised spec
            spec = {"title": title, "blocks": rr["blocks"], "connections": rr.get("connections", spec["connections"])}
        elif "connections" in rr:                  # connections-only revision
            spec = assemble(title, s1.get("blocks", []), s2.get("components", {}), rr["connections"])
        issues = full_verify(spec)
    log("final issues:", issues)
    return spec, issues, rounds

def verify(spec):
    """Return a list of human-readable issues (empty = passes)."""
    issues = []
    cids = all_comp_ids(spec.get("blocks", []))
    if len(cids) < 2:
        issues.append("fewer than 2 components placed")
    def chk(bs):
        for b in bs:
            has = bool(b.get("components")) or any(_nonempty(s) for s in b.get("blocks", []))
            if not has: issues.append(f"block '{b.get('label')}' is empty")
            chk(b.get("blocks", []))
    def _nonempty(b):
        return bool(b.get("components")) or any(_nonempty(s) for s in b.get("blocks", []))
    chk(spec.get("blocks", []))
    connected = set()
    for e in spec.get("connections", []):
        if e["source"] not in cids or e["target"] not in cids:
            issues.append(f"connection {e['source']}->{e['target']} references missing component")
        else:
            connected.add(e["source"]); connected.add(e["target"])
    iso = [c for c in cids if c not in connected]
    if iso:                                            # ANY unconnected component (was: >60% — far too lenient)
        issues.append(f"{len(iso)} unconnected component(s) {iso[:6]}: wire each into the flow")
    if cids and len(spec.get("connections", [])) < len(cids) - 1:   # connection floor (~connected graph)
        issues.append(f"under-connected: {len(spec.get('connections',[]))} connections for {len(cids)} "
                      f"components (expect >= {len(cids)-1})")
    # duplicate detection: same id reused (bad nesting) or the same SAP service placed twice
    all_ids, labels = [], []
    def _collect(bs):
        for b in bs:
            for c in b.get("components", []): all_ids.append(c.get("id")); labels.append(c.get("label", ""))
            _collect(b.get("blocks", []))
    _collect(spec.get("blocks", []))
    dup_ids = sorted({i for i in all_ids if all_ids.count(i) > 1})
    if dup_ids:
        issues.append(f"duplicate component id(s) {dup_ids}: each component must appear once")
    dup_lbl = sorted({l for l in labels if l and labels.count(l) > 1})
    if dup_lbl:
        issues.append(f"duplicate service(s) {dup_lbl}: wire the existing one, do NOT add duplicate services")
    # structural render check
    try:
        xml = open(build_hier(spec, "/tmp/_verify.drawio")).read()
        ok, info = validate(xml)
        if not ok: issues.append(f"does not render: {info}")
    except Exception as ex:
        issues.append(f"render error: {repr(ex)[:80]}")
    # SAP best-practice rules (ported from the PAA architecture-validator)
    issues += check_rules(spec)
    return issues