#!/usr/bin/env python3
"""
compose_gen — render a team's docker-compose.yml from the challenge registry.

For every ENABLED challenge, a canonical per-challenge compose template lives
at services/<name>/docker-compose.yml (see gen_canonical_composes.py).  The
renderer:

  - replaces `build:` blocks with `image: services-<name>` for the MAIN
    service (sidecars keep their images/builds),
  - rewrites container_name / hostname to the per-team suffix,
  - rewrites host ports (<ORG>:<INT>, <ORG+22>:22) to the team's ports,
  - replaces PASSWORD_<ORG> placeholders with the team's challenge password,
  - normalizes flag volume to ../receiver/flags/<name>.txt.

Multi-container challenges (gemas-notes, gemas-fetcher, kode-viewer,
anti-alchemy, tempest-poc) keep their sidecar services.
"""
import json
import re
import subprocess
from pathlib import Path

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

# Challenges whose compose has multiple top-level services; the FIRST service
# listed is the MAIN challenge service (gets image: reuse + challenge ports),
# the rest are sidecars.
SIDECAR_NAMES = {
    "anti-alchemy":  ["anti-alchemy-db"],
    "gemas-fetcher": ["mongodb"],
    "gemas-notes":   ["database", "validation-service"],
    "kode-viewer":   ["redis"],
    "tempest-poc":   ["backend"],
}

def _load_registry() -> dict:
    try:
        return json.loads((TEAMS_DIR / "challenge_registry.json").read_text())
    except Exception:
        return {"sets": {}, "challenges": []}

def read_template(name: str) -> str:
    p = SERVICES_SRC / name / "docker-compose.yml"
    if not p.exists():
        raise FileNotFoundError(f"Tidak ada template compose untuk {name} di {p}")
    return p.read_text()

def _image_exists(image_name: str) -> bool:
    """True if the docker image is already present locally.

    The renderer only swaps a service's `build:` block for `image: services-<name>`
    when that image actually exists. Otherwise compose would try to PULL a local
    build artifact and fail with "pull access denied for services-<name>".
    """
    r = subprocess.run(["docker", "image", "inspect", image_name],
                       capture_output=True, text=True, timeout=30)
    return r.returncode == 0


def _parse_services(text: str):
    names = []
    for line in text.splitlines():
        m = re.match(r"^  ([A-Za-z0-9_-]+):\s*$", line)
        if m and not line.startswith("    "):
            names.append(m.group(1))
    return names

def render_team_compose(idx: int, state: dict) -> str:
    ports = state["ports"]
    passwords = state["chall_passwords"]
    blocks = []
    for ch in enabled_challenges():
        name = ch["name"]
        text = read_template(name)
        main = _parse_services(text)
        main = main[0] if main else name
        sidecars = set(SIDECAR_NAMES.get(name, []))
        org = int(ch.get("org_port", 10000))
        tc = ports[name]["chall"]
        ts = ports[name]["ssh"]

        # --- rewrite container_name / hostname per team ---
        out_lines = []
        for line in text.splitlines():
            s = line.strip()
            if s.startswith("container_name:"):
                cname = s.split(":", 1)[1].strip()
                line = f"    container_name: {cname}_team{idx}"
            elif s.startswith("hostname:"):
                hname = s.split(":", 1)[1].strip()
                if hname == name:
                    line = f"    hostname: {name}_team{idx}"
                else:
                    # sidecar hostname also suffixed to keep per-team network unique
                    line = f"    hostname: {hname}_team{idx}"
            out_lines.append(line)
        text = "\n".join(out_lines)

        # --- ports: rewrite ONLY the main challenge service's ports ---
        # The main service is the one whose ports map to org/org+22.
        # (sidecar "ports_chall" cases: gemas-notes validation-service exposes
        # 12000:80 — handled by rewriting ANY "<org>:" / "<org+22>:" occurrence.)
        text = re.sub(rf'"({org}):', f'"{tc}:', text)
        text = re.sub(rf'"({org + 22}):', f'"{ts}:', text)

        # --- build: -> image for MAIN only (only when the image exists) ---
        # Replace the main service's build block with image: services-<name> so
        # teams share one image. If that image was never built, KEEP the build
        # block so `docker compose up --build` builds it instead of trying to
        # pull a nonexistent local image.
        use_image = _image_exists(f"services-{name}")
        lines = text.splitlines()
        i = 0
        in_main = False
        cur = None
        out = []
        while i < len(lines):
            line = lines[i]
            m = re.match(r"^  ([A-Za-z0-9_-]+):\s*$", line)
            if m and not line.startswith("    "):
                cur = m.group(1)
                in_main = (cur == main)
                out.append(line)
                i += 1
                continue
            # inside a service block
            if in_main and use_image and line.strip() == "build:":
                # skip build block (context/args/dockerfile...) until next key at same indent
                out.append(f"    image: services-{name}")
                i += 1
                while i < len(lines) and (lines[i].startswith("      ") or lines[i].strip() == ""):
                    i += 1
                continue
            out.append(line)
            i += 1
        text = "\n".join(out)

        # --- PASSWORD placeholder -> team password ---
        text = re.sub(r"\$PASSWORD_" + str(org) + r"\b", passwords[name], text)
        text = re.sub(r"PASSWORD_" + str(org) + r"\b", passwords[name], text)

        # --- build context -> per-team challenge subdir ---
        # Canonical templates are written for the SHARED services/ tree where a
        # challenge's files sit in services/<name>/. The per-team compose lives in
        # teamN/services/, so a bare `context: .` would resolve to the team dir
        # (no Dockerfile). Point the MAIN service's build at ./<name>.
        if re.search(r"^    build:$", text, re.M):
            text = re.sub(r"^    build:\n      context: \.$",
                          f"    build:\n      context: ./{name}", text, count=1, flags=re.M)

        # --- flag volume normalization ---
        # Replace any ./flag.txt / ../receiver/flags/<name>.txt with the per-team flag mount
        text = re.sub(r"\./flag\.txt(:\w+)?", f"../receiver/flags/{name}.txt", text)
        blocks.append(text)
    header = "version: '3.8'\nservices:\n"
    body = []
    for b in blocks:
        if not b.strip():
            continue
        # strip a leading "services:" header from each block (they are fragments)
        b = re.sub(r"^services:\n", "", b)
        body.append(b)
    return header + "\n".join(body) + "\n"

def enabled_challenges() -> list:
    return [c for c in _load_registry().get("challenges", []) if c.get("enabled")]