#!/usr/bin/env python3
"""Generate canonical docker-compose.yml templates under services/<name>/ for
every challenge in the registry (gemastik18 + imported XVI/XVII).

The generated per-challenge file is the SOURCE for compose_gen — i.e. the
per-team compose is rendered from these templates. Uses:

  services/<name>/docker-compose.yml  (canonical, uniform)

If a template already exists for gemastik18 challenges (from the shared
compose), keep it. For imported challenges, build one from the registry.
"""
import json
import re
from pathlib import Path

BASE = Path("/opt/gemastik18-final")
SVC = BASE / "services"
TEAMS_DIR = BASE / "teams"

def load_registry():
    return json.loads((TEAMS_DIR / "challenge_registry.json").read_text())

def internal_port_for(ch: dict) -> int:
    """Best-effort: the container's internal listening port."""
    name = ch["name"]
    known = {
        # web-ish default 8000, others from upstream compose
        "art": 8080, "xl": 3000, "gemas-notes": 80, "pasta": 8000,
        "burvesigner": 80, "hirnfick": 8000, "gemas-fetcher": 8000,
        "s3": 80, "crawlback": 80, "back-to-basic": 8000,
        "anti-alchemy": 5000, "asmr": 8000, "bit-canvas": 8000,
        "fjb": 80, "gift-card": 5000, "gift-voucher": 5000,
        "gleam-drive": 8000, "go-green": 8000, "kode-viewer": 3000,
        "more-less": 8000, "tempest-poc": 80, "ticketer": 5000,
    }
    return known.get(name, 8000)

def sidecar_for(ch: dict) -> dict:
    """Return {sidecar_name: {image, env, volumes, cmd...}} or {}."""
    name = ch["name"]
    if name == "gemas-notes":
        return {
            "database": {
                "image": "mysql:8",
                "container_name": None,  # per-team
                "environment": ["MYSQL_ROOT_PASSWORD=why_my_random_string_password_doesnot_working",
                                "MYSQL_DATABASE=gemasnotes"],
                "volumes": ["./gemas-notes/src/db/init.sql:/docker-entrypoint-initdb.d/init.sql"],
            },
            "validation-service": {
                "build": "./gemas-notes/src/validation-service",
                "container_name": None,
                "depends_on": ["database"],
                "ports_chall": 80,   # challenge port exposed here (REST API)
            },
        }
    if name == "gemas-fetcher":
        return {
            "mongodb": {
                "image": "mongo:4.4",
                "container_name": None,
                "environment": ["MONGO_INITDB_ROOT_USERNAME=ctf",
                                "MONGO_INITDB_ROOT_PASSWORD=asjdkjk23j1k3dsdn2h233j3jj",
                                "MONGO_INITDB_DATABASE=web_fetcher"],
            }
        }
    if name == "kode-viewer":
        return {
            "redis": {
                "image": "redis:alpine",
                "container_name": None,
                "volumes": ["./redis-data:/data"],
            }
        }
    if name == "anti-alchemy":
        return {
            "anti-alchemy-db": {
                "image": "postgres:16.3-alpine",
                "container_name": None,
                "environment": ["POSTGRES_USER=postgres", "POSTGRES_PASSWORD=password"],
                "volumes": ["./db/dump.sql:/docker-entrypoint-initdb.d/init.sql"],
            }
        }
    if name == "tempest-poc":
        return {
            "backend": {
                "build": "./tempest-poc/backend",
                "container_name": None,
                "ports_ssh": 22,
                "extra_hosts": True,
            },
            "frontend": {
                "build": "./tempest-poc/frontend",
                "container_name": None,
                "ports_chall": 80,
                "extra_hosts": True,
            }
        }
    return {}

