import json
import yaml

# Custom list class to force inline (flow style) YAML formatting
class InlineList(list):
    pass

def inline_list_representer(dumper, data):
    return dumper.represent_sequence('tag:yaml.org,2002:seq', data, flow_style=True)

yaml.add_representer(InlineList, inline_list_representer)

# Load the channels from channels.json
with open("channels.json", "r") as f:
    channels = json.load(f)

# Instead of multiple services, pass all channels as a JSON string to one container
channels_json_str = json.dumps(channels)

compose = {
    "services": {
        "transcriptor": {
            "image": "t0is/madmonq-transcriptor-image:cuda",
            "environment": [
                f"CHANNELS_JSON={channels_json_str}",
                "TIMEDELTA_DAYS=11",
                "TIMEDELTA_DAYS_EXACT=false",
                "CLIP_CREATE_FROM_CHAT=false",
                "TWITCH_CLIENT_ID=a0fuj6tm5ct79clvim9816orphqkov",
                "TWITCH_CLIENT_SECRET=h7whj3yspxgj1909sgcafx6iz1p1es"
            ],
            "volumes": [
                "/shared/transcriptor/clips:/app/clips",
                "/shared/transcriptor/vods:/app/vods",
                "/shared/transcriptor/audio:/app/audio",
                "/shared/transcriptor/chat:/app/chat",
                "/shared/transcriptor/models:/app/models",
                "/shared/transcriptor/transcripts:/app/transcripts"
            ],
            "deploy": {
                "resources": {
                    "reservations": {
                        "devices": [
                            {
                                "driver": "nvidia",
                                "count": "all",
                                "capabilities": InlineList(["gpu"])
                            }
                        ]
                    }
                }
            }
        },
        "downloader": {
            "image": "t0is/madmonq-transcriptor-image:download-only",
            "environment": [
                f"CHANNELS_JSON={channels_json_str}",
                "TIMEDELTA_DAYS=11",
                "TIMEDELTA_DAYS_EXACT=false",
                "CLIP_CREATE_FROM_CHAT=false",
                "TWITCH_CLIENT_ID=a0fuj6tm5ct79clvim9816orphqkov",
                "TWITCH_CLIENT_SECRET=h7whj3yspxgj1909sgcafx6iz1p1es"
            ],
            "volumes": [
                "/shared/transcriptor/clips:/app/clips",
                "/shared/transcriptor/vods:/app/vods",
                "/shared/transcriptor/audio:/app/audio",
                "/shared/transcriptor/chat:/app/chat",
                "/shared/transcriptor/models:/app/models",
                "/shared/transcriptor/transcripts:/app/transcripts"
            ]
        }
    }
}

# Write the docker-compose file
with open("docker-compose.yml", "w") as f:
    yaml.dump(compose, f, default_flow_style=False)

print("docker-compose.yml generated successfully.")