#!/usr/bin/env python3
import json
import os
import subprocess
import sys
import time
import urllib.parse
import urllib.request

API = "https://api.telegram.org/bot" + os.environ["EMERGENCY_BOT_TOKEN"]
ALLOWED_USERS = {item.strip() for item in os.environ["EMERGENCY_ALLOWED_USER_IDS"].split(",") if item.strip()}
STATE_FILE = "{{ emergency_bot_state_dir }}/updates.json"
BUTTONS = {
    "inline_keyboard": [[
        {"text": "Enable SSH", "callback_data": "start"},
        {"text": "Status", "callback_data": "status"},
        {"text": "Stop", "callback_data": "stop"},
    ]]
}


HTTP = urllib.request.build_opener(
    urllib.request.ProxyHandler({"https": os.environ["EMERGENCY_TELEGRAM_PROXY"]})
)


def api(method, payload=None, timeout=35):
    if payload is not None:
        payload = {
            key: json.dumps(value) if isinstance(value, (dict, list)) else value
            for key, value in payload.items()
        }
    data = None if payload is None else urllib.parse.urlencode(payload).encode()
    with HTTP.open(API + "/" + method, data=data, timeout=timeout) as response:
        result = json.load(response)
    if not result.get("ok"):
        raise RuntimeError(result.get("description", "Telegram API request failed"))
    return result["result"]


def control(command):
    args = [
        "/usr/bin/ssh", "-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=yes",
        "-o", "HostKeyAlias=mini-pc-emergency", "-o", "UserKnownHostsFile=" + os.environ["EMERGENCY_KNOWN_HOSTS"],
        "-i", os.environ["EMERGENCY_CONTROL_KEY"],
        os.environ["EMERGENCY_MINI_PC_USER"] + "@" + os.environ["EMERGENCY_MINI_PC_HOST"], command,
    ]
    result = subprocess.run(args, capture_output=True, text=True, timeout=25, check=False)
    output = (result.stdout or result.stderr).strip()
    if result.returncode:
        return "Control command failed: " + (output or "unknown error")[:1000]
    return output or "ok"


def reply(chat_id, text):
    payload = {"chat_id": chat_id, "text": text[:3500], "reply_markup": BUTTONS}
    if text.startswith("started; expires in") and "\nConnect:\n" in text:
        status, command = text.split("\nConnect:\n", 1)
        payload["text"] = status + "\nConnect:\n`" + command + "`"
        payload["parse_mode"] = "Markdown"
    api("sendMessage", payload)


def allowed(chat, sender):
    return chat.get("type") == "private" and str(sender.get("id")) in ALLOWED_USERS


def handle_command(chat_id, command):
    if command == "help":
        reply(chat_id, "Choose an action or use /emergency ssh, /emergency status, /emergency stop")
    elif command:
        print("accepted emergency command: " + command, file=sys.stderr, flush=True)
        reply(chat_id, control(command))


def load_offset():
    try:
        with open(STATE_FILE, encoding="utf-8") as state_file:
            return int(json.load(state_file).get("offset", 0))
    except (FileNotFoundError, ValueError, json.JSONDecodeError):
        return 0


def save_offset(offset):
    temporary = STATE_FILE + ".tmp"
    with open(temporary, "w", encoding="utf-8") as state_file:
        json.dump({"offset": offset}, state_file)
    os.replace(temporary, STATE_FILE)


def command_for(text):
    parts = text.strip().split()
    if not parts:
        return None
    command_name = parts[0].split("@", 1)[0]
    if command_name == "/start":
        return "help"
    if command_name != "/emergency":
        return None
    if len(parts) != 2:
        return "help"
    return {"ssh": "start", "stop": "stop", "status": "status"}.get(parts[1], "help")


def run():
    offset = load_offset()
    while True:
        try:
            for update in api("getUpdates", {"offset": offset, "timeout": 30}, timeout=40):
                offset = update["update_id"] + 1
                save_offset(offset)
                callback = update.get("callback_query")
                if callback:
                    callback_message = callback.get("message", {})
                    chat = callback_message.get("chat", {})
                    sender = callback.get("from", {})
                    command = callback.get("data")
                    if not allowed(chat, sender) or command not in {"start", "status", "stop"}:
                        print("ignored callback: chat_type=" + str(chat.get("type")) + " sender=" + str(sender.get("id")), file=sys.stderr, flush=True)
                        api("answerCallbackQuery", {"callback_query_id": callback["id"], "text": "Not authorized", "show_alert": True})
                    else:
                        api("answerCallbackQuery", {"callback_query_id": callback["id"]})
                        handle_command(chat["id"], command)
                    continue
                message = update.get("message", {})
                chat = message.get("chat", {})
                sender = message.get("from", {})
                text = message.get("text", "")
                if not allowed(chat, sender):
                    print("ignored update: chat_type=" + str(chat.get("type")) + " sender=" + str(sender.get("id")), file=sys.stderr, flush=True)
                    continue
                command = command_for(text)
                handle_command(chat["id"], command)
        except Exception as error:
            print("emergency-bot error: " + str(error), file=sys.stderr, flush=True)
            time.sleep(10)


if __name__ == "__main__":
    run()