def render(ch: dict) -> str:
    name = ch["name"]
    org = int(ch.get("org_port", 10000))
    internal = internal_port_for(ch)
    sidecars = sidecar_for(ch)
    # For tempest-poc the main (first) service is frontend; compose_gen's
    # SIDECAR_NAMES marks backend as sidecar.  Order matters: put frontend
    # before backend in the generated file so compose_gen picks frontend as main.
    if name == "tempest-poc":
        sidecars = {"frontend": sidecars.pop("frontend"), **sidecars}
    lines = ["services:"]
    # main service
    cont = f"{name}_container"
    lines += [
        f"  {name}:",
        f"    container_name: {cont}",
        f"    hostname: {name}",
        "    restart: always",
        "    build:",
        "      context: .",
        "      args:",
        f"        - PASSWORD=$PASSWORD_{org}",
    ]
    # volumes: flag + bashrc + preexec (uniform unless challenge differs)
    flag_path = f"../receiver/flags/{name}.txt:/flag.txt:ro"
    # gift-card/gift-voucher mount to /ctf/<name>/flag.txt
    if name in ("gift-card", "gift-voucher"):
        flag_path = f"../receiver/flags/{name}.txt:/ctf/{name}/flag.txt:ro"
    if name == "pasta":
        flag_path = f"../receiver/flags/{name}.txt:/ctf/pasta/flag.txt:ro"
    lines += [
        "    volumes:",
        f"      - {flag_path}",
        "      - ../utils/bashrc:/root/.bashrc:ro",
        "      - ../utils/preexec.sh:/root/.preexec.sh:ro",
    ]
    # ports: external (org / org+22) but rewritten per team by composer
    if name == "gemas-notes":
        # main service = note-service (Go, ssh only). The challenge port is
        # exposed by validation-service (the REST API the checker hits).
        pass
    elif name == "tempest-poc":
        # main (frontend) exposes chall port; backend (sidecar) has ssh.
        pass
    else:
        lines += [
            "    ports:",
            f"      - \"{org}:{internal}\"",
            f"      - \"{org + 22}:22\"",
        ]
    lines.append("    extra_hosts:")
    lines.append('      - "host.docker.internal:host-gateway"')
    # env (challenge-specific)
    if name == "carbeat":
        lines.append("    environment:")
        lines.append("      - FLAG=GEMASTIK18{local_flag}")
    if name == "phew":
        lines.append("    environment:")
        lines.append("      - FLAG=GEMASTIK18{local_flag}")
    if name == "sheesh":
        lines.append("    environment:")
        lines.append("      - FLAG=GEMASTIK18{local_flag}")
    if name == "anti-alchemy":
        lines.append("    environment:")
        lines.append("      - DB_NAME=postgres")
        lines.append("      - DB_USER=postgres")
        lines.append("      - DB_PASS=password")
        lines.append("      - DB_HOST=anti-alchemy-db")
        lines.append("      - DB_PORT=5432")
        lines.append(f"      - SECRET_KEY=$PASSWORD_{org}")
        lines.append("    depends_on:")
        lines.append("      - anti-alchemy-db")
    if name == "gemas-fetcher":
        lines.append("    environment:")
        lines.append("      - MONGO_URI=mongodb://ctf:asjdkjk23j1k3dsdn2h233j3jj@mongodb:27017/web_fetcher?authSource=admin")
        lines.append(f"      - APP_URI=http://0.0.0.0:{org}")
        lines.append("    depends_on:")
        lines.append("      - mongodb")
    if name == "kode-viewer":
        lines.append("    environment:")
        lines.append("      - REDIS_HOST=redis")
        lines.append("      - REDIS_PORT=6379")
        lines.append("    depends_on:")
        lines.append("      - redis")
    if name == "gemas-notes":
        lines.append("    ports:")
        lines.append(f"      - \"{org + 22}:22\"")
        lines.append("    depends_on:")
        lines.append("      - database")
        lines.append("      - validation-service")
        # note: gemas-notes internal port differs (validation-service is the 8000 listener?)
        # upstream maps validation-service:12000:80 and gemas-notes has no port except ssh.
        # Our single main service is gemas-notes which exposes ssh on 22.
    if name == "burvesigner":
        lines.append("    volumes:")
        lines.append("      - ../receiver/files/burvesigner.priv:/priv.data:ro")
    # sidecars
    for sname, sc in sidecars.items():
        cont_s = f"{sname}_container"
        lines.append(f"  {sname}:")
        if sc.get("build"):
            lines += ["    build:",
                      f"      context: {sc['build']}",
                      f"      dockerfile: Dockerfile"]
        elif sc.get("image"):
            lines.append(f"    image: {sc['image']}")
        if sc.get("container_name") is not None:
            lines.append(f"    container_name: {cont_s}")
        if sc.get("environment"):
            lines.append("    environment:")
            for e in sc["environment"]:
                lines.append(f"      - {e}")
        if sc.get("volumes"):
            lines.append("    volumes:")
            for v in sc["volumes"]:
                lines.append(f"      - {v}")
        if sc.get("depends_on"):
            lines.append("    depends_on:")
            for dep in sc["depends_on"]:
                lines.append(f"      - {dep}")
        if sc.get("ports_chall"):
            lines.append("    ports:")
            lines.append(f"      - \"{org}:{sc['ports_chall']}\"")
        if sc.get("ports_ssh"):
            lines.append("    ports:")
            lines.append(f"      - \"{org + 22}:{sc['ports_ssh']}\"")
        if sc.get("extra_hosts"):
            lines.append("    extra_hosts:")
            lines.append('      - "host.docker.internal:host-gateway"')
        if sname == "frontend" and name == "tempest-poc":
            # tempest frontend already handled via ports_chall
            pass
        elif sname == "frontend":
            lines += ["    ports:", f"      - \"{org}:80\""]
            lines.append("    extra_hosts:")
            lines.append('      - "host.docker.internal:host-gateway"')
    return "\n".join(lines) + "\n"

def main():
    reg = load_registry()
    for ch in reg["challenges"]:
        name = ch["name"]
        dst = SVC / name / "docker-compose.yml"
        if name in ("blogpost", "carbeat", "cdn", "phew", "sheesh", "warmup"):
            # keep existing split template (from shared compose)
            tpl = SVC / name / "compose.template.yml"
            if tpl.exists() and not dst.exists():
                text = tpl.read_text()
                lines = [l for l in text.splitlines() if l.strip() and not l.strip().startswith("#")]
                out = []
                for l in lines:
                    if l.strip().startswith("# ---") and out:
                        break
                    out.append(l)
                dst.write_text("\n".join(out).rstrip() + "\n")
                print(f"kept template: {name}")
            continue
        if name == "tempest-poc":
            dst.write_text(render_tempest(ch))
            print(f"wrote canonical compose (tempest): {name}")
            continue
        dst.write_text(render(ch))
        print(f"wrote canonical compose: {name}")


def render_tempest(ch: dict) -> str:
    """tempest-poc: frontend (nginx) is the MAIN challenge service — it gets
    the challenge container name + SSH + challenge port; backend is a sidecar
    service that frontend proxies to (no host port needed)."""
    name = ch["name"]
    org = int(ch.get("org_port", 10000))
    return f"""services:
  {name}:
    container_name: {name}_container
    hostname: {name}
    restart: always
    build:
      context: ./tempest-poc/frontend
      dockerfile: Dockerfile
    ports:
      - "{org}:80"
      - "{org + 22}:22"
    extra_hosts:
      - "host.docker.internal:host-gateway"
  backend:
    build:
      context: ./tempest-poc/backend
      dockerfile: Dockerfile
"""

if __name__ == "__main__":
    main()