"""Shared drawing primitives for the COM Security Masterclass diagrams.

Global style (from DIAGRAM_SPEC.md):
  white canvas; square-corner nodes fill #F6F3EE border #46505C lw=1.2;
  hub/dark node fill #20242A text #EDEBE6; normal arrows #46505C lw=1.2 -|>;
  attack arrows #8F2D24 lw=2.2; defense #256D5B dashed; research #9A6A1F;
  text #20242A DejaVu Sans; code/paths/CLSIDs DejaVu Sans Mono;
  trust boundary dashed #6B7280; step circles fill #8F2D24 white bold 9pt;
  title top-left bold 13pt with thin #E5E0D8 rule; caption bottom-left 9pt #6B7280.
Coordinate system: 10 units per inch, origin bottom-left, equal aspect.
"""
import os
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle, Circle, Polygon, FancyArrowPatch

OUT = os.path.dirname(os.path.abspath(__file__))

# ---- palette -------------------------------------------------------------
INK    = "#20242A"   # primary text / hub fill
SUB    = "#6B7280"   # secondary text / boundaries
FILL   = "#F6F3EE"   # node fill
EDGE   = "#46505C"   # node border / normal arrows
ATTACK = "#8F2D24"   # attack path
DEF    = "#256D5B"   # defense / detection (teal)
RES    = "#9A6A1F"   # research / next-step (amber)
RULE   = "#E5E0D8"   # title rule / grid
HUBTXT = "#EDEBE6"   # text on dark hub

SANS = "DejaVu Sans"
MONO = "DejaVu Sans Mono"

U = 10.0  # units per inch

DASH = (0, (5, 3))       # dashed arrows
BDASH = (0, (4, 3))      # boundary dashes

# arrow kinds -> (color, lw, linestyle, mutation_scale)
KINDS = {
    "normal":   (EDGE,   1.2, "solid", 12),
    "attack":   (ATTACK, 2.2, "solid", 16),
    "defense":  (DEF,    1.2, DASH,    12),
    "research": (RES,    1.2, "solid", 13),
}


class B:
    """A box: center + size, with side-midpoint accessors."""
    def __init__(self, cx, cy, w, h):
        self.cx, self.cy, self.w, self.h = cx, cy, w, h

    @property
    def l(self): return self.cx - self.w / 2.0
    @property
    def r(self): return self.cx + self.w / 2.0
    @property
    def t(self): return self.cy + self.h / 2.0
    @property
    def b(self): return self.cy - self.h / 2.0

    def side(self, s, f=0.5):
        """Midpoint (or fractionally-offset point) of a side.
        f in [0,1] shifts along the side (0 = bottom/left end)."""
        if s == "l":
            return (self.l, self.b + f * self.h)
        if s == "r":
            return (self.r, self.b + f * self.h)
        if s == "t":
            return (self.l + f * self.w, self.t)
        if s == "b":
            return (self.l + f * self.w, self.b)
        raise ValueError(s)


# ---- text metrics ---------------------------------------------------------
def msize(text, fs=9.5, mono=False, weight="normal", padx=2.1, pady=1.6):
    """Estimated box size (w,h) in units for the given text."""
    lines = text.split("\n")
    cw = fs * (0.0865 if mono else 0.0890)
    if weight == "bold":
        cw *= 1.06
    w = max(len(l) for l in lines) * cw + 2 * padx
    h = len(lines) * fs * 0.205 + 2 * pady
    return w, h


# ---- canvas ---------------------------------------------------------------
def canvas(title, w=11.0, h=6.2, caption=None):
    fig = plt.figure(figsize=(w, h))
    fig.patch.set_facecolor("white")
    ax = fig.add_axes([0, 0, 1, 1])
    ax.set_xlim(0, w * U)
    ax.set_ylim(0, h * U)
    ax.set_aspect("equal")
    ax.axis("off")
    ax.set_facecolor("white")
    ty = h * U - 2.7
    ax.text(1.2, ty, title, fontsize=13, fontweight="bold", color=INK,
            family=SANS, ha="left", va="center", zorder=5)
    ax.plot([1.2, w * U - 1.2], [ty - 2.6, ty - 2.6], color=RULE, lw=1.0,
            solid_capstyle="butt", zorder=1)
    if caption:
        ax.text(1.2, 1.8, caption, fontsize=9, color=SUB, family=SANS,
                ha="left", va="center", zorder=5)
    return fig, ax


