#!/usr/bin/env python3
"""diagram.py — sinh SƠ ĐỒ (Mermaid/Graphviz/…) thành ảnh PNG để gửi vào chat qua MEDIA:.

Mục đích: cho AI Mentor chèn diagram vào bài giảng (mental model, ma trận, flow) — chữ chính xác,
khác với image_gen (AI vẽ, chữ hay méo). KHÔNG bắt model tự base64 — gọi 1 lệnh là xong.

Render qua dịch vụ URL (không cần cài gì trên VPS): mermaid.ink (mặc định) + fallback kroki.io
(kroki hỗ trợ thêm graphviz, plantuml, excalidraw…). Tải PNG về /tmp rồi trả đường dẫn để dùng `MEDIA:`.

⚠️ Diagram gửi qua dịch vụ NGOÀI — chỉ dùng cho nội dung CHUNG (giáo trình), KHÔNG cho dữ liệu khách nhạy cảm.

Lệnh:
  python3 lib/diagram.py mermaid "<code mermaid>"          -> tải PNG, in JSON {ok, path, url}
  python3 lib/diagram.py mermaid "<code>" --url            -> chỉ in URL ảnh (gửi qua send_photo URL)
  python3 lib/diagram.py mermaid-file <đường_dẫn.mmd>      -> đọc code từ file rồi render
  python3 lib/diagram.py <kind> "<code>"                   -> kind: mermaid|graphviz|plantuml|excalidraw (qua kroki)
"""
import sys
import os
import json
import base64
import zlib
import hashlib
import shutil
import subprocess
import urllib.request

OUT_DIR = "/tmp/ai_os_diagrams"


def render_local_dot(code):
    """Render Graphviz dot -> PNG bytes, HOÀN TOÀN LOCAL (không gọi dịch vụ ngoài) → dùng được cho
    dữ liệu khách NHẠY CẢM. Cần binary `dot` (graphviz). Thiếu -> RuntimeError('no_local_renderer')."""
    exe = shutil.which("dot")
    if not exe:
        raise RuntimeError("no_local_renderer: chưa cài graphviz (dot)")
    p = subprocess.run([exe, "-Tpng"], input=code.encode("utf-8"),
                       capture_output=True, timeout=20)
    if p.returncode != 0 or not p.stdout[:8].startswith(b"\x89PNG"):
        raise RuntimeError("dot lỗi: %s" % (p.stderr.decode("utf-8", "replace")[:200]))
    return p.stdout


def mermaid_ink_url(code, theme="default"):
    payload = {"code": code, "mermaid": {"theme": theme}}
    b64 = base64.urlsafe_b64encode(json.dumps(payload).encode("utf-8")).decode()
    return "https://mermaid.ink/img/%s?type=png" % b64


def kroki_url(code, kind="mermaid"):
    comp = zlib.compress(code.encode("utf-8"), 9)
    b64 = base64.urlsafe_b64encode(comp).decode()
    return "https://kroki.io/%s/png/%s" % (kind, b64)


def _fetch(url, timeout=30):
    req = urllib.request.Request(url, headers={"User-Agent": "curl/8"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return r.read()


def render(code, kind="mermaid"):
    """Trả (png_bytes, url). Thử mermaid.ink (nếu mermaid) rồi kroki."""
    candidates = []
    if kind == "mermaid":
        candidates.append(mermaid_ink_url(code))
    candidates.append(kroki_url(code, kind))
    last_err = "unknown"
    for u in candidates:
        try:
            data = _fetch(u)
            if data[:8].startswith(b"\x89PNG"):
                return data, u
            last_err = "phản hồi không phải PNG từ %s" % u
        except Exception as e:
            last_err = "%s (%s)" % (e, u)
    raise RuntimeError(last_err)


def save_png(data):
    os.makedirs(OUT_DIR, exist_ok=True)
    h = hashlib.sha1(data).hexdigest()[:12]
    p = os.path.join(OUT_DIR, "diagram_%s.png" % h)
    with open(p, "wb") as f:
        f.write(data)
    return p


def main():
    args = sys.argv[1:]
    if not args:
        print(json.dumps({"ok": False, "error": "thiếu lệnh"}, ensure_ascii=False)); sys.exit(2)
    cmd = args[0]
    url_only = "--url" in args
    rest = [a for a in args[1:] if a != "--url"]

    # LOCAL render (sovereign, cho dữ liệu nhạy cảm) — dot từ arg hoặc file
    if cmd in ("local-dot", "local-dot-file"):
        if not rest:
            print(json.dumps({"ok": False, "error": "thiếu dot"}, ensure_ascii=False)); sys.exit(2)
        code = open(rest[0], encoding="utf-8").read() if cmd == "local-dot-file" else rest[0]
        try:
            data = render_local_dot(code)
            print(json.dumps({"ok": True, "path": save_png(data), "bytes": len(data), "render": "local"},
                             ensure_ascii=False))
        except Exception as e:
            err = str(e)
            reason = "no_local_renderer" if err.startswith("no_local_renderer") else "render_error"
            print(json.dumps({"ok": False, "reason": reason, "error": err}, ensure_ascii=False)); sys.exit(1)
        return

    if cmd == "mermaid-file":
        if not rest:
            print(json.dumps({"ok": False, "error": "thiếu đường dẫn file"}, ensure_ascii=False)); sys.exit(2)
        code = open(rest[0], encoding="utf-8").read()
        kind = "mermaid"
    else:
        kind = cmd
        if not rest:
            print(json.dumps({"ok": False, "error": "thiếu code diagram"}, ensure_ascii=False)); sys.exit(2)
        code = rest[0]

    try:
        if url_only:
            url = mermaid_ink_url(code) if kind == "mermaid" else kroki_url(code, kind)
            print(json.dumps({"ok": True, "url": url}, ensure_ascii=False)); return
        data, url = render(code, kind)
        path = save_png(data)
        print(json.dumps({"ok": True, "path": path, "url": url, "bytes": len(data)}, ensure_ascii=False))
    except Exception as e:
        print(json.dumps({"ok": False, "error": str(e)}, ensure_ascii=False)); sys.exit(1)


if __name__ == "__main__":
    main()