#!/usr/bin/env python3
"""
Modify the reference.conf inside the Ditto gateway JAR.
Inserts ditto { gateway { authentication { pre-authentication { enabled = true } } } }
at the root level, before the final closing brace.
"""
import zipfile
import shutil
import os
import subprocess
import sys

JAR_PATH = "/tmp/ditto-jar-mod/ditto-gateway-service-3.8.12-allinone.jar"
AUTH_BLOCK = """
### Custom Ditto auth override - pre-authentication enabled
ditto {
  gateway {
    authentication {
      pre-authentication {
        enabled = true
      }
      devops {
        secured = false
        devops-authentication-method = "basic"
        password = "ditto-devops-secret"
        status-secured = false
        status-authentication-method = "basic"
        statusPassword = "ditto-status-secret"
      }
    }
  }
}
"""

def main():
    os.makedirs("/tmp/ditto-jar-mod", exist_ok=True)

    print("=== Step 1: Extracting JAR ===")
    result = subprocess.run(
        ["docker", "run", "--rm", "eclipse/ditto-gateway:latest",
         "cat", "/opt/ditto/ditto-gateway-service-3.8.12-allinone.jar"],
        capture_output=True, check=True
    )
    with open(JAR_PATH, "wb") as f:
        f.write(result.stdout)
    print(f"JAR: {len(result.stdout)} bytes")

    print("=== Step 2: Modifying reference.conf ===")
    with zipfile.ZipFile(JAR_PATH, 'r') as zin:
        ref_conf = zin.read("reference.conf").decode("utf-8")

    lines = ref_conf.split('\n')
    total_lines = len(lines)
    print(f"reference.conf: {total_lines} lines")

    # Find the LAST closing brace at root level (depth 0)
    # Track depth through the entire file
    depth = 0
    last_root_close_idx = None

    for i, line in enumerate(lines):
        stripped = line.strip()
        if not stripped or stripped.startswith('#'):
            continue

        # Count braces (simple approach - count { and })
        # This isn't perfect for HOCON but works for our case
        for ch in stripped:
            if ch == '{':
                depth += 1
            elif ch == '}':
                depth -= 1

        if depth == 0:
            last_root_close_idx = i

    if last_root_close_idx is None:
        print("ERROR: Could not find root-level closing brace!")
        sys.exit(1)

    insert_idx = last_root_close_idx
    print(f"Last root-level '}}' at line {insert_idx + 1}: '{lines[insert_idx].strip()}'")

    # Check if we need a comma before our block
    # Look at the non-empty line before insert_idx
    prev_idx = insert_idx - 1
    while prev_idx >= 0 and lines[prev_idx].strip() == '':
        prev_idx -= 1

    if prev_idx >= 0:
        prev_stripped = lines[prev_idx].strip()
        if prev_stripped.endswith('}'):
            # Need to add a comma
            lines[prev_idx] = lines[prev_idx].rstrip()
            if not lines[prev_idx].endswith(','):
                lines[prev_idx] += ','
            print(f"Added comma to line {prev_idx + 1}")

    # Insert our block
    auth_lines = AUTH_BLOCK.split('\n')
    new_lines = lines[:insert_idx] + auth_lines + lines[insert_idx:]
    modified_conf = '\n'.join(new_lines)

    print(f"Modified: {total_lines} -> {len(new_lines)} lines")

    # Verify brace balance
    depth = 0
    for i, line in enumerate(new_lines):
        stripped = line.strip()
        if not stripped or stripped.startswith('#'):
            continue
        for ch in stripped:
            if ch == '{':
                depth += 1
            elif ch == '}':
                depth -= 1
        if depth < 0:
            print(f"ERROR: Negative depth at line {i+1}")
            sys.exit(1)

    print(f"Brace depth check: {depth} (should be 0)")
    if depth != 0:
        print("ERROR: Unbalanced braces!")
        sys.exit(1)

    # Verify ditto is at root level
    for i, line in enumerate(new_lines):
        if line.strip() == 'ditto {':
            indent = len(line) - len(line.lstrip())
            print(f"'ditto {{' at line {i+1}, indent: {indent}")
            if indent != 0:
                print(f"WARNING: Expected indent 0, got {indent}")
            break

    print("=== Step 3: Creating unsigned JAR ===")
    skip_files = set()
    with zipfile.ZipFile(JAR_PATH, 'r') as zin:
        for name in zin.namelist():
            if name.startswith("META-INF/"):
                upper = name.upper()
                if upper.endswith(".SF") or upper.endswith(".RSA") or upper.endswith(".DSA") or upper == "MANIFEST.MF":
                    skip_files.add(name)

        with zipfile.ZipFile(JAR_PATH + ".new", 'w', zipfile.ZIP_DEFLATED) as zout:
            for item in zin.infolist():
                if item.filename in skip_files:
                    continue
                data = zin.read(item.filename)
                if item.filename == "reference.conf":
                    data = modified_conf.encode("utf-8")
                info = zipfile.ZipInfo(filename=item.filename, date_time=item.date_time)
                info.compress_type = zipfile.ZIP_DEFLATED
                info.external_attr = item.external_attr
                zout.writestr(info, data)

    shutil.move(JAR_PATH + ".new", JAR_PATH)

    print("=== Step 4: Verifying ===")
    with zipfile.ZipFile(JAR_PATH, 'r') as z:
        ref = z.read("reference.conf").decode("utf-8")
        assert "pre-authentication" in ref
        sig = [n for n in z.namelist() if n.startswith("META-INF/") and n.upper().endswith(('.SF', '.RSA', '.DSA'))]
        print(f"OK Signature files: {len(sig)}")
        print(f"OK JAR size: {os.path.getsize(JAR_PATH)}")

    print("\n=== DONE ===")

if __name__ == "__main__":
    main()