#!/usr/bin/env python3

import json
import os
import shlex
import subprocess
import sys

from openai import OpenAI


# =========================
# KONFIGURASI
# =========================

API_KEY = os.environ.get("KENARI_API_KEY")

if not API_KEY:
    print("ERROR: KENARI_API_KEY belum diset.")
    print('Contoh: export KENARI_API_KEY="kn-xxxxxxxx"')
    sys.exit(1)

BASE_URL = "https://kenari.id/v1"

# Bisa diganti dengan model Kenari lain
MODEL = os.environ.get(
    "KENARI_MODEL",
    "step-3-7-flash:free"
)

client = OpenAI(
    base_url=BASE_URL,
    api_key=API_KEY
)


# =========================
# BATASAN COMMAND
# =========================

ALLOWED_COMMANDS = {
    "pwd",
    "ls",
    "find",
    "cat",
    "head",
    "tail",
    "grep",
    "df",
    "du",
    "free",
    "uptime",
    "uname",
    "whoami",
    "id",
    "groups",
    "ps",
    "top",
    "ip",
    "ss",
    "ping",
    "which",
    "whereis",
    "file",
    "stat",
    "mount",
    "lsblk",
    "lscpu",
    "lsusb",
    "lspci",
    "systemctl",
    "journalctl",
}


# Command yang selalu meminta konfirmasi
DANGEROUS_COMMANDS = {
    "rm",
    "rmdir",
    "mv",
    "cp",
    "chmod",
    "chown",
    "kill",
    "pkill",
    "killall",
    "reboot",
    "shutdown",
    "poweroff",
    "systemctl",
    "apt",
    "apt-get",
    "dpkg",
    "pacman",
    "dnf",
    "yum",
    "mount",
    "umount",
    "mkfs",
    "dd",
    "fdisk",
    "parted",
    "wipefs",
    "format",
    "curl",
    "wget",
    "git",
}


# =========================
# VALIDASI COMMAND
# =========================

def get_base_command(command):
    try:
        parts = shlex.split(command)
        if not parts:
            return None
        return parts[0]
    except ValueError:
        return None


def is_dangerous(command):
    base = get_base_command(command)

    if base is None:
        return True

    # Tolak shell nesting
    dangerous_shell_chars = [
        "&&",
        "||",
        ";",
        "|",
        ">",
        ">>",
        "<",
        "$(",
        "`",
    ]

    for token in dangerous_shell_chars:
        if token in command:
            return True

    # sudo selalu butuh konfirmasi
    if base == "sudo":
        return True

    return base in DANGEROUS_COMMANDS


def command_allowed(command):
    base = get_base_command(command)

    if base is None:
        return False

    if base in ALLOWED_COMMANDS:
        return True

    return False


# =========================
# EKSEKUSI COMMAND
# =========================

def run_command(command):

    print()
    print("┌──────────────────────────────────────")
    print("│ COMMAND:")
    print("│", command)
    print("└──────────────────────────────────────")

    dangerous = is_dangerous(command)

    if dangerous:
        print()
        print("⚠️  COMMAND BERISIKO")
        print("Command ini membutuhkan konfirmasi.")

        answer = input("Jalankan? [y/N]: ").strip().lower()

        if answer != "y":
            return {
                "success": False,
                "output": "Command ditolak oleh pengguna."
            }

    elif not command_allowed(command):
        return {
            "success": False,
            "output": (
                "Command tidak ada dalam daftar command yang "
                "diizinkan."
            )
        }

    try:
        result = subprocess.run(
            command,
            shell=True,
            executable="/bin/bash",
            capture_output=True,
            text=True,
            timeout=30,
        )

        output = result.stdout

        if result.stderr:
            output += "\nSTDERR:\n" + result.stderr

        return {
            "success": result.returncode == 0,
            "returncode": result.returncode,
            "output": output[-12000:]
        }

    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "output": "Command timeout setelah 30 detik."
        }

    except Exception as e:
        return {
            "success": False,
            "output": f"Exception: {e}"
        }


# =========================
# TOOL DEFINITION
# =========================

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "run_linux_command",
            "description": (
                "Menjalankan command Linux pada komputer lokal. "
                "Gunakan hanya jika diperlukan untuk menjawab "
                "permintaan pengguna."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": (
                            "Satu command Linux yang akan dijalankan."
                        )
                    }
                },
                "required": ["command"],
                "additionalProperties": False
            }
        }
    }
]


# =========================
# SYSTEM PROMPT
# =========================

SYSTEM_PROMPT = """
Kamu adalah Linux system assistant.

Kamu dapat menggunakan tool run_linux_command
untuk memeriksa sistem Linux pengguna.

ATURAN:

1. Jangan menjalankan command yang tidak diperlukan.
2. Jangan menggunakan sudo.
3. Jangan mencoba melewati batasan keamanan tool.
4. Untuk operasi perubahan sistem, gunakan command yang sesuai
   tetapi tool lokal akan meminta konfirmasi pengguna.
5. Jangan mengarang output command.
6. Setelah tool memberikan hasil, jelaskan hasilnya kepada pengguna.
7. Untuk pertanyaan informasi sistem, gunakan command Linux
   jika informasi tersebut belum tersedia.
8. Jangan menghapus file atau mengubah konfigurasi hanya karena
   menurutmu itu mungkin memperbaiki masalah.
9. Jika command ditolak pengguna, jangan mencoba command lain
   untuk melakukan hal yang sama.
"""


# =========================
# AGENT LOOP
# =========================

def ask_agent(user_message):

    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT
        },
        {
            "role": "user",
            "content": user_message
        }
    ]

    for _ in range(8):

        response = client.chat.completions.create(
            model=MODEL,
            messages=messages,
            tools=TOOLS,
            tool_choice="auto",
        )

        message = response.choices[0].message

        # Tambahkan response assistant
        messages.append(message)

        # Tidak ada tool call
        if not message.tool_calls:
            print()
            print("AI:")
            print(message.content or "")
            return

        # Jalankan semua tool call
        for tool_call in message.tool_calls:

            if tool_call.function.name != "run_linux_command":
                continue

            try:
                arguments = json.loads(
                    tool_call.function.arguments
                )

                command = arguments["command"]

            except Exception as e:

                result = {
                    "success": False,
                    "output": f"Tool arguments error: {e}"
                }

                command = ""

            else:
                result = run_command(command)

            messages.append(
                {
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": json.dumps(
                        result,
                        ensure_ascii=False
                    )
                }
            )

    print("Agent berhenti: terlalu banyak langkah.")


# =========================
# CHAT LOOP
# =========================

def main():

    print("=" * 55)
    print("       KENARI LINUX SYSTEM AGENT")
    print("=" * 55)

    print()
    print("Model :", MODEL)
    print("API   :", BASE_URL)
    print()
    print("Ketik 'exit' untuk keluar.")
    print()

    while True:

        try:
            user_input = input("Kamu > ").strip()

        except (KeyboardInterrupt, EOFError):
            print()
            break

        if not user_input:
            continue

        if user_input.lower() in {
            "exit",
            "quit",
            "keluar"
        }:
            break

        try:
            ask_agent(user_input)

        except Exception as e:
            print()
            print("ERROR:")
            print(e)


if __name__ == "__main__":
    main()