import json
import os
import random
import paramiko
import telebot
import string
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Configuration and Sensitive Data
TOKEN = os.getenv("BOT_TOKEN")  # Ensure you have BOT_TOKEN in your .env file
if not TOKEN:
    raise ValueError("Bot token is not set. Please set BOT_TOKEN in your environment variables.")

bot = telebot.TeleBot(TOKEN)

# Load configuration from config.json
with open("config.json", "r") as file:
    config = json.load(file)

plan_names = list(config["plans"].keys())
server_counter = 0

# User Database
users_db = {"users": []}

# ---------------------- Utility Functions ----------------------

def load_users():
    global users_db
    try:
        with open("db.json", "r") as file:
            users_db = json.load(file)
    except FileNotFoundError:
        save_users()

def save_users():
    with open("db.json", "w") as file:
        json.dump(users_db, file, indent=4)

def add_user(user_id, username):
    if user_id not in [user["id"] for user in users_db["users"]]:
        users_db["users"].append({
            "id": user_id,
            "username": username,
            "is_admin": False,
            "servers": [],
            "balance": 0
        })
        save_users()

def set_admin(user_id, status=True):
    for user in users_db["users"]:
        if user["id"] == user_id:
            user["is_admin"] = status
            save_users()
            return True
    return False

def is_admin(user_id):
    return any(user["id"] == user_id and user.get("is_admin", False) for user in users_db["users"])

def get_unique_port():
    port = random.randint(4000, 10000)
    while port in config["rdp_ports"]:
        port = random.randint(4000, 10000)
    config["rdp_ports"].append(port)
    return port

def generate_random_password(length=12):
    characters = string.ascii_letters + string.digits + string.punctuation
    # Ensure password has at least one lowercase, uppercase, digit, and special character
    password = [
        random.choice(string.ascii_lowercase),
        random.choice(string.ascii_uppercase),
        random.choice(string.digits),
        random.choice(string.punctuation)
    ]
    password += random.choices(characters, k=length - 4)
    random.shuffle(password)
    return ''.join(password)

def create_docker_compose(plan, server_id, version):
    plan_config = config["plans"][plan]
    rdp_port = get_unique_port()
    vnc_port = get_unique_port()
    password = generate_random_password()

    docker_compose_content = f"""
version: '3'
services:
  windows:
    image: dockurr/windows
    container_name: windows{server_id}
    environment:
      VERSION: "{version}"
      DISK_SIZE: "{plan_config['DISK']}"
      RAM_SIZE: "{plan_config['RAM']}"
      CPU_CORES: "{plan_config['CPU']}"
      USERNAME: "admin"
      PASSWORD: "{password}"
    devices:
      - /dev/kvm
    cap_add:
      - NET_ADMIN
    ports:
      - "{vnc_port}:8006/tcp"
      - "{rdp_port}:3389/tcp"
      - "{rdp_port}:3389/udp"
    stop_grace_period: 2m
    restart: unless-stopped
    """

    return docker_compose_content, rdp_port, vnc_port, password

def save_docker_compose(ssh_client, docker_compose_content, server_id):
    try:
        sftp = ssh_client.open_sftp()
        folder_path = f"/root/windows{server_id}"
        try:
            sftp.mkdir(folder_path)
        except IOError:
            # Directory already exists
            pass
        with sftp.open(f"{folder_path}/docker-compose.yml", "w") as file:
            file.write(docker_compose_content)
        sftp.close()
    except Exception as e:
        print(f"Error saving docker-compose.yml: {e}")

def connect_via_ssh(server):
    ssh_client = paramiko.SSHClient()
    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        ssh_client.connect(
            server["host"],
            username=server["username"],
            password=server["password"],
            timeout=10
        )
    except Exception as e:
        print(f"SSH connection failed: {e}")
        return None
    return ssh_client