# ---- primitives -----------------------------------------------------------
def box(ax, cx, cy, w, h, text=None, fs=9.5, mono=False, dark=False,
        fill=None, edge=None, lw=1.2, ls="solid", tc=None, weight="normal",
        zorder=2, linespacing=1.45):
    """Square-corner node. Returns B."""
    fill = fill if fill is not None else (HUB_TXT_FILL(dark))
    edge = edge if edge is not None else EDGE
    tc = tc if tc is not None else (HUBTXT if dark else INK)
    r = Rectangle((cx - w / 2.0, cy - h / 2.0), w, h, facecolor=fill,
                  edgecolor=edge, lw=lw, linestyle=ls, zorder=zorder,
                  joinstyle="miter")
    ax.add_patch(r)
    if text is not None:
        ax.text(cx, cy, text, fontsize=fs, family=(MONO if mono else SANS),
                color=tc, ha="center", va="center", weight=weight,
                zorder=zorder + 1, linespacing=linespacing)
    return B(cx, cy, w, h)


def HUB_TXT_FILL(dark):
    return "#20242A" if dark else FILL


def mbox(ax, cx, cy, text, fs=9.5, mono=False, weight="normal", padx=2.1,
         pady=1.6, **kw):
    """Auto-sized box fitted to text."""
    w, h = msize(text, fs=fs, mono=mono, weight=weight, padx=padx, pady=pady)
    return box(ax, cx, cy, w, h, text, fs=fs, mono=mono, weight=weight, **kw)


def label(ax, x, y, text, fs=9.5, mono=False, color=INK, ha="center",
          va="center", weight="normal", zorder=4, linespacing=1.4):
    ax.text(x, y, text, fontsize=fs, family=(MONO if mono else SANS),
            color=color, ha=ha, va=va, weight=weight, zorder=zorder,
            linespacing=linespacing)


def arrow(ax, p1, p2, kind="normal", label=None, t=0.5, lx=0.0, ly=0.0,
          fs=9.0, mono=False, lc=None, ha="center", va="center",
          connectionstyle=None, ms=None, zorder=3, lweight="normal"):
    """Straight (or connectionstyle-curved) arrow with -|> head."""
    color, lw, ls, mscale = KINDS[kind]
    a = FancyArrowPatch(p1, p2, arrowstyle="-|>",
                        mutation_scale=ms if ms else mscale, color=color,
                        lw=lw, linestyle=ls, shrinkA=0, shrinkB=0,
                        zorder=zorder, connectionstyle=connectionstyle,
                        joinstyle="miter", capstyle="butt")
    ax.add_patch(a)
    if label is not None:
        mx = p1[0] + (p2[0] - p1[0]) * t + lx
        my = p1[1] + (p2[1] - p1[1]) * t + ly
        ax.text(mx, my, label, fontsize=fs, family=(MONO if mono else SANS),
                color=lc if lc else INK, ha=ha, va=va, zorder=4,
                weight=lweight, linespacing=1.35)
    return a


def conn(ax, b1, s1, b2, s2, f1=0.5, f2=0.5, **kw):
    """Arrow between side points of two boxes."""
    return arrow(ax, b1.side(s1, f1), b2.side(s2, f2), **kw)


def elbow(ax, pts, kind="normal", label=None, li=0, t=0.5, lx=0.0, ly=0.0,
          fs=9.0, mono=False, lc=None, ha="center", va="center", ms=None,
          zorder=3):
    """Right-angle polyline through pts with an arrowhead at the end."""
    color, lw, ls, mscale = KINDS[kind]
    for i in range(len(pts) - 2):
        ax.plot([pts[i][0], pts[i + 1][0]], [pts[i][1], pts[i + 1][1]],
                color=color, lw=lw, linestyle=ls, solid_capstyle="butt",
                zorder=zorder)
    arrow(ax, pts[-2], pts[-1], kind=kind, ms=ms if ms else mscale,
          zorder=zorder)
    if label is not None:
        i = min(li, len(pts) - 2)
        p1, p2 = pts[i], pts[i + 1]
        mx = p1[0] + (p2[0] - p1[0]) * t + lx
        my = p1[1] + (p2[1] - p1[1]) * t + ly
        ax.text(mx, my, label, fontsize=fs, family=(MONO if mono else SANS),
                color=lc if lc else INK, ha=ha, va=va, zorder=4,
                linespacing=1.35)


def step_circle(ax, x, y, n, r=1.9):
    """Numbered step marker: filled #8F2D24 circle, white bold 9pt number."""
    c = Circle((x, y), r, facecolor=ATTACK, edgecolor="white", lw=0.8,
               zorder=6)
    ax.add_patch(c)
    ax.text(x, y - 0.12, str(n), fontsize=9, family=SANS, color="white",
            weight="bold", ha="center", va="center", zorder=7)


def boundary(ax, x, y, w, h, text):
    """Dashed trust-boundary rectangle; label top-left inside, 9pt."""
    r = Rectangle((x, y), w, h, facecolor="none", edgecolor=SUB, lw=1.0,
                  linestyle=BDASH, zorder=1, joinstyle="miter")
    ax.add_patch(r)
    ax.text(x + 1.2, y + h - 1.8, text, fontsize=9, family=SANS, color=SUB,
            ha="left", va="center", zorder=4)


def seq_lifeline(ax, x, ytop, ybot):
    """Sequence-diagram lifeline: vertical dashed grey line."""
    ax.plot([x, x], [ybot, ytop], color=SUB, lw=1.0, linestyle=BDASH,
            zorder=1, solid_capstyle="butt")


def diamond(ax, cx, cy, w, h, text, fs=9.0, tc=INK, zorder=2):
    p = Polygon([(cx - w / 2.0, cy), (cx, cy + h / 2.0), (cx + w / 2.0, cy),
                 (cx, cy - h / 2.0)], closed=True, facecolor=FILL,
                edgecolor=EDGE, lw=1.2, zorder=zorder, joinstyle="miter")
    ax.add_patch(p)
    ax.text(cx, cy, text, fontsize=fs, family=SANS, color=tc, ha="center",
            va="center", zorder=zorder + 1, linespacing=1.35)
    return B(cx, cy, w, h)


def note(ax, cx, cy, text, kind="defense", fs=9.0, mono=False, dashed=True,
         padx=2.0, pady=1.5):
    """Callout box: defense -> teal dashed; research -> amber; attack -> red."""
    color = {"defense": DEF, "research": RES, "attack": ATTACK}[kind]
    w, h = msize(text, fs=fs, mono=mono, padx=padx, pady=pady)
    return box(ax, cx, cy, w, h, text, fs=fs, mono=mono, fill="white",
               edge=color, lw=1.2, ls=(DASH if dashed else "solid"), tc=color)


def tick(ax, p1, p2, text, off=(0.0, 2.2), fs=8.5, t=0.5):
    """Thin red tick across an edge + red abuse label (d08)."""
    import math
    mx = p1[0] + (p2[0] - p1[0]) * t
    my = p1[1] + (p2[1] - p1[1]) * t
    dx, dy = p2[0] - p1[0], p2[1] - p1[1]
    L = math.hypot(dx, dy) or 1.0
    px, py = -dy / L * 2.1, dx / L * 2.1
    ax.plot([mx - px, mx + px], [my - py, my + py], color=ATTACK, lw=2.2,
            zorder=4, solid_capstyle="butt")
    ax.text(mx + off[0], my + off[1], text, fontsize=fs, family=SANS,
            color=ATTACK, ha="center", va="center", zorder=4, weight="bold")


def save(fig, name):
    path = os.path.join(OUT, name)
    fig.savefig(path, dpi=150, bbox_inches="tight", facecolor="white")
    plt.close(fig)
    return path
