Horoscope/main.py

from os import makedirs
from os.path import isfile, join
from pathlib import Path
from random import choice
from sys import argv
from urllib.request import urlretrieve

from PIL import Image, ImageDraw, ImageEnhance, ImageFilter, ImageFont

zodiac_signs = [
    "Belier",
    "Cancer",
    "Balance",
    "Capricorne",
    "Taureau",
    "Lion",
    "Scorpion",
    "Verseau",
    "Gemeaux",
    "Vierge",
    "Sagittaire",
    "Poissons",
]


image_dir = join(Path(argv[0]).parent.resolve(), "images")


def generate_horoscope() -> dict[str, dict[str, str]]:
    """Generate horoscope predictions"""
    horoscope = {}
    for sign in zodiac_signs:
        love = choice(["+++", "++", "+", "-", "- -"])
        work = choice(["+++", "++", "+", "-", "- -"])
        horoscope[sign] = {"love": love, "work": work}
    return horoscope


def get_path(image: str):
    """Return the path of an image"""
    return join(image_dir, f"{image.lower()}.png")


def get_sign_image(image: str):
    """Get sign image"""
    sign_image = Image.open(get_path(image)).convert("RGBA")

    # Add thickness
    alpha = sign_image.split()[3]
    bold_mask = alpha.filter(ImageFilter.MaxFilter())
    sign_image.putalpha(bold_mask)

    # Add contrast
    enhancer = ImageEnhance.Contrast(sign_image)
    enhanced_image = enhancer.enhance(2.0)

    enhanced_image.thumbnail((100, 105))
    return enhanced_image


def vertical_text(text: str, font: ImageFont.ImageFont | ImageFont.FreeTypeFont):
    """Image with vertical text"""
    text_img = Image.new("RGBA", (135, 24), (255, 255, 255, 0))

    text_draw = ImageDraw.Draw(text_img)
    text_draw.text((0, 0), f"{text:^15}", font=font, fill="black")
    text_draw = text_img.rotate(90, expand=True)

    return text_draw


def create_horoscope_image(horoscope: dict[str, dict[str, str]]):
    """Generate horoscope images"""
    img_width, img_height = 1200, 425
    image = Image.new("RGBA", (img_width, img_height), color=(255, 255, 255, 0))
    draw = ImageDraw.Draw(image)

    size_emoji = (30, 30)

    font_sign = ImageFont.load_default(24)
    font_prediction = ImageFont.load_default(40)

    x, y = 10, 10
    for sign, prediction in horoscope.items():
        sign_image = get_sign_image(sign)
        image.paste(sign_image, (x + 20, y), sign_image)

        # Sign text
        text_draw = vertical_text(sign, font_sign)
        image.paste(text_draw, (x - 10, y - 10), text_draw)

        # Love prediction
        love = Image.open(join(image_dir, "love.png")).convert("RGBA")
        love.thumbnail(size_emoji)
        image.paste(love, (x + 125, y + 20), love)
        draw.text(
            (x + 160, y + 10),
            prediction["love"],
            fill="black",
            font=font_prediction,
        )

        # Work prediction
        work = Image.open(join(image_dir, "work.png")).convert("RGBA")
        work.thumbnail(size_emoji)
        image.paste(work, (x + 125, y + 60), work)
        draw.text(
            (x + 160, y + 50),
            prediction["work"],
            fill="black",
            font=font_prediction,
        )

        x += 320
        if x > 1200:
            x = 10
            y += 150

    return image


def download_images():
    """Download images from source"""
    url = "https://git.mylloon.fr/Anri/Horoscope/raw/branch/main/images/"

    makedirs(image_dir, exist_ok=True)
    for el in zodiac_signs + ["love", "work"]:
        path = get_path(el)
        if not isfile(path):
            image = el.lower() + ".png"
            print(f"Download {image}...")
            urlretrieve(url + image, path)


if __name__ == "__main__":
    if len(argv) > 2 or (len(argv) == 2 and argv[1] != "online"):
        exit(1)

    if len(argv) == 2:
        print("Fetch missing images...")
        download_images()

    # Generate new horoscope
    new_horoscope = generate_horoscope()

    # Create and save the image
    horoscope_image = create_horoscope_image(new_horoscope)

    png = "nouvel_horoscope.png"
    horoscope_image.save(png)

    print(f"Nouvel horoscope généré et sauvegardé sous '{png}'")