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)

# Group channels by language
channels_cs = [ch for ch in channels if ch.get("language") == "cs" or ch.get("language") == "sk" ]
channels_en = [ch for ch in channels if ch.get("language") == "en"]
channels_others = [ch for ch in channels if ch.get("language") not in ["cs", "en"]]

# Create JSON strings for each group
channels_cs_json_str = json.dumps(channels_cs)
channels_en_json_str = json.dumps(channels_en)
channels_others_json_str = json.dumps(channels_others)

# Also, full channels for the download-only container
channels_json_str = json.dumps(channels)

compose = {
    "services": {
        "transcriptor_cs": {
            "image": "t0is/madmonq-transcriptor-image:cuda",
            "environment": [
                f"CHANNELS_JSON={channels_cs_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"])
                            }
                        ]
                    }
                }
            }
        },
        "transcriptor_en": {
            "image": "t0is/madmonq-transcriptor-image:cuda",
            "environment": [
                f"CHANNELS_JSON={channels_en_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"])
                            }
                        ]
                    }
                }
            }
        },
        "transcriptor_others": {
            "image": "t0is/madmonq-transcriptor-image:cuda",
            "environment": [
                f"CHANNELS_JSON={channels_others_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.")