import asyncio
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(1)

async def limited_task(task):
    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():
                    await docker_pull(f"{image.source}:{tag}")
                    await docker_tag(f"{image.source}:{tag}", f"{image.target}:{tag}")
                    await docker_push(f"{image.target}:{tag}")
                tasks.append(limited_task(task()))
        else:
            async def task():
                await docker_pull(image.source)
                await docker_tag(image.source, image.target)
                await docker_push(image.target)
            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
            
    print(f"{len(results)} tasks completed. {len(results) - failed_tasks} succeed, {failed_tasks} failed.")

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