import os
import asyncio
import sqlite3
import re
from datetime import datetime
from pathlib import Path
import subprocess
import openai
from google import genai
from google.genai import types
from dotenv import load_dotenv

# --- FARBEN FÜR DAS TERMINAL ---
class Color:
    USER = '\033[94m'      # Blau
    JARVIS = '\033[92m'    # Grün
    SYSTEM = '\033[93m'    # Gelb
    EXEC = '\033[95m'      # Magenta für SSH
    ERROR = '\033[91m'     # Rot
    RESET = '\033[0m'      # Zurücksetzen
    BOLD = '\033[1m'

# --- PFADE & SETUP ---
BASE_DIR = Path(__file__).resolve().parent
ROOT_DIR = BASE_DIR.parent
CONFIG_DIR = ROOT_DIR / "config"
DATA_DIR = ROOT_DIR / "data"
WORKSPACE_DIR = ROOT_DIR / "workspace"

ENV_FILE = CONFIG_DIR / ".env"
load_dotenv(ENV_FILE)

DB_PATH = DATA_DIR / "cluster.db"
PROMPT_FILE = CONFIG_DIR / "system_prompt.txt"
NOTES_FILE = WORKSPACE_DIR / "NOTIZEN.md"
TODO_FILE = WORKSPACE_DIR / "TODO.md"

WEB_USER_NAME = os.getenv("WEB_USER_NAME", "Admin")

# Ordner & Dateien anlegen
for d in [WORKSPACE_DIR, DATA_DIR, CONFIG_DIR]:
    d.mkdir(parents=True, exist_ok=True)

for f in [NOTES_FILE, TODO_FILE]:
    if not f.exists():
        f.write_text(f"# {f.name}\nHier fängt dein Gedächtnis an, J.A.R.V.I.S.\n", encoding="utf-8")

# --- KI KONFIGURATION ---
AI_PROVIDER = os.getenv("AI_PROVIDER", "google").lower()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "")
NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY", "")
OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://127.0.0.1:11434/v1")

GOOGLE_MODEL = os.getenv("GOOGLE_MODEL", "gemini-2.0-flash")
OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o")
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama3")
NVIDIA_MODEL = os.getenv("NVIDIA_MODEL", "moonshotai/kimi-k2.5")

# --- DATENBANK ---
def init_db():
    conn = sqlite3.connect(DB_PATH)
    conn.execute('''
        CREATE TABLE IF NOT EXISTS nodes (
            id INTEGER PRIMARY KEY AUTOINCREMENT, 
            name TEXT, ip TEXT UNIQUE, user TEXT, sudo_password TEXT,
            os TEXT DEFAULT 'Unbekannt', arch TEXT DEFAULT 'Unbekannt',
            docker_installed INTEGER DEFAULT 0, status TEXT
        )
    ''')
    conn.commit()
    conn.close()

init_db()

def get_db():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn

def get_system_prompt(current_user=WEB_USER_NAME):
    prompt_path = CONFIG_DIR / "system_prompt.txt"
    if prompt_path.exists():
        prompt = prompt_path.read_text(encoding="utf-8")
    else:
        prompt = f"Hallo {current_user}, ich bin J.A.R.V.I.S., dein lokaler Systemassistent."

    prompt = prompt.replace("{user_name}", current_user)

    conn = get_db()
    nodes = conn.execute('SELECT * FROM nodes').fetchall()
    conn.close()
    
    node_info = ""
    for n in nodes:
        docker_str = "Ja" if n['docker_installed'] else "Nein"
        node_info += f"- Name: {n['name']}, IP: {n['ip']}, User: {n['user']}, OS: {n['os']}, Arch: {n['arch']}, Docker: {docker_str}\n"
    
    prompt = prompt.replace("{node_info}", node_info)
    prompt = prompt.replace("{workspace_dir}", str(WORKSPACE_DIR))
    prompt = prompt.replace("{notes_file}", str(NOTES_FILE))
    prompt = prompt.replace("{todo_file}", str(TODO_FILE))
    
    return prompt

# --- KI KOMMUNIKATION ---
async def get_ai_response(user_msg, system_prompt, history_list):
    try:
        if AI_PROVIDER in ["openai", "ollama", "nvidia"]:
            messages = [{"role": "system", "content": system_prompt}] + history_list
            
            if AI_PROVIDER == "ollama":
                url = OLLAMA_BASE_URL if OLLAMA_BASE_URL.endswith('/v1') else OLLAMA_BASE_URL.rstrip('/') + '/v1'
                key, model_to_use = "ollama", OLLAMA_MODEL
            elif AI_PROVIDER == "nvidia":
                url, key, model_to_use = "https://integrate.api.nvidia.com/v1", NVIDIA_API_KEY, NVIDIA_MODEL
            else:
                url, key, model_to_use = None, OPENAI_API_KEY, OPENAI_MODEL

            client = openai.AsyncOpenAI(base_url=url, api_key=key)
            response = await client.chat.completions.create(model=model_to_use, messages=messages)
            return response.choices[0].message.content
            
        elif AI_PROVIDER == "google":
            if not GOOGLE_API_KEY: return "Fehler: GOOGLE_API_KEY fehlt!"
            client = genai.Client(api_key=GOOGLE_API_KEY)
            
            google_history = [
                types.Content(role="user" if msg["role"] == "user" else "model", parts=[types.Part.from_text(text=msg["content"])])
                for msg in history_list[:-1]
            ]
            
            chat = client.chats.create(
                model=GOOGLE_MODEL,
                config=types.GenerateContentConfig(system_instruction=system_prompt),
                history=google_history
            )
            return chat.send_message(user_msg).text
            
    except Exception as e:
        return f"{Color.ERROR}Fehler bei der KI-Anfrage: {e}{Color.RESET}"

