import os
import pty
import fcntl
import subprocess
import sqlite3
import asyncio
import openai
import google.generativeai as genai
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

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"

# --- KI KONFIGURATION ---
AI_PROVIDER = "ollama" # "openai", "google" oder "ollama"
OPENAI_API_KEY = "dein-key"
GOOGLE_API_KEY = "dein-key"
OLLAMA_BASE_URL = "http://x.x.x.x:11434/v1" # IP deines Ollama-Servers
# System Prompt: Sagt der KI, wer sie ist und was sie kann
SYSTEM_PROMPT = """Du bist der Pi-Orchestrator KI-Assistent. 
Deine Aufgabe ist es, dem Nutzer zu helfen, Raspberry Pis zu verwalten und Docker-Container zu steuern.
Du kannst Linux-Befehle generieren. Antworte präzise und hilfsbereit."""

# --- KI FUNKTIONEN ---

async def get_ai_response(user_input):
    if AI_PROVIDER == "openai":
        client = openai.OpenAI(api_key=OPENAI_API_KEY)
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_input}]
        )
        return response.choices[0].message.content

    elif AI_PROVIDER == "ollama":
        # Ollama nutzt das OpenAI-Format, braucht aber keinen Key
        client = openai.OpenAI(base_url=OLLAMA_BASE_URL, api_key="ollama")
        response = client.chat.completions.create(
            model="llama3", # Oder dein bevorzugtes Modell
            messages=[{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_input}]
        )
        return response.choices[0].message.content

    elif AI_PROVIDER == "google":
        genai.configure(api_key=GOOGLE_API_KEY)
        model = genai.GenerativeModel('gemini-1.5-flash')
        # Gemini braucht einen etwas anderen Aufbau für System Prompts
        response = model.generate_content(f"{SYSTEM_PROMPT}\n\nNutzer: {user_input}")
        return response.text

    return "Fehler: Kein KI-Provider konfiguriert."


# --- 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 websocket_chat(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            # Hole Antwort von der gewählten KI
            ai_msg = await get_ai_response(data)
            await websocket.send_text(ai_msg)
    except Exception as e:
        print(f"Chat Error: {e}")

# --- Status in DB aktualisieren Helper ---
def update_node_status(ip, new_status):
    conn = get_db()
    conn.execute('UPDATE nodes SET status = ? WHERE ip = ?', (new_status, ip))
    conn.commit()
    conn.close()

# --- Erweiterter Remote-Task ---
async def run_remote_task(ip, user, cmd, task_name):
    await manager.broadcast(f"🚀 KI-Task: {task_name} auf {ip}")
    ssh_cmd = f"ssh -o StrictHostKeyChecking=no {user}@{ip} '{cmd}'"
    process = subprocess.Popen(ssh_cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    
    for line in process.stdout:
        await manager.broadcast(f"🛠️ {line.strip()[:80]}")
    
    # WICHTIG: Wenn Docker installiert wurde, Status ändern!
    if "docker" in task_name.lower():
        update_node_status(ip, "Docker Aktiv")
        await manager.broadcast(f"✨ Datenbank aktualisiert: {ip} ist nun 'Docker Aktiv'")
    
    await manager.broadcast(f"✅ {task_name} fertig.")

# --- 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)