def execute_docker_command(server_id, command):
    server = config["servers"][0]
    ssh_client = connect_via_ssh(server)
    if ssh_client:
        try:
            stdin, stdout, stderr = ssh_client.exec_command(f"docker {command} windows{server_id}")
            exit_status = stdout.channel.recv_exit_status()
            if exit_status != 0:
                error = stderr.read().decode()
                print(f"Error executing docker {command}: {error}")
        except Exception as e:
            print(f"Error executing docker command: {e}")
        finally:
            ssh_client.close()

def validate_integer(value):
    try:
        return int(value)
    except ValueError:
        return None

def validate_float(value):
    try:
        return float(value)
    except ValueError:
        return None

# ---------------------- Bot Command Handlers ----------------------

@bot.message_handler(commands=["start"])
def start_command(message):
    add_user(message.chat.id, message.from_user.username)
    bot.send_message(message.chat.id, "👋 Привет! Используйте /help для списка доступных команд.")

@bot.message_handler(commands=["help"])
def help_command(message):
    help_text = (
        "🛠️ **Доступные команды:**\n"
        "/shop - Приобрести сервер\n"
        "/admin - Админ-панель (только для админов)\n"
        "/menu - Меню управления серверами\n"
        "/balance - Проверить баланс"
    )
    bot.send_message(message.chat.id, help_text, parse_mode="Markdown")

@bot.message_handler(commands=["balance"])
def balance_command(message):
    user = next((user for user in users_db["users"] if user["id"] == message.chat.id), None)
    if user:
        bot.send_message(message.chat.id, f"💵 Ваш баланс: ${user['balance']}")
    else:
        bot.send_message(message.chat.id, "🚫 У вас нет баланса.")

@bot.message_handler(commands=["shop"])
def shop_command(message):
    markup = telebot.types.InlineKeyboardMarkup()
    for plan in plan_names:
        markup.add(telebot.types.InlineKeyboardButton(text=plan.capitalize(), callback_data=f"plan:{plan}"))
    bot.send_message(message.chat.id, "🛒 Выберите тарифный план:", reply_markup=markup)

@bot.message_handler(commands=["admin"])
def admin_command(message):
    if is_admin(message.chat.id):
        admin_text = (
            "🔧 **Админ-панель:**\n"
            "/setadmin [ID] - Назначить пользователя админом\n"
            "/removeadmin [ID] - Удалить пользователя из админов\n"
            "/listusers - Показать всех пользователей\n"
            "/givebalance [ID] [AMOUNT] - Выдать баланс пользователю"
        )
        bot.send_message(message.chat.id, admin_text, parse_mode="Markdown")
    else:
        bot.send_message(message.chat.id, "🚫 У вас нет доступа к админ-панели")

@bot.message_handler(commands=["setadmin", "removeadmin"])
def set_admin_command(message):
    if is_admin(message.chat.id):
        parts = message.text.split()
        if len(parts) != 2:
            bot.send_message(message.chat.id, "⚠️ Использование: /setadmin [ID] или /removeadmin [ID]")
            return
        command, user_id_str = parts
        user_id = validate_integer(user_id_str)
        if user_id is None:
            bot.send_message(message.chat.id, "⚠️ ID пользователя должно быть числом.")
            return
        if command == "/setadmin":
            success = set_admin(user_id, True)
            if success:
                bot.send_message(message.chat.id, f"✅ Пользователь {user_id} назначен администратором.")
            else:
                bot.send_message(message.chat.id, "🚫 Пользователь не найден.")
        elif command == "/removeadmin":
            success = set_admin(user_id, False)
            if success:
                bot.send_message(message.chat.id, f"✅ Пользователь {user_id} удален из администраторов.")
            else:
                bot.send_message(message.chat.id, "🚫 Пользователь не найден.")
    else:
        bot.send_message(message.chat.id, "🚫 У вас нет прав.")

