#!/usr/bin/env python3
"""
VPS Gateway Patch — добавляет /provider/add и /provider/remove эндпоинты.
Запустить на VPS: python3 /opt/th_gateway/deploy_gateway_patch.py
"""
import json, os, shutil
from datetime import datetime

GATEWAY_PATH = "/opt/th_gateway/th_gateway.py"
STATE_PATH = "/opt/th_gateway/th_gateway_state.json"
BACKUP_DIR = "/opt/th_gateway/backups"

# Backup
os.makedirs(BACKUP_DIR, exist_ok=True)
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
shutil.copy(GATEWAY_PATH, f"{BACKUP_DIR}/th_gateway.py.{ts}.bak")
shutil.copy(STATE_PATH, f"{BACKUP_DIR}/th_gateway_state.json.{ts}.bak")
print(f"Backup: {BACKUP_DIR}/th_gateway.py.{ts}.bak")

with open(GATEWAY_PATH) as f:
    code = f.read()

if "STATE_FILE" not in code:
    print("ERROR: STATE_FILE not found in gateway code")
    exit(1)

old_block = '''        elif path == "/register":
            count = int(params.get("count", 5))'''

new_block = '''        elif path == "/provider/add":
            try:
                data = json.loads(body)
                provider = params.get("provider", data.get("provider", ""))
                api_key = [REDACTED]"key", "")
                label = data.get("label", "added")
                if not provider or not api_key:
                    [REDACTED]"error": "provider and key required"}, 400)
                    return
                with open(STATE_FILE) as f:
                    st = json.load(f)
                if provider not in st["providers"]:
                    self._json({"error": f"unknown provider: {provider}"}, 400)
                    return
                for k in st["providers"][provider].get("keys", []):
                    if k["key"] == api_key:
                        [REDACTED]"status": "duplicate", "key": api_key[:20] + "..."})
                        return
                st["providers"][provider]["keys"].append({
                    "key": api_key, "label": label,
                    "added": datetime.now().isoformat(), "working": True, "last_used": None
                })
                with open(STATE_FILE, "w") as f:
                    json.dump(st, f, indent=2)
                self._json({"status": "added", "provider": provider, "key": api_key[:20] + "..."})
            except Exception as e:
                self._json({"error": str(e)}, 500)

        elif path == "/provider/remove":
            try:
                data = json.loads(body)
                provider = params.get("provider", data.get("provider", ""))
                api_key = [REDACTED]"key", "")
                if not provider or not api_key:
                    [REDACTED]"error": "provider and key required"}, 400)
                    return
                with open(STATE_FILE) as f:
                    st = json.load(f)
                if provider not in st["providers"]:
                    self._json({"error": f"unknown provider: {provider}"}, 400)
                    return
                before = len(st["providers"][provider].get("keys", []))
                st["providers"][provider]["keys"] = [
                    k for k in st["providers"][provider].get("keys", [])
                    if k["key"] != api_key
                ]
                after = len(st["providers"][provider]["keys"])
                with open(STATE_FILE, "w") as f:
                    json.dump(st, f, indent=2)
                self._json({"status": "removed" if after < before else "not_found",
                            "provider": provider, "before": before, "after": after})
            except Exception as e:
                self._json({"error": str(e)}, 500)

        elif path == "/register":
            count = int(params.get("count", 5))'''

if old_block in code:
    code = code.replace(old_block, new_block)
    print("Patch 1: /provider/add + /provider/remove added")
else:
    print("Patch 1: SKIPPED — already patched or code changed")

old_root = '''    "POST /register?count=N": "start registering N accounts",'''
new_root = '''    "POST /provider/add?provider=X": "add API key to provider",
    "POST /provider/remove?provider=X": "remove API key from provider",
    "POST /register?count=N": "start registering N accounts",'''

if old_root in code:
    code = code.replace(old_root, new_root)
    print("Patch 2: root endpoint updated")
else:
    print("Patch 2: SKIPPED — already patched or code changed")

with open(GATEWAY_PATH, "w") as f:
    f.write(code)

print("Code updated. Restarting th-gateway...")
os.system("systemctl restart th-gateway")
print("Done!")