import os
import random
from wand.image import Image
from wand.drawing import Drawing
from wand.color import Color

def load_dataset(captcha_dir: str) -> list:
    dataset = []
    for entry in os.scandir(captcha_dir):
        if not entry.is_dir():
            continue
        count = sum(
            1 for f in os.scandir(entry.path)
            if f.is_file() and f.name.endswith(".png")
        )
        if count > 0:
            dataset.append([entry.name, count])

    if len(dataset) < 2:
        raise ValueError(
            f"captcha_dir '{captcha_dir}' must contain at least 2 category "
            f"sub-folders with .png images; found {len(dataset)}."
        )
    return dataset


def generate_captcha_image(
    captcha_dir: str,
    font_path: str = "./data/fonts/captcha.ttf",
    minimum_ans: int = 3,
    maximum_ans: int = 5
) -> tuple[list[int], bytes]:
    captcha_dataset = load_dataset(captcha_dir)
    pool = list(range(1, 16)) #[1, 2, ..., 15]
    picks = random.randint(minimum_ans, maximum_ans)
    answer_pos = []

    for i in range(picks): # pick random answer slots first
        idx = random.randint(0, len(pool)-1)
        answer_pos.append(pool.pop(idx))

    c = len(captcha_dataset)
    chosen = captcha_dataset[random.randint(0, c-1)] # choose a dataset

    choices = [ds for ds in captcha_dataset if ds[0] != chosen[0]] # wrong answers
    grid = []
    for i in range(16):
        if i in answer_pos:
            grid.append(chosen)
        else:
            grid.append(choices[random.randint(0, len(choices) - 1)])

    # top banner bg
    if random.randint(0, 1) == 0:
        theme = {"bg": "#ebdbb2", "fg": "#1d2021"}
    else:
        theme = {"bg": "#1d2021", "fg": "#ebdbb2"}

    noise_types = ["gaussian", "laplacian"]
    distort_types = ["affine", "shepards"]

    with Image(width=400, height=427, background=Color(theme["bg"])) as im:
        im.format = "jpeg"
        im.background_color = Color(theme["bg"])

        for idx in range(16):
            row = idx // 4
            col = idx % 4

            img_path = os.path.join(
                captcha_dir,
                grid[idx][0],
                f"{random.randint(1, grid[idx][1])}.png",
            )

            with Image(filename=img_path) as tile:
                tile.background_color = Color("black")
                tile.alpha_channel = "remove"
                if random.randint(0, 1) == 1:
                    tile.flop()

                tile.distort(
                    distort_types[random.randint(0, 1)],
                    [
                        0,   0,   random.randint(-15, 15), random.randint(-15, 15),
                        100, 0,   random.randint(80, 120), random.randint(-15, 15),
                        100, 100, random.randint(80, 120), random.randint(80, 120),
                        0,   100, random.randint(-15, 15), random.randint(80, 120),
                    ],
                    best_fit=False,
                )
                tile.noise(noise_types[random.randint(0, 1)])
                im.composite(tile, left=col * 100, top=(row * 100) + 27)

        label = f"Pick {picks} images of {chosen[0].replace('_', ' ')}"

        with Drawing() as draw:
            draw.font = font_path
            draw.font_size = 24
            draw.fill_color = Color(theme["fg"])
            full_metrics = draw.get_font_metrics(im, label)
            pos_x = 200 - (full_metrics.text_width / 2)

            for char in label:
                char_metrics = draw.get_font_metrics(im, char)
                draw.push()
                draw.translate(pos_x, 22)
                draw.rotate(random.randint(-15, 15))
                draw.text(0, 0, char)
                draw.pop()

                pos_x += char_metrics.text_width
            draw(im)
        im.compression_quality = 90
        return answer_pos, im.make_blob("jpeg")




if __name__ == "__main__":
    import sys
    positions, blob = generate_captcha_image(
        captcha_dir="../captcha",
        font_path="../data/fonts/captcha.ttf",
    )
    out = "captcha_out.jpg"
    with open(out, "wb") as fh:
        fh.write(blob)
    print(f"Answer positions : {positions}", file=sys.stderr)
    print(f"Image saved to   : {out}", file=sys.stderr)