from bs4 import BeautifulSoup
from mkdocs.structure.files import Files
from typing import NamedTuple
from mkdocs.utils import copy_file
import os
from functools import partial

class Resize(NamedTuple):
    source_path: str
    url: str
    width: int

    def get_resize_url(self):
        return f"/_resized/w/{self.width}/{self.url}"

    def get_dest_path(self):
        return self.get_resize_url().lstrip("/")

RESIZES = set()

def on_pre_build(config):
    RESIZES.clear()

def on_page_content(html, page, config, files: Files):
    soup = BeautifulSoup(html, "html.parser")

    for tag in soup.find_all("img", src=True):
        src = tag.get("src")
        source_file = files.get_file_from_path(src)
        if source_file is None:
            continue

        resize = Resize(source_file.abs_src_path, source_file.url, 1000)
        RESIZES.add(resize)

        tag["src"] = resize.get_resize_url()

    return str(soup)

def resize_image(files: Files, image, width):
    source_file = files.get_file_from_path(image)
    if source_file is None:
        return None

    resize = Resize(source_file.abs_src_path, source_file.url, width)
    RESIZES.add(resize)

    return resize.get_resize_url()

def on_env(env, config, files):
    env.filters["resize_image"] = partial(resize_image, files)

    return env


def on_post_build(config):
    output_path = config["site_dir"]
    for resize in RESIZES:
        # TODO: Actually resize
        copy_file(resize.source_path, os.path.join(output_path, resize.get_dest_path()))