import os
import pty
import fcntl
import subprocess
import sqlite3
import asyncio
import openai
import re
from google import genai
from google.genai import types
import json
from fastapi import FastAPI, WebSocket, BackgroundTasks, Request, Form, WebSocketDisconnect
from fastapi.responses import RedirectResponse
from fastapi.templating import Jinja2Templates
from fastapi.staticfiles import StaticFiles
from dotenv import load_dotenv

# Lade Umgebungsvariablen aus der .env Datei
load_dotenv()

app = FastAPI()
static_path = os.path.join(os.path.dirname(__file__), "static")
app.mount("/static", StaticFiles(directory=static_path), name="static")
templates = Jinja2Templates(directory="templates")

SSH_KEY = os.path.expanduser("~/.ssh/id_rsa")
DB_PATH = "cluster.db"
chat_history = []

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

# Modelle aus .env laden (mit Standardwerten als Fallback)
GOOGLE_MODEL = os.getenv("GOOGLE_MODEL", "gemini-2.5-flash")
OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o")
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama3")

def get_system_prompt():
    conn = get_db()
    nodes = conn.execute('SELECT * FROM nodes').fetchall()
    conn.close()
    
    node_info = ""
    for n in nodes:
        node_info += f"- Name: {n['name']}, IP: {n['ip']}, User: {n['user']}\n"   
    return f"""Du bist der Pi-Orchestrator KI-Assistent. Deine Aufgabe ist es, Befehle auf Raspberry Pis auszuführen. Du KANNST und SOLLST Befehle ausführen! Hier sind die aktuell verbundenen Nodes: {node_info} . WENN der Nutzer dich bittet, etwas zu tun (z.B. ping, update, docker installieren), dann formuliere erst eine kurze Antwort und frage nochmal nach ob du dies dann auf dem gewünschten node durchführen sollst. Du prüfst vorher noch, ob du auf dem gewünschten node sudo rechte ohne eingäbe eines passwortes hast, sofern diese notwendig sind. Erst nach einer positiven Bestätigung darfst du es ausführen und fügst am Ende die Befehle in genau diesem XML-Format hinzu: <EXECUTE target="IP_ADRESSE">befehl</EXECUTE> .WICHTIG: Verwende als target IMMER die IP-Adresse des Nodes. Bei Befehlen wie 'ping' oder 'top', die nicht enden, MUSS ein Limit gesetzt werden (z.B. ping -c 4 IP).Beispielantwort für 'sende einen ping an pi-06': Ich starte den Ping an pi-06 für dich. <EXECUTE target="192.168.1.10">ping -c 4 192.168.1.10</EXECUTE>"""

# --- KI FUNKTIONEN ---

async def get_ai_response(user_input, system_prompt):
    global chat_history
    
    # 1. Die neue User-Nachricht dem Gedächtnis hinzufügen
    chat_history.append({"role": "user", "content": user_input})
    
    # 2. Gedächtnis auf die letzten 30 Nachrichten begrenzen
    chat_history = chat_history[-30:]
    
    ai_msg = ""
    
    try:
        if AI_PROVIDER == "openai" or AI_PROVIDER == "ollama":
            # Für OpenAI / Ollama
            messages = [{"role": "system", "content": system_prompt}] + chat_history
            url = OLLAMA_BASE_URL if AI_PROVIDER == "ollama" else None
            key = "ollama" if AI_PROVIDER == "ollama" else OPENAI_API_KEY
            model_to_use = OLLAMA_MODEL if AI_PROVIDER == "ollama" else OPENAI_MODEL
            
            client = openai.OpenAI(base_url=url, api_key=key)
            response = client.chat.completions.create(
                model=model_to_use,
                messages=messages
            )
            ai_msg = response.choices[0].message.content
            
        elif AI_PROVIDER == "google":
            # Für Google Gemini
            if not GOOGLE_API_KEY:
                return "Fehler: GOOGLE_API_KEY fehlt in der .env Datei!"
                
            client = genai.Client(api_key=GOOGLE_API_KEY)
            
            # Wir müssen unser Array in das spezielle Google-Format umwandeln
            google_history = []
            
            # Alle Nachrichten AUSSER der allerletzten (die aktuelle User-Frage) in die History packen
            for msg in chat_history[:-1]:
                role = "user" if msg["role"] == "user" else "model"
                google_history.append(
                    types.Content(role=role, parts=[types.Part.from_text(text=msg["content"])])
                )
            
            # Chat MIT dem übersetzten Gedächtnis starten
            chat = client.chats.create(
                model=GOOGLE_MODEL,
                config=types.GenerateContentConfig(system_instruction=system_prompt),
                history=google_history
            )
            
            # Jetzt erst die neue Nachricht an den Chat mit Gedächtnis schicken
            response = chat.send_message(user_input)
            ai_msg = response.text
            
    except Exception as e:
        ai_msg = f"Fehler bei der KI-Anfrage: {e}"
        print(f"KI Fehler: {e}")

    # 3. Die Antwort der KI ebenfalls ins Gedächtnis aufnehmen
    chat_history.append({"role": "assistant", "content": ai_msg})
    
    return ai_msg

