from PIL import Image
import io
from web import Requester
from concurrent.futures import ThreadPoolExecutor


class WebImg:
    def __init__(self, task_name: str, url: str):
        self.task_name = task_name
        self.url = url
        self.data = bytearray()

    def fill_img_data(self):
        """
        填充图片data字段
        请求图片url, 获取图片的字节码填充到data字段
        """
        requester = Requester()
        data = requester.fetch_image(self.url)
        if data is None:
            print(f"{self.task_name}, 下载图片失败")
            raise Exception(f"{self.task_name}, 下载图片失败")
        self.data = data


class ImgManager:
    def __init__(self, img_url_list: list[str], task_name: str):
        self.task_name = task_name
        self.img_list = [WebImg(task_name, img_url) for img_url in img_url_list]

    def batch_fill_image_data(self):
        """
        批量填充图片data字段.
        这个函数会批量请求图片url, 获取图片的字节码填充到data字段
        """
        with ThreadPoolExecutor() as executor:
            executor.map(lambda web_img: web_img.fill_img_data(), self.img_list)

    def concatenate_images_vertically(self):
        """
        垂直拼接长图片
        """
        try:
            # 计算拼接后的长图宽度和总高度
            max_width = max(
                Image.open(io.BytesIO(web_img.data)).width for web_img in self.img_list
            )
            total_height = sum(
                Image.open(io.BytesIO(web_img.data)).height for web_img in self.img_list
            )

            # 创建一张新的长图
            long_image = Image.new(
                "RGB", (max_width, total_height), color=(255, 255, 255)
            )

            # 依次将图片在垂直方向上拼接起来
            y_offset = 0
            for web_img in self.img_list:
                img = Image.open(io.BytesIO(web_img.data))
                img_width, img_height = img.size
                x_offset = (max_width - img_width) // 2  # 居中拼接
                long_image.paste(img, (x_offset, y_offset))
                y_offset += img_height

            return long_image

        except Exception as e:
            task_name = self.task_name
            print(f"{task_name}, 拼接图片失败：{e}")
            return None

    def save_long_image(self):
        long_image = self.concatenate_images_vertically()  # 垂直拼接长图片
        long_image.save(f"output/{self.task_name}.png")  # 保存长图到本地