# test with:

# docker image build -t headless-fox Docker && docker run -v $PWD/static:/opt/static -ti headless-fox time python annotate_header.py screenshot.png "/opt/static/output.png" "content-type" "Tutaj jest content-type" "etag" "Tutaj jest etag z długim opisem co ma wiele słów i wychodzi poza network inspector"

import os
import sys
import pytesseract
import uuid
import json
from pytesseract import Output
from PIL import Image, ImageDraw, ImageFont

output_file_relative = sys.argv[
    1
]  # this is also the existing source screenshot to annotate. It will be updated in-place

output_file = "/opt/static/" + output_file_relative
domain = sys.argv[2]
needles = sys.argv[3:]

base_url = os.getenv("BASE_URL")

# generator
def partition(lst, size):
    for i in range(0, len(lst), size):
        yield lst[i : i + size]


# print(d)

with Image.open(output_file) as im:
    x_offset = 2054
    y_offset = 313
    cropped = im.crop((x_offset, y_offset, 2875, 1558))
    cropped_filename = "/opt/static/" + uuid.uuid4().hex + ".png"
    cropped.save(cropped_filename)
    d = pytesseract.image_to_data(cropped_filename, output_type=Output.DICT)
    os.remove(cropped_filename)
    draw = ImageDraw.Draw(im)
    n_boxes = len(d["level"])
    font = ImageFont.truetype("/usr/share/fonts/noto/NotoSansDisplay-Medium.ttf", 48)
    found_needles = []
    for i in range(n_boxes):
        (x, y, w, h, text) = (
            d["left"][i],
            d["top"][i],
            d["width"][i],
            d["height"][i],
            d["text"][i],
        )
        for [needle, comment] in partition(needles, 2):
            if needle.lower() in text.lower():
                found_needles.append(needle)
                # modify y so it's aligned not with the top of the text, but with the midline
                y = y + h / 2
                radius = 30
                # offset both y and x
                y = y + y_offset
                x = x + x_offset
                fill = "red"
                line_length = 200
                draw.regular_polygon(
                    ((x - radius - 5, y), radius), n_sides=3, rotation=270, fill=fill
                )
                draw.line((x - radius - 5, y, x - line_length, y), fill=fill, width=10)
                text_w, text_h = draw.textsize(comment, font)
                text_padding = 10
                draw.rectangle(
                    [
                        (x - line_length - text_w - text_padding, y - text_h / 2),
                        (x - line_length + text_padding, y + text_h / 2),
                    ],
                    fill="white",
                )
                draw.text(
                    (x - line_length - 10, y),
                    comment,
                    fill=fill,
                    anchor="rm",
                    font=font,
                )
    im = im.resize((im.width // 2, im.height // 2))
    im.save(output_file, "PNG")
    print(json.dumps({"new_file":
                      {"url": base_url + "/static/" + output_file_relative,
                       "domain": domain,
                       "found_headers": found_needles}}))