# --- DATENBANK INITIALISIERUNG ---
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, 
            status TEXT
        )
    ''')
    conn.commit()
    conn.close()

init_db()

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

# --- WebSocket Manager für Logs & Chat ---
class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []
    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)
    def disconnect(self, websocket: WebSocket):
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)
    async def broadcast(self, message: str):
        for connection in self.active_connections:
            try:
                await connection.send_text(message)
            except:
                pass

manager = ConnectionManager()

# --- SSH Handshake (Nur Key kopieren) ---
async def bootstrap_ssh_only(ip, user, password):
    await manager.broadcast(f"🔑 Initialisiere SSH-Handshake für {ip}...")
    # Nutzt sshpass um den Key einmalig mit Passwort zu hinterlegen
    ssh_copy_cmd = f"sshpass -p '{password}' ssh-copy-id -o StrictHostKeyChecking=no -i {SSH_KEY}.pub {user}@{ip}"
    
    process = subprocess.Popen(ssh_copy_cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    for line in process.stdout:
        await manager.broadcast(f"SSH: {line.strip()}")
    
    conn = get_db()
    conn.execute('UPDATE nodes SET status = "Bereit (Kein Docker)" WHERE ip = ?', (ip,))
    conn.commit()
    conn.close()
    await manager.broadcast(f"✅ Node {ip} ist verbunden. KI kann nun Befehle senden.")

# --- Routen ---

@app.get("/")
async def index(request: Request):
    conn = get_db()
    nodes = conn.execute('SELECT * FROM nodes').fetchall()
    conn.close()
    return templates.TemplateResponse("index.html", {"request": request, "nodes": nodes})

@app.post("/add_node")
async def add_node(background_tasks: BackgroundTasks, name: str = Form(...), ip: str = Form(...), user: str = Form(...), password: str = Form(...)):
    conn = get_db()
    try:
        conn.execute('INSERT INTO nodes (name, ip, user, status) VALUES (?, ?, ?, ?)', (name, ip, user, "Kopplung..."))
        conn.commit()
        background_tasks.add_task(bootstrap_ssh_only, ip, user, password)
    except sqlite3.IntegrityError: pass
    finally: conn.close()
    return RedirectResponse(url="/", status_code=303)

@app.post("/remove_node/{node_id}")
async def remove_node(node_id: int):
    conn = get_db()
    conn.execute('DELETE FROM nodes WHERE id = ?', (node_id,))
    conn.commit()
    conn.close()
    return RedirectResponse(url="/", status_code=303)

# --- WebSockets ---

@app.websocket("/ws/install_logs")
async def log_websocket(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        manager.disconnect(websocket)

@app.websocket("/ws/terminal/{ip}")
async def terminal_websocket(websocket: WebSocket, ip: str):
    await websocket.accept()
    conn = get_db()
    node = conn.execute('SELECT * FROM nodes WHERE ip = ?', (ip,)).fetchone()
    conn.close()

    if not node:
        await websocket.send_text("\r\nFehler: Node nicht gefunden.\r\n")
        await websocket.close()
        return

    # Pseudo-Terminal für interaktive SSH-Session
    master_fd, slave_fd = pty.openpty()
    proc = await asyncio.create_subprocess_exec(
        "ssh", "-o", "StrictHostKeyChecking=no", "-t", f"{node['user']}@{ip}",
        stdin=slave_fd, stdout=slave_fd, stderr=slave_fd
    )

    async def pty_to_ws():
        # Setzt den Master-FD auf non-blocking
        fl = fcntl.fcntl(master_fd, fcntl.F_GETFL)
        fcntl.fcntl(master_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
        try:
            while True:
                await asyncio.sleep(0.01)
                try:
                    data = os.read(master_fd, 1024).decode(errors='ignore')
                    if data:
                        await websocket.send_text(data)
                except BlockingIOError:
                    continue
        except Exception: pass

    async def ws_to_pty():
        try:
            while True:
                data = await websocket.receive_text()
                os.write(master_fd, data.encode())
        except Exception: pass

    try:
        await asyncio.gather(pty_to_ws(), ws_to_pty())
    finally:
        if proc.returncode is None: proc.terminate()
        os.close(master_fd)
        os.close(slave_fd)

# --- WEBSOCKET CHAT UPDATE ---

@app.websocket("/ws/chat")
async def chat_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            user_msg = await websocket.receive_text()
            sys_prompt = get_system_prompt()
            ai_response = await get_ai_response(user_msg, sys_prompt)
            
            # Befehle extrahieren
            commands_to_run = re.findall(r'<EXECUTE target="(.*?)">(.*?)</EXECUTE>', ai_response, re.IGNORECASE | re.DOTALL)
            clean_chat_msg = re.sub(r'<EXECUTE.*?>.*?</EXECUTE>', '', ai_response, flags=re.IGNORECASE | re.DOTALL).strip()
            
            if clean_chat_msg:
                await websocket.send_text(clean_chat_msg)
            
            if commands_to_run:
                # Liste für alle laufenden Tasks erstellen
                tasks = []
                for target_ip, cmd in commands_to_run:
                    conn = get_db()
                    node = conn.execute('SELECT * FROM nodes WHERE ip = ? OR name = ?', (target_ip.strip(), target_ip.strip())).fetchone()
                    conn.close()
                    
                    if node:
                        # Wir erstellen den Task, starten ihn aber noch nicht separat
                        tasks.append(run_remote_task(node['ip'], node['user'], cmd.strip(), "KI-Kommando"))
                
                if tasks:
                    # Dem Nutzer im Chat kurz Bescheid geben
                    await websocket.send_text("ℹ️ *Warte auf Rückmeldungen der Nodes...*")
                    
                    # Jetzt werden alle SSH-Befehle gleichzeitig gestartet und abgewartet
                    await asyncio.gather(*tasks)
                    
                    # Sobald asyncio.gather fertig ist, geht es hier weiter mit dem Follow-up:
                    follow_up_prompt = "Die Befehle wurden ausgeführt. Bitte fasse die Ergebnisse kurz zusammen."
                    ai_summary = await get_ai_response(follow_up_prompt, sys_prompt)
                    
                    await websocket.send_text("--- Ergebnis-Zusammenfassung ---")
                    await websocket.send_text(ai_summary)

    except Exception as e:
        print(f"Chat Fehler: {e}")

async def run_remote_task(ip, user, cmd, task_name):
    global chat_history # Zugriff auf das Gedächtnis der KI
    
    await manager.broadcast(f"🚀 KI-Task gestartet: {cmd} auf {ip}")
    
    ssh_cmd = f"ssh -o StrictHostKeyChecking=no {user}@{ip} '{cmd}'"
    
    process = await asyncio.create_subprocess_shell(
        ssh_cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.STDOUT
    )
    
    # Hier speichern wir die komplette Terminal-Ausgabe
    full_output = ""
    
    while True:
        line = await process.stdout.readline()
        if not line:
            break
        decoded_line = line.decode('utf-8', errors='ignore').strip()
        if decoded_line:
            await manager.broadcast(f"🛠️ {decoded_line}")
            full_output += decoded_line + "\n"
            
    await process.wait()
    
    # --- NEU: Feedback an die KI ---
    # Wir bereiten den Bericht vor
    if not full_output.strip():
        full_output = "Befehl wurde ohne Ausgabe ausgeführt (Exit Code 0)."
        
    system_report = f"[SYSTEM-RÜCKMELDUNG] Der Befehl '{cmd}' auf Node {ip} wurde beendet. Ausgabe des Terminals:\n{full_output}"
    
    # Wir schmuggeln den Bericht als "User"-Nachricht in den Verlauf,
    # damit die KI beim nächsten Mal weiß, was passiert ist.
    chat_history.append({"role": "user", "content": system_report})
    # -------------------------------
    
    if "docker" in cmd.lower() and "install" in cmd.lower():
        # Kleiner Bonus: Nur updaten, wenn wirklich installiert wird
        await manager.broadcast(f"✨ Bitte aktualisiere den Status für {ip} manuell über das Refresh-Icon.")
        
    await manager.broadcast(f"✅ Befehl auf {ip} abgeschlossen.")
    
# --- Neuer Endpunkt: Manueller Refresh-Check ---
@app.get("/refresh_status/{node_id}")
async def refresh_status(node_id: int):
    conn = get_db()
    node = conn.execute('SELECT * FROM nodes WHERE id = ?', (node_id,)).fetchone()
    
    if node:
        # Kurzer Check via SSH, ob Docker antwortet
        check_cmd = "command -v docker >/dev/null 2>&1 && echo 'Docker Aktiv' || echo 'Bereit (Kein Docker)'"
        ssh_cmd = f"ssh -o StrictHostKeyChecking=no -o ConnectTimeout=2 {node['user']}@{node['ip']} \"{check_cmd}\""
        try:
            # Wir führen den Befehl aus
            new_status = subprocess.check_output(ssh_cmd, shell=True).decode().strip()
        except Exception:
            new_status = "Offline/Fehler"
            
        conn.execute('UPDATE nodes SET status = ? WHERE id = ?', (new_status, node_id))
        conn.commit()
        conn.close()
        return {"status": new_status} # Wir senden nur den Status zurück
    
    conn.close()
    return {"status": "Unbekannt"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)