@bot.message_handler(commands=["givebalance"])
def give_balance_command(message):
    if is_admin(message.chat.id):
        parts = message.text.split()
        if len(parts) != 3:
            bot.send_message(message.chat.id, "⚠️ Использование: /givebalance [ID] [AMOUNT]")
            return
        _, user_id_str, amount_str = parts
        user_id = validate_integer(user_id_str)
        amount = validate_float(amount_str)
        if user_id is None or amount is None:
            bot.send_message(message.chat.id, "⚠️ ID должно быть числом и сумма должна быть числом.")
            return
        user = next((user for user in users_db["users"] if user["id"] == user_id), None)
        if user:
            user["balance"] += amount
            save_users()
            bot.send_message(message.chat.id, f"✅ Выдано ${amount} пользователю {user_id}.")
        else:
            bot.send_message(message.chat.id, "🚫 Пользователь не найден.")
    else:
        bot.send_message(message.chat.id, "🚫 У вас нет прав.")

@bot.message_handler(commands=["menu"])
def menu_command(message):
    user = next((user for user in users_db["users"] if user["id"] == message.chat.id), None)
    if user:
        if user["servers"]:
            markup = telebot.types.InlineKeyboardMarkup()
            for server in user["servers"]:
                markup.add(telebot.types.InlineKeyboardButton(
                    text=f"🖥️ Сервер {server['id']}",
                    callback_data=f"control:{server['id']}"
                ))
            bot.send_message(message.chat.id, "🖥️ Ваши серверы:", reply_markup=markup)
        else:
            bot.send_message(message.chat.id, "🚫 У вас нет серверов.")
    else:
        bot.send_message(message.chat.id, "🚫 У вас нет доступа.")

# ---------------------- Callback Query Handlers ----------------------

@bot.callback_query_handler(func=lambda call: call.data.startswith("plan:"))
def plan_callback(call):
    global server_counter
    plan = call.data.split(":")[1]

    user = next((user for user in users_db["users"] if user["id"] == call.message.chat.id), None)
    if user and user["balance"] >= config["plans"][plan]["price"]:
        server_counter += 1
        version_markup = telebot.types.InlineKeyboardMarkup()
        versions = ["2022", "2025", "2019", "win10", "win11"]
        for version in versions:
            if version.startswith("win"):
                version_display = f"Windows {version[-2:].upper()}"
            else:
                version_display = f"Windows Server {version}"
            # Include 'plan' in the callback_data for each version button
            version_markup.add(
                telebot.types.InlineKeyboardButton(
                    text=version_display,
                    callback_data=f"version:{version}:plan:{plan}"
                )
            )
        bot.send_message(call.message.chat.id, "📋 **Выберите версию Windows:**", reply_markup=version_markup, parse_mode="Markdown")
    else:
        bot.send_message(call.message.chat.id, "🚫 У вас недостаточно средств для покупки этого тарифа.")

@bot.callback_query_handler(func=lambda call: call.data.startswith("version:"))
def version_callback(call):
    data_parts = call.data.split(":")
    if len(data_parts) != 4:
        bot.send_message(call.message.chat.id, "⚠️ Некорректные данные. Попробуйте снова.")
        return

    _, version, _, plan = data_parts

    global server_counter
    server_id = server_counter

    # Create Docker compose content
    docker_compose_content, rdp_port, vnc_port, password = create_docker_compose(plan, server_id, version)

    # Save docker-compose.yml to the server
    server = config["servers"][0]
    ssh_client = connect_via_ssh(server)
    if not ssh_client:
        bot.send_message(call.message.chat.id, "🚫 Не удалось подключиться к серверу. Пожалуйста, попробуйте позже.")
        return

    save_docker_compose(ssh_client, docker_compose_content, server_id)

    # Start the Docker container
    try:
        ssh_client.exec_command(f"cd /root/windows{server_id} && docker compose up -d")
    except Exception as e:
        print(f"Error starting Docker container: {e}")
        bot.send_message(call.message.chat.id, "⚠️ Не удалось запустить контейнер. Пожалуйста, попробуйте позже.")
        ssh_client.close()
        return

    ssh_client.close()

    # Update user data
    user = next((user for user in users_db["users"] if user["id"] == call.message.chat.id), None)
    if user:
        user["servers"].append({
            "id": server_id,
            "plan": plan,
            "status": "Запущен",
            "rdp_port": rdp_port,
            "vnc_port": vnc_port,
            "password": password
        })

        # Deduct the price from the user's balance
        user["balance"] -= config["plans"][plan]["price"]
        save_users()

        # Send server details to the user
        bot.send_message(
            call.message.chat.id,
            f"✅ **Сервер создан!**\n\n"
            f"🌐 **RDP IP:** {server['host']}:{rdp_port}\n"
            f"🌐 **VNC:** [http://{server['host']}:{vnc_port}](http://{server['host']}:{vnc_port})\n"
            f"🔑 **Пароль:** `{password}`",
            parse_mode="Markdown"
        )
    else:
        bot.send_message(call.message.chat.id, "🚫 Пользователь не найден.")

@bot.callback_query_handler(func=lambda call: call.data.startswith("control:"))
def control_callback(call):
    server_id_str = call.data.split(":")[1]
    server_id = validate_integer(server_id_str)
    if server_id is None:
        bot.send_message(call.message.chat.id, "⚠️ Некорректный ID сервера.")
        return

    user = next((user for user in users_db["users"] if user["id"] == call.message.chat.id), None)
    if user:
        server = next((s for s in user["servers"] if s["id"] == server_id), None)
        if server:
            server_status = f"**Статус:** {server['status']}\n"
            control_markup = telebot.types.InlineKeyboardMarkup()
            control_markup.row(
                telebot.types.InlineKeyboardButton(text="🛑 Остановить", callback_data=f"stop:{server_id}"),
                telebot.types.InlineKeyboardButton(text="🔄 Перезапустить", callback_data=f"restart:{server_id}"),
                telebot.types.InlineKeyboardButton(text="▶️ Запустить", callback_data=f"start:{server_id}")
            )
            control_markup.add(
                telebot.types.InlineKeyboardButton(text="❌ Удалить", callback_data=f"delete:{server_id}")
            )
            bot.send_message(
                call.message.chat.id,
                f"🔧 **Управление сервером {server_id}:**\n{server_status}",
                reply_markup=control_markup,
                parse_mode="Markdown"
            )
        else:
            bot.send_message(call.message.chat.id, "🚫 Сервер не найден.")
    else:
        bot.send_message(call.message.chat.id, "🚫 У вас нет доступа к этому серверу.")

@bot.callback_query_handler(func=lambda call: call.data.startswith("start:"))
def start_server_callback(call):
    server_id_str = call.data.split(":")[1]
    server_id = validate_integer(server_id_str)
    if server_id is None:
        bot.send_message(call.message.chat.id, "⚠️ Некорректный ID сервера.")
        return

    execute_docker_command(server_id, "start")

    user = next((user for user in users_db["users"] if user["id"] == call.message.chat.id), None)
    if user:
        server = next((s for s in user["servers"] if s["id"] == server_id), None)
        if server:
            server["status"] = "Запущен"
            save_users()
            bot.send_message(call.message.chat.id, f"✅ Сервер {server_id} запущен.")
        else:
            bot.send_message(call.message.chat.id, "🚫 Сервер не найден.")
    else:
        bot.send_message(call.message.chat.id, "🚫 Пользователь не найден.")

@bot.callback_query_handler(func=lambda call: call.data.startswith("stop:"))
def stop_server_callback(call):
    server_id_str = call.data.split(":")[1]
    server_id = validate_integer(server_id_str)
    if server_id is None:
        bot.send_message(call.message.chat.id, "⚠️ Некорректный ID сервера.")
        return

    execute_docker_command(server_id, "stop")

    user = next((user for user in users_db["users"] if user["id"] == call.message.chat.id), None)
    if user:
        server = next((s for s in user["servers"] if s["id"] == server_id), None)
        if server:
            server["status"] = "Остановлен"
            save_users()
            bot.send_message(call.message.chat.id, f"🛑 Сервер {server_id} остановлен.")
        else:
            bot.send_message(call.message.chat.id, "🚫 Сервер не найден.")
    else:
        bot.send_message(call.message.chat.id, "🚫 Пользователь не найден.")

@bot.callback_query_handler(func=lambda call: call.data.startswith("restart:"))
def restart_server_callback(call):
    server_id_str = call.data.split(":")[1]
    server_id = validate_integer(server_id_str)
    if server_id is None:
        bot.send_message(call.message.chat.id, "⚠️ Некорректный ID сервера.")
        return

    execute_docker_command(server_id, "restart")

    user = next((user for user in users_db["users"] if user["id"] == call.message.chat.id), None)
    if user:
        server = next((s for s in user["servers"] if s["id"] == server_id), None)
        if server:
            server["status"] = "Запущен"
            save_users()
            bot.send_message(call.message.chat.id, f"🔄 Сервер {server_id} перезапущен.")
        else:
            bot.send_message(call.message.chat.id, "🚫 Сервер не найден.")
    else:
        bot.send_message(call.message.chat.id, "🚫 Пользователь не найден.")

@bot.callback_query_handler(func=lambda call: call.data.startswith("delete:"))
def delete_server_callback(call):
    server_id_str = call.data.split(":")[1]
    server_id = validate_integer(server_id_str)
    if server_id is None:
        bot.send_message(call.message.chat.id, "⚠️ Некорректный ID сервера.")
        return

    user = next((user for user in users_db["users"] if user["id"] == call.message.chat.id), None)

    if user:
        server = next((s for s in user["servers"] if s["id"] == server_id), None)
        if server:
            markup = telebot.types.InlineKeyboardMarkup()
            markup.add(
                telebot.types.InlineKeyboardButton(text="✅ Да", callback_data=f"confirm_delete:{server_id}"),
                telebot.types.InlineKeyboardButton(text="❌ Нет", callback_data="cancel_delete")
            )
            bot.send_message(
                call.message.chat.id,
                f"❗ Вы уверены, что хотите удалить сервер {server_id}?",
                reply_markup=markup
            )
        else:
            bot.send_message(call.message.chat.id, "🚫 Сервер не найден.")
    else:
        bot.send_message(call.message.chat.id, "🚫 У вас нет доступа к этому серверу.")

@bot.callback_query_handler(func=lambda call: call.data.startswith("confirm_delete:"))
def confirm_delete_callback(call):
    server_id_str = call.data.split(":")[1]
    server_id = validate_integer(server_id_str)
    if server_id is None:
        bot.send_message(call.message.chat.id, "⚠️ Некорректный ID сервера.")
        return

    user = next((user for user in users_db["users"] if user["id"] == call.message.chat.id), None)

    if user:
        server = next((s for s in user["servers"] if s["id"] == server_id), None)
        if server:
            # Stop and remove the Docker container
            execute_docker_command(server_id, "stop")
            execute_docker_command(server_id, "rm")
            # Optionally, remove docker-compose files from the server
            server_info = config["servers"][0]
            ssh_client = connect_via_ssh(server_info)
            if ssh_client:
                try:
                    ssh_client.exec_command(f"rm -rf /root/windows{server_id}")
                except Exception as e:
                    print(f"Failed to remove server files: {e}")
                finally:
                    ssh_client.close()

            # Remove server from user's list
            user["servers"] = [s for s in user["servers"] if s["id"] != server_id]
            save_users()
            bot.send_message(call.message.chat.id, f"✅ Сервер {server_id} успешно удален.")
        else:
            bot.send_message(call.message.chat.id, "🚫 Сервер не найден.")
    else:
        bot.send_message(call.message.chat.id, "🚫 Пользователь не найден.")

@bot.callback_query_handler(func=lambda call: call.data == "cancel_delete")
def cancel_delete_callback(call):
    bot.send_message(call.message.chat.id, "❌ Удаление сервера отменено.")

# ---------------------- Start Bot ----------------------

def main():
    load_users()
    print("Bot is polling...")
    try:
        bot.polling(none_stop=True)
    except Exception as e:
        print(f"Bot polling failed: {e}")

if __name__ == "__main__":
    main()

check, if bot automatically starts dockervm of each version of server