# --- SSH BEFEHLSAUSFÜHRUNG ---
async def run_remote_task(target, cmd):
    print(f"\n{Color.EXEC}🚀 J.A.R.V.I.S. führt aus auf {target}:{Color.RESET}\n> {cmd}\n")
    
    conn = get_db()
    n = conn.execute('SELECT * FROM nodes WHERE ip=? OR name=?', (target, target)).fetchone()
    conn.close()

    if not n:
        msg = f"{Color.ERROR}⚠️ Node '{target}' nicht in der Datenbank gefunden.{Color.RESET}"
        print(msg)
        return msg

    try:
        proc = await asyncio.create_subprocess_shell(
            f"ssh -o StrictHostKeyChecking=no -o LogLevel=ERROR {n['user']}@{n['ip']} '{cmd}'",
            stdout=asyncio.subprocess.PIPE, 
            stderr=asyncio.subprocess.STDOUT
        )
        stdout, _ = await proc.communicate()
        output = stdout.decode('utf-8', errors='ignore').strip()
        
        print(f"{Color.EXEC}💻 Output:\n{output or '✅ Ohne Output ausgeführt.'}{Color.RESET}\n")
        return output
    except Exception as e:
        err = f"{Color.ERROR}❌ Fehler: {e}{Color.RESET}"
        print(err)
        return err

# --- HAUPT-TERMINAL-LOOP ---
async def main_chat_loop():
    # Bunter Boot-Screen
    print(f"{Color.SYSTEM}===================================================={Color.RESET}")
    print(f"{Color.BOLD}{Color.JARVIS}🤖 J.A.R.V.I.S. Terminal Interface geladen.{Color.RESET}")
    print(f"{Color.SYSTEM}Provider: {AI_PROVIDER.upper()}")
    print(f"Tippe 'exit' oder 'quit' um zu beenden.{Color.RESET}")
    print(f"{Color.SYSTEM}===================================================={Color.RESET}\n")
    
    chat_history = []
    system_prompt = get_system_prompt()

    while True:
        # User Eingabe abwarten (Blau und Fett)
        user_msg = await asyncio.to_thread(input, f"\n{Color.BOLD}{Color.USER}Du: {Color.RESET}")
        
        if user_msg.lower().strip() in ['exit', 'quit']:
            print(f"{Color.SYSTEM}J.A.R.V.I.S. wird beendet. Auf Wiedersehen!{Color.RESET}")
            break
        if not user_msg.strip():
            continue

        now = datetime.now().strftime("%d.%m.%Y %H:%M")
        chat_history.append({"role": "user", "content": user_msg, "timestamp": now})

        # Ladeindikator
        print(f"{Color.SYSTEM}J.A.R.V.I.S. verarbeitet...{Color.RESET}", end="\r")
        
        # KI antworten lassen
        ai_response = await get_ai_response(user_msg, system_prompt, chat_history)
        
        # Nach <EXECUTE> Tags suchen
        commands = re.findall(r'<EXECUTE target="(.*?)">(.*?)</EXECUTE>', ai_response, re.I | re.S)
        clean_msg = re.sub(r'<EXECUTE.*?>.*?</EXECUTE>', '', ai_response, flags=re.I | re.S).strip()

        # Ladezeile säubern und KI-Antwort ausgeben (Grün)
        print("\r" + " " * 30 + "\r", end="") 
        print(f"{Color.BOLD}{Color.JARVIS}🤖 J.A.R.V.I.S.:{Color.RESET}")
        print(f"{Color.JARVIS}{clean_msg}{Color.RESET}")

        chat_history.append({"role": "assistant", "content": clean_msg, "timestamp": now})

        # Ausführung der gefundenen Befehle
        if commands:
            for target, cmd in commands:
                output = await run_remote_task(target.strip(), cmd.strip())
                # System-Output ebenfalls in die Historie packen
                sys_now = datetime.now().strftime("%d.%m.%Y %H:%M")
                chat_history.append({
                    "role": "user", 
                    "content": f"[SYSTEM] Befehl '{cmd}' auf {target} abgeschlossen. Output:\n{output}", 
                    "timestamp": sys_now
                })

        # RAM-Schutz: Historie nicht ins Unendliche wachsen lassen
        if len(chat_history) > 20:
            chat_history = chat_history[-20:]

if __name__ == "__main__":
    try:
        asyncio.run(main_chat_loop())
    except KeyboardInterrupt:
        print(f"\n{Color.ERROR}J.A.R.V.I.S. hart beendet.{Color.RESET}")