#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import re
import struct
import sys
import zlib
from pathlib import Path
from typing import Iterable


SCRIPT_DIR = Path(__file__).resolve().parent
VENDOR_ROOT = SCRIPT_DIR / "_vendor"
if str(VENDOR_ROOT) not in sys.path:
    sys.path.insert(0, str(VENDOR_ROOT))

from qrcode.constants import ERROR_CORRECT_M
from qrcode.main import QRCode


DEFAULT_OUTPUT_DIR = SCRIPT_DIR.parent / "output" / "qr"
SERIAL_PREFIXES = ("TalkingQ-", "TQ_")
SAFE_NAME_RE = re.compile(r"[^A-Za-z0-9._-]+")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate a binding QR PNG for the banban device bind flow."
    )
    parser.add_argument("device_id", help="Device ID written into the QR payload.")
    parser.add_argument(
        "serial_number",
        help=f"Device serial number. It must start with one of {SERIAL_PREFIXES!r}.",
    )
    parser.add_argument(
        "-o",
        "--output",
        help="Output PNG path. Defaults to talkingq-url/output/qr/<device_id>.png",
    )
    parser.add_argument(
        "--box-size",
        type=int,
        default=10,
        help="Pixel size of one QR module. Default: 10",
    )
    return parser.parse_args()


def validate_inputs(device_id: str, serial_number: str, box_size: int) -> tuple[str, str]:
    normalized_device_id = device_id.strip()
    normalized_serial_number = serial_number.strip()

    if not normalized_device_id:
        raise ValueError("device_id cannot be empty")
    if not normalized_serial_number:
        raise ValueError("serial_number cannot be empty")
    if not normalized_serial_number.startswith(SERIAL_PREFIXES):
        raise ValueError(f"serial_number must start with one of {SERIAL_PREFIXES}")
    if box_size <= 0:
        raise ValueError("box_size must be greater than 0")

    return normalized_device_id, normalized_serial_number


def build_payload(device_id: str, serial_number: str) -> str:
    return json.dumps(
        {"device_id": device_id, "serial_number": serial_number},
        ensure_ascii=False,
        separators=(",", ":"),
    )


def build_qr_matrix(payload: str) -> list[list[bool]]:
    qr = QRCode(error_correction=ERROR_CORRECT_M, border=4, box_size=10)
    qr.add_data(payload)
    qr.make(fit=True)
    return qr.get_matrix()


def write_png(path: Path, matrix: list[list[bool]], box_size: int) -> None:
    width = len(matrix[0]) * box_size
    height = len(matrix) * box_size
    raw_rows = bytearray()

    for row in matrix:
        expanded_row = bytearray()
        for cell in row:
            pixel = 0 if cell else 255
            expanded_row.extend([pixel] * box_size)
        row_bytes = bytes(expanded_row)
        for _ in range(box_size):
            raw_rows.append(0)
            raw_rows.extend(row_bytes)

    ihdr = struct.pack("!IIBBBBB", width, height, 8, 0, 0, 0, 0)
    compressed = zlib.compress(bytes(raw_rows), level=9)

    with path.open("wb") as fp:
        fp.write(b"\x89PNG\r\n\x1a\n")
        write_png_chunk(fp, b"IHDR", ihdr)
        write_png_chunk(fp, b"IDAT", compressed)
        write_png_chunk(fp, b"IEND", b"")


def write_png_chunk(fp, chunk_type: bytes, data: bytes) -> None:
    fp.write(struct.pack("!I", len(data)))
    fp.write(chunk_type)
    fp.write(data)
    crc = zlib.crc32(chunk_type)
    crc = zlib.crc32(data, crc)
    fp.write(struct.pack("!I", crc & 0xFFFFFFFF))


def resolve_output_path(device_id: str, output_arg: str | None) -> Path:
    if output_arg:
        return Path(output_arg).expanduser().resolve()

    safe_name = SAFE_NAME_RE.sub("_", device_id).strip("._") or "bind_qr"
    return (DEFAULT_OUTPUT_DIR / f"{safe_name}.png").resolve()


def write_payload_copy(txt_path: Path, payload: str) -> None:
    txt_path.write_text(payload, encoding="utf-8")


def main() -> int:
    args = parse_args()
    try:
        device_id, serial_number = validate_inputs(args.device_id, args.serial_number, args.box_size)
    except ValueError as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1
    device_id = 'TalkingQ_XQSN00001005'
    serial_number = 'TQ_XQSN000000001005'
    payload = build_payload(device_id, serial_number)
    output_path = resolve_output_path(device_id, args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    matrix = build_qr_matrix(payload)
    write_png(output_path, matrix, args.box_size)
    write_payload_copy(output_path.with_suffix(".txt"), payload)

    print(output_path)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())