from pathlib import Path
import struct
import zlib


try:
    from PIL import Image, ImageDraw, ImageFont
except ImportError:
    Image = None
    ImageDraw = None
    ImageFont = None


DEFAULT_SCREENSHOT_TIMEOUT_MS = 8000


def _normalize_format(fmt: str | None) -> str:
    fmt = (fmt or "png").lower()
    return "jpeg" if fmt == "jpg" else fmt


def _extract_ieee4882_payload(raw: bytes) -> bytes:
    """Extract payload from an IEEE488.2 definite-length block if present."""
    if not raw or raw[:1] != b"#":
        return raw

    try:
        digits = int(raw[1:2].decode("ascii"))
    except Exception:
        return raw

    if digits <= 0 or len(raw) < 2 + digits:
        return raw

    try:
        size = int(raw[2:2 + digits].decode("ascii"))
    except Exception:
        return raw

    start = 2 + digits
    end = start + size
    return raw[start:end] if end <= len(raw) else raw[start:]


def _extract_png_blob(blob: bytes):
    """Extract a PNG stream from mixed bytes by signature and IEND markers."""
    png_sig = b"\x89PNG\r\n\x1a\n"
    iend_full = b"\x00\x00\x00\x00IEND\xaeB`\x82"

    start = blob.find(png_sig)
    if start < 0:
        return None

    end = blob.find(iend_full, start)
    if end < 0:
        return blob[start:]

    return blob[start:end + len(iend_full)]


def _verify_image(path: Path) -> bool:
    """Return True if Pillow can decode this image file."""
    if Image is None:
        return path.exists() and path.stat().st_size > 0

    try:
        with Image.open(path) as image:
            image.verify()
        return True
    except Exception:
        return False


def _sanitize_png_bytes(blob: bytes):
    """Remove PNG ancillary chunks with invalid CRC while keeping critical chunks."""
    png_sig = b"\x89PNG\r\n\x1a\n"
    if not blob.startswith(png_sig):
        return None

    out = bytearray(png_sig)
    pos = len(png_sig)
    saw_iend = False

    while pos + 12 <= len(blob):
        length = struct.unpack(">I", blob[pos:pos + 4])[0]
        chunk_type = blob[pos + 4:pos + 8]
        end = pos + 12 + length
        if end > len(blob):
            return None

        data = blob[pos + 8:pos + 8 + length]
        crc = struct.unpack(">I", blob[pos + 8 + length:end])[0]
        crc_calc = zlib.crc32(chunk_type)
        crc_calc = zlib.crc32(data, crc_calc) & 0xFFFFFFFF

        is_critical = 65 <= chunk_type[0] <= 90
        if crc == crc_calc or is_critical:
            out.extend(blob[pos:end])

        pos = end
        if chunk_type == b"IEND":
            saw_iend = True
            break

    return bytes(out) if saw_iend else None


def _repair_png_file(path: Path) -> bool:
    """Attempt to rewrite PNG by stripping broken ancillary chunks."""
    try:
        raw = path.read_bytes()
    except Exception:
        return False

    repaired = _sanitize_png_bytes(raw)
    if not repaired:
        return False

    path.write_bytes(repaired)
    return _verify_image(path)


def capture_screenshot_safe(osc, output_path, fmt: str = "png", timeout_ms: int = DEFAULT_SCREENSHOT_TIMEOUT_MS):
    """Capture a screenshot with bounded timeout and PNG repair fallback."""
    output_path = Path(output_path)
    fmt = _normalize_format(fmt)

    try:
        osc.get_screenshot(str(output_path), fmt)
    except TypeError:
        osc.get_screenshot(str(output_path))

    if _verify_image(output_path):
        return output_path
    if fmt == "png" and _repair_png_file(output_path):
        return output_path

    old_timeout = osc.visa_resource.timeout
    osc.visa_resource.timeout = timeout_ms
    try:
        osc.visa_write(f":disp:data? on,off,{fmt}")
        raw = osc.visa_resource.read_raw()
    finally:
        osc.visa_resource.timeout = old_timeout

    payload = _extract_ieee4882_payload(bytes(raw))

    if fmt == "png":
        png = _extract_png_blob(payload)
        output_path.write_bytes(png if png else payload)
        if _repair_png_file(output_path):
            return output_path
    else:
        output_path.write_bytes(payload)

    if not _verify_image(output_path):
        raise RuntimeError(f"Screenshot capture failed or produced invalid image: {output_path}")

    return output_path


def apply_channel_labels_watermark(input_image, output_image, labels):
    """Overlay channel labels in the top-right corner of a screenshot."""
    if Image is None or ImageDraw is None or ImageFont is None:
        raise RuntimeError("Pillow is required for watermarking. Install with: pip install pillow")

    input_image = Path(input_image)
    output_image = Path(output_image)

    image = Image.open(input_image).convert("RGBA")
    overlay = Image.new("RGBA", image.size, (0, 0, 0, 0))
    draw = ImageDraw.Draw(overlay)

    font = ImageFont.load_default()
    line_height = 18
    margin = 16
    channel_colors = {
        1: (255, 234, 0, 255),
        2: (50, 220, 255, 255),
        3: (236, 75, 75, 255),
        4: (171, 120, 255, 255),
    }

    label_rows = [(channel, text) for channel, text in labels.items() if text]
    if not label_rows:
        image.save(output_image)
        return output_image

    max_text_width = max(draw.textlength(text, font=font) for _, text in label_rows)
    box_width = int(max_text_width) + 28
    box_height = len(label_rows) * line_height + 14
    x0 = image.width - margin - box_width
    y0 = margin
    x1 = image.width - margin
    y1 = y0 + box_height

    draw.rounded_rectangle((x0, y0, x1, y1), radius=8, fill=(0, 0, 0, 135))

    y = y0 + 8
    for channel, text in label_rows:
        color = channel_colors.get(channel, (255, 255, 255, 255))
        tx = x1 - 12 - int(draw.textlength(text, font=font))
        draw.text((tx, y), text, fill=color, font=font)
        y += line_height

    out = Image.alpha_composite(image, overlay).convert("RGB")
    out.save(output_image)
    return output_image


def capture_labeled_screenshot(
    osc,
    output_path,
    labels,
    labeled_output_path=None,
    fmt: str = "png",
    timeout_ms: int = DEFAULT_SCREENSHOT_TIMEOUT_MS,
):
    """Capture a screenshot and write a labeled copy with channel annotations."""
    output_path = Path(output_path)
    if labeled_output_path is None:
        labeled_output_path = output_path.with_name(f"{output_path.stem}_labeled{output_path.suffix}")
    else:
        labeled_output_path = Path(labeled_output_path)

    capture_screenshot_safe(osc, output_path, fmt=fmt, timeout_ms=timeout_ms)
    apply_channel_labels_watermark(output_path, labeled_output_path, labels)
    return labeled_output_path