import asyncio
from types import CoroutineType
from typing import Any, Awaitable, Callable, Coroutine
import aiofiles
import yaml
from pydantic import BaseModel

class Config(BaseModel):
    images: list["Images"] = []

class Images(BaseModel):
    source: str
    target: str
    tags: list[str] = []

    
async def run_command(command: str) -> int | None:
    process = await asyncio.create_subprocess_shell(
        command,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE
    )
    stdout, stderr = await process.communicate()

    if stdout:
        print(f"[STDOUT]\n{stdout.decode()}")
    if stderr:
        print(f"[STDERR]\n{stderr.decode()}")
        
    return process.returncode
    
async def docker_pull(image: str) -> int | None:
    print(f"Pulling image {image}...")
    return await run_command(f"docker pull {image}")
    
async def docker_tag(source: str, target: str) -> int | None:
    print(f"Tagging image {source} as {target}...")
    return await run_command(f"docker tag {source} {target}")
    
async def docker_push(image: str) -> int | None:
    print(f"Pushing image {image}...")
    return await run_command(f"docker push {image}")

semaphore = asyncio.Semaphore(5)

async def limited_task[T: Any](task: Callable[[], Coroutine[None, None, T]]) -> T:
    async with semaphore:
        return await task()


async def main():
    async with aiofiles.open('images.yaml', 'r') as file:
        config = await file.read()

    config = yaml.safe_load(config)
    config = Config(**config)

    print("Loaded configuration:")
    print(config)
    
    tasks = []
    for image in config.images:
        if len(image.tags) > 0:
            for tag in image.tags:
                # 定义任务函数
                async def task():
                    if r := await docker_pull(f"{image.source}:{tag}"):
                        if r != 0:
                            return r
                    await docker_tag(f"{image.source}:{tag}", f"{image.target}:{tag}")
                    if r := await docker_push(f"{image.target}:{tag}"):
                        if r != 0:
                            return r
                    await docker_push(f"{image.target}:{tag}")
                    if r := await docker_push(image.target):
                        if r != 0:
                            return r
                    return 0
                        
                # 将协程对象添加到任务列表
                tasks.append(limited_task(task))
        else:
            async def task():
                if r := await docker_pull(image.source):
                    if r != 0:
                        return r
                if r:= await docker_tag(image.source, image.target):
                    if r != 0:
                        return r
                if r := await docker_push(image.target):
                    if r != 0:
                        return r
                return 0
            tasks.append(limited_task(task))

    results = await asyncio.gather(*tasks)
    failed_tasks = 0
    for result in results:
        if result is not None and result != 0:
            failed_tasks += 1
    if failed_tasks > 0:
        raise Exception(f"{failed_tasks} tasks failed.")
    print(f"{len(results)} tasks completed. {len(results) - failed_tasks} succeed, {failed_tasks} failed.")

if __name__ == "__main__":
    asyncio.run(main())