import struct
import zlib
from pathlib import Path

from diagrams import Cluster, Diagram, Edge
from diagrams.generic.blank import Blank
from diagrams.onprem.client import Client, User
from diagrams.onprem.compute import Server
from diagrams.onprem.container import Docker
from diagrams.onprem.database import Mongodb
from diagrams.onprem.inmemory import Redis
from diagrams.onprem.network import Caddy, Internet
from diagrams.onprem.queue import Rabbitmq
from diagrams.onprem.security import Vault

ROOT = Path(__file__).resolve().parents[2]
OUT = ROOT / "docs" / "assets" / "diagrams"
TARGET_RATIO = 16 / 9


def graph_attr() -> dict[str, str]:
    """Return shared Graphviz settings for clean 16:9 diagram output."""

    return {
        "bgcolor": "white",
        "pad": "0.42",
        "ranksep": "1.05",
        "nodesep": "0.72",
        "fontname": "Arial",
        "fontsize": "18",
        "dpi": "220",
        "size": "16,9!",
        "ratio": "fill",
        "splines": "spline",
        "outputorder": "edgesfirst",
    }


def node_attr() -> dict[str, str]:
    """Return readable node text styling while preserving provider icons."""

    return {
        "fontname": "Arial",
        "fontsize": "14",
        "fontcolor": "#172033",
        "labelloc": "b",
        "margin": "0.12",
    }


def edge_attr() -> dict[str, str]:
    """Return restrained edge styling for presentation diagrams."""

    return {
        "fontname": "Arial",
        "fontsize": "11",
        "color": "#64748B",
        "fontcolor": "#334155",
        "arrowsize": "0.75",
        "penwidth": "1.6",
    }


def cluster_attr() -> dict[str, str]:
    """Return subtle cluster styling that keeps related services visually grouped."""

    return {
        "bgcolor": "#F8FAFC",
        "fontname": "Arial",
        "fontsize": "16",
        "fontcolor": "#334155",
        "pencolor": "#CBD5E1",
    }


def diagram(name: str, filename: str, *, direction: str = "LR") -> Diagram:
    """Create a 16:9 icon-based diagram using the diagrams package."""

    return Diagram(
        name,
        filename=str(OUT / filename),
        show=False,
        direction=direction,
        outformat="png",
        curvestyle="curved",
        graph_attr=graph_attr(),
        node_attr=node_attr(),
        edge_attr=edge_attr(),
    )


def normalize_png(filename: str) -> None:
    """Center a rendered PNG on an exact 16:9 white canvas."""

    path = OUT / f"{filename}.png"
    png = _read_png(path)
    target_width, target_height = _target_size(png.width, png.height)
    if (png.width, png.height) == (target_width, target_height):
        return

    channels = png.channels
    white_pixel = b"\xff" * channels
    rows = [bytearray(white_pixel * target_width) for _ in range(target_height)]
    offset_x = (target_width - png.width) // 2
    offset_y = (target_height - png.height) // 2
    for source_y, source_row in enumerate(png.rows):
        target_row = rows[offset_y + source_y]
        start = offset_x * channels
        target_row[start : start + len(source_row)] = source_row

    _write_png(path, PngImage(target_width, target_height, png.color_type, rows))


class PngImage:
    """Small in-memory PNG representation for deterministic 16:9 padding."""

    def __init__(
        self,
        width: int,
        height: int,
        color_type: int,
        rows: list[bytearray],
    ) -> None:
        self.width = width
        self.height = height
        self.color_type = color_type
        self.rows = rows

    @property
    def channels(self) -> int:
        """Return channel count for supported truecolor PNG types."""

        return 4 if self.color_type == 6 else 3


def _target_size(width: int, height: int) -> tuple[int, int]:
    """Calculate the smallest exact 16:9 canvas that contains the diagram."""

    if width / height < TARGET_RATIO:
        return round(height * TARGET_RATIO), height
    return width, round(width / TARGET_RATIO)


def _read_png(path: Path) -> PngImage:
    """Read a non-interlaced 8-bit RGB/RGBA PNG produced by Graphviz."""

    data = path.read_bytes()
    if not data.startswith(b"\x89PNG\r\n\x1a\n"):
        raise ValueError(f"{path} is not a PNG file")

    position = 8
    width = height = color_type = 0
    compressed = bytearray()
    while position < len(data):
        length = struct.unpack(">I", data[position : position + 4])[0]
        chunk_type = data[position + 4 : position + 8]
        chunk_data = data[position + 8 : position + 8 + length]
        position += 12 + length
        if chunk_type == b"IHDR":
            width, height, bit_depth, color_type, _, _, interlace = struct.unpack(
                ">IIBBBBB", chunk_data
            )
            if bit_depth != 8 or color_type not in {2, 6} or interlace != 0:
                raise ValueError(f"Unsupported PNG encoding in {path}")
        elif chunk_type == b"IDAT":
            compressed.extend(chunk_data)
        elif chunk_type == b"IEND":
            break

    channels = 4 if color_type == 6 else 3
    stride = width * channels
    raw = zlib.decompress(bytes(compressed))
    rows: list[bytearray] = []
    previous = bytearray(stride)
    cursor = 0
    for _ in range(height):
        filter_type = raw[cursor]
        cursor += 1
        row = bytearray(raw[cursor : cursor + stride])
        cursor += stride
        _unfilter(row, previous, filter_type, channels)
        rows.append(row)
        previous = row

    return PngImage(width, height, color_type, rows)


def _unfilter(row: bytearray, previous: bytearray, filter_type: int, channels: int) -> None:
    """Reverse PNG scanline filters used by Graphviz output."""

    for index in range(len(row)):
        left = row[index - channels] if index >= channels else 0
        up = previous[index]
        upper_left = previous[index - channels] if index >= channels else 0
        match filter_type:
            case 0:
                value = row[index]
            case 1:
                value = row[index] + left
            case 2:
                value = row[index] + up
            case 3:
                value = row[index] + ((left + up) // 2)
            case 4:
                value = row[index] + _paeth(left, up, upper_left)
            case _:
                raise ValueError(f"Unsupported PNG filter type: {filter_type}")
        row[index] = value & 0xFF


def _paeth(left: int, up: int, upper_left: int) -> int:
    """Return the PNG Paeth predictor for a single channel."""

    prediction = left + up - upper_left
    left_distance = abs(prediction - left)
    up_distance = abs(prediction - up)
    upper_left_distance = abs(prediction - upper_left)
    if left_distance <= up_distance and left_distance <= upper_left_distance:
        return left
    if up_distance <= upper_left_distance:
        return up
    return upper_left


def _write_png(path: Path, png: PngImage) -> None:
    """Write a simple unfiltered PNG with the same color type as the source."""

    raw = bytearray()
    for row in png.rows:
        raw.append(0)
        raw.extend(row)

    ihdr = struct.pack(">IIBBBBB", png.width, png.height, 8, png.color_type, 0, 0, 0)
    chunks = [
        _chunk(b"IHDR", ihdr),
        _chunk(b"IDAT", zlib.compress(bytes(raw), level=6)),
        _chunk(b"IEND", b""),
    ]
    path.write_bytes(b"\x89PNG\r\n\x1a\n" + b"".join(chunks))


def _chunk(chunk_type: bytes, data: bytes) -> bytes:
    """Create one PNG chunk with CRC."""

    return (
        struct.pack(">I", len(data))
        + chunk_type
        + data
        + struct.pack(">I", zlib.crc32(chunk_type + data) & 0xFFFFFFFF)
    )


def link(label: str) -> Edge:
    """Create a consistent labeled connector."""

    return Edge(label=label)


def render_service_architecture() -> None:
    with diagram("Service Architecture", "service_architecture"):
        with Cluster("Client boundary", graph_attr=cluster_attr()):
            browser = Client("Browser")
            demo = Docker("demo-app\nFastAPI BFF")

        with Cluster("Application services", graph_attr=cluster_attr()):
            auth = Server("auth-service\ngRPC API")
            worker = Docker("sms-worker")

        with Cluster("State and messaging", graph_attr=cluster_attr()):
            mongo = Mongodb("MongoDB\nusers + sessions")
            redis = Redis("Redis\nOTP TTL + limits")
            rabbit = Rabbitmq("RabbitMQ\nSMS jobs")

        providers = Internet("Kavenegar / SMS.ir")

        browser >> link("HTTPS / local HTTP") >> demo >> link("gRPC") >> auth
        auth >> link("documents") >> mongo
        auth >> link("ephemeral state") >> redis
        auth >> link("durable job") >> rabbit >> link("consume") >> worker
        worker >> link("provider API") >> providers
    normalize_png("service_architecture")


def render_clean_architecture() -> None:
    with diagram("Clean Architecture Layers", "clean_architecture", direction="TB"):
        with Cluster("Transport", graph_attr=cluster_attr()):
            demo = Client("Demo BFF")
            grpc = Server("gRPC servicer")

        with Cluster("Application", graph_attr=cluster_attr()):
            auth = Server("AuthService\nuse cases")
            security = Vault("JWT + OTP\nsecurity helpers")

        with Cluster("Domain", graph_attr=cluster_attr()):
            entities = Blank("Entities\nUser, Session, Role")
            ports = Blank("Ports\nrepositories + SMS")

        with Cluster("Infrastructure", graph_attr=cluster_attr()):
            mongo = Mongodb("Mongo repositories")
            redis = Redis("Redis OTP store")
            rabbit = Rabbitmq("RabbitMQ publisher")
            sms = Internet("SMS strategies")

        demo >> link("calls") >> grpc >> link("executes") >> auth
        auth >> link("uses") >> entities
        auth >> link("depends on") >> ports
        auth >> link("delegates crypto") >> security
        ports << link("implements") << [mongo, redis, rabbit, sms]
    normalize_png("clean_architecture")


def render_otp_flow() -> None:
    with diagram("OTP Login Flow", "otp_login_flow"):
        user = User("User")

        with Cluster("Request OTP", graph_attr=cluster_attr()):
            demo_request = Client("demo-app")
            auth_request = Server("RequestOtp")
            redis_store = Redis("Redis\nstore HMAC + TTL")
            rabbit = Rabbitmq("RabbitMQ\npublish SMS job")

        with Cluster("Delivery", graph_attr=cluster_attr()):
            worker = Docker("sms-worker")
            sms = Internet("SMS provider")

        with Cluster("Verify OTP", graph_attr=cluster_attr()):
            demo_verify = Client("demo-app")
            auth_verify = Server("VerifyOtp")
            redis_verify = Redis("Redis\ncompare hash")
            mongo = Mongodb("MongoDB\nuser + session")

        user >> link("1. mobile") >> demo_request >> link("2. gRPC") >> auth_request
        auth_request >> link("3. hash only") >> redis_store
        auth_request >> link("4. queued") >> rabbit >> link("5. consume") >> worker
        worker >> link("6. send OTP") >> sms
        user >> link("7. code") >> demo_verify >> link("8. gRPC") >> auth_verify
        auth_verify >> link("9. verify") >> redis_verify
        auth_verify >> link("10. issue session") >> mongo
        auth_verify >> link("11. token pair") >> demo_verify
    normalize_png("otp_login_flow")


def render_refresh_flow() -> None:
    with diagram("Refresh Token Rotation", "refresh_rotation_flow"):
        client = Client("Client")
        auth = Server("auth-service\nRefreshToken")
        mongo = Mongodb("MongoDB\nrefresh_sessions")

        with Cluster("Rotation result", graph_attr=cluster_attr()):
            old = Vault("Old hash\nrevoked")
            new = Vault("New hash\nactive")

        client >> link("1. opaque refresh token") >> auth
        auth >> link("2. hash lookup") >> mongo >> link("3. active session") >> auth
        auth >> link("4. revoke") >> old
        auth >> link("5. create") >> new >> link("6. persist") >> mongo
        auth >> link("7. new token pair") >> client
    normalize_png("refresh_rotation_flow")


def render_sms_strategy() -> None:
    with diagram("SMS Provider Strategy", "sms_provider_strategy", direction="TB"):
        worker = Docker("sms-worker")
        factory = Server("create_sms_client()\nprovider factory")
        port = Blank("SmsClient port\nsend_otp(...)")

        with Cluster("Concrete strategies", graph_attr=cluster_attr()):
            kavenegar = Internet("Kavenegar")
            sms_ir = Internet("SMS.ir")
            debug = Redis("Debug provider\nlocal Redis")

        worker >> link("startup wiring") >> factory >> link("returns interface") >> port
        port >> link("SMS_PROVIDER=kavenegar") >> kavenegar
        port >> link("SMS_PROVIDER=sms_ir") >> sms_ir
        port >> link("SMS_PROVIDER=debug") >> debug
    normalize_png("sms_provider_strategy")


def render_production_deployment() -> None:
    with diagram("Production Deployment", "production_deployment"):
        internet = Internet("Internet\ngapido.amiirkhl.ir")
        caddy = Caddy("Caddy\n80 / 443 only")

        with Cluster("Private Docker network", graph_attr=cluster_attr()):
            demo = Docker("demo-app\nFastAPI UI")
            auth = Server("auth-service\nprivate gRPC")
            mongo = Mongodb("MongoDB\nprivate")
            redis = Redis("Redis\nprivate")
            rabbit = Rabbitmq("RabbitMQ\nprivate")
            worker = Docker("sms-worker")

        sms = Internet("SMS provider")

        internet >> link("HTTPS") >> caddy >> link("reverse proxy") >> demo
        demo >> link("gRPC") >> auth
        auth >> link("sessions") >> mongo
        auth >> link("OTP state") >> redis
        auth >> link("SMS job") >> rabbit >> link("consume") >> worker
        worker >> link("provider API") >> sms
    normalize_png("production_deployment")


def main() -> None:
    """Regenerate all 16:9 icon-based diagrams for docs and slides."""

    OUT.mkdir(parents=True, exist_ok=True)
    render_service_architecture()
    render_clean_architecture()
    render_otp_flow()
    render_refresh_flow()
    render_sms_strategy()
    render_production_deployment()


if __name__ == "__main__":
    main()