import os
from dotenv import load_dotenv

# Dynamically select configuration file based on APP_ENV
app_env = os.getenv("APP_ENV", "development")
if app_env == "production":
    load_dotenv(".env.production")
else:
    load_dotenv(".env.development")

# Fallback to load default .env if any variables are not yet defined
load_dotenv()

from fastapi import FastAPI, HTTPException, Depends, Header
from fastapi.responses import PlainTextResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Dict, Any, List
import uvicorn
import base64
import hmac
import hashlib
import json
from datetime import datetime, timedelta

app = FastAPI(title="RMM Central API")

# Enable CORS for frontend dashboard queries (React dev server runs on a separate port)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# JWT-like HMAC Secure Token Utilities
JWT_SECRET_KEY = os.getenv("JWT_SECRET_KEY", "seekright_rmm_secret_key_2026_default")

def generate_token(username: str) -> str:
    # Expire in 1 day (24 hours)
    expiry = (datetime.utcnow() + timedelta(days=1)).isoformat()
    payload = {
        "username": username,
        "expires": expiry
    }
    payload_json = json.dumps(payload)
    payload_b64 = base64.urlsafe_b64encode(payload_json.encode('utf-8')).decode('utf-8').rstrip('=')
    
    sig = hmac.new(JWT_SECRET_KEY.encode('utf-8'), payload_b64.encode('utf-8'), hashlib.sha256).hexdigest()
    return f"{payload_b64}.{sig}"

def verify_token(token: str) -> bool:
    try:
        parts = token.split(".")
        if len(parts) != 2:
            return False
        payload_b64, sig = parts
        
        # Verify signature
        expected_sig = hmac.new(JWT_SECRET_KEY.encode('utf-8'), payload_b64.encode('utf-8'), hashlib.sha256).hexdigest()
        if not hmac.compare_digest(sig, expected_sig):
            return False
            
        # Decode and check expiry
        padding = 4 - (len(payload_b64) % 4)
        if padding < 4:
            payload_b64 += "=" * padding
        payload_json = base64.urlsafe_b64decode(payload_b64.encode('utf-8')).decode('utf-8')
        payload = json.loads(payload_json)
        
        expires_dt = datetime.fromisoformat(payload["expires"])
        if datetime.utcnow() > expires_dt:
            return False # Expired
            
        return payload["username"] == "root"
    except Exception:
        return False

# FastAPI dependency to secure UI endpoints
async def get_current_user(authorization: str = Header(None)):
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing or invalid authentication credentials")
    token = authorization.split(" ")[1]
    if not verify_token(token):
        raise HTTPException(status_code=401, detail="Authentication token is invalid or has expired")
    return "root"

import os
import json
from datetime import datetime

# In-memory storage for client hardware telemetry
client_telemetry = {}

from pymongo import MongoClient

# Database Connection: Configured via environment variable with local fallback
MONGODB_URI = os.getenv("MONGODB_URI", "mongodb://localhost:27017")
DB_NAME = os.getenv("MONGODB_DB", "rmm_db")

mongo_client = MongoClient(MONGODB_URI)
db = mongo_client[DB_NAME]

clients_collection = db["clients"]
logs_collection = db["logs"]

MAX_ENTRIES_PER_CLIENT = 100  # Keep the latest 100 historical logs per client node

ROCKETCHAT_WEBHOOK_URL = os.getenv("ROCKETCHAT_WEBHOOK_URL", "")

# Alert states are now stored directly in the MongoDB client documents 
# to support multi-worker environments safely.

def send_rocketchat_notification(text: str, color: str = "#808080"):
    """
    Dispatches a formatted notification payload to the Rocket.Chat incoming webhook.
    """
    if not ROCKETCHAT_WEBHOOK_URL:
        try:
            print(f"[Rocket.Chat Simulation] {text}")
        except UnicodeEncodeError:
            # Fallback for Windows consoles that do not support printing unicode emojis
            sanitized_text = text.encode('ascii', errors='backslashreplace').decode('ascii')
            print(f"[Rocket.Chat Simulation] {sanitized_text}")
        return

    payload = {
        "text": text,
        "attachments": [
            {
                "color": color,
                "ts": datetime.now().isoformat()
            }
        ]
    }
    
    try:
        import urllib.request
        import json
        req = urllib.request.Request(
            ROCKETCHAT_WEBHOOK_URL,
            data=json.dumps(payload).encode('utf-8'),
            headers={'Content-Type': 'application/json'}
        )
        with urllib.request.urlopen(req, timeout=3.0) as response:
            pass
    except Exception as e:
        print(f"[!] Error sending Rocket.Chat notification: {e}")

def append_to_logs(log_type: str, client_id: str, detail: Any):
    """
    Appends a new event log grouped by client_id in MongoDB logs collection.
    Ensures fair-share log limits per system so noisy clients never overwrite others.
    """
    try:
        entry = {
            "client_id": client_id,
            "timestamp": datetime.now().isoformat(),
            "type": log_type,
            "detail": detail
        }
        logs_collection.insert_one(entry)

        # Enforce fair-share logging limits per system (FIFO circular buffer)
        count = logs_collection.count_documents({"client_id": client_id})
        if count > MAX_ENTRIES_PER_CLIENT:
            oldest_docs = logs_collection.find(
                {"client_id": client_id},
                {"_id": 1}
            ).sort("timestamp", 1).limit(count - MAX_ENTRIES_PER_CLIENT)
            
            ids_to_delete = [doc["_id"] for doc in oldest_docs]
            if ids_to_delete:
                logs_collection.delete_many({"_id": {"$in": ids_to_delete}})
    except Exception as e:
        print(f"[!] Error writing log to MongoDB: {e}")

def mark_client_active(client_id: str):
    try:
        now_str = datetime.now().isoformat()
        # Atomically check if transitioned from offline to online (active is not True)
        old_doc = clients_collection.find_one_and_update(
            {"_id": client_id, "active": {"$ne": True}},
            {"$set": {"active": True, "last_seen": now_str, "last_down_alert_time": None}},
            return_document=False
        )
        if old_doc:
            # Client transitioned from inactive/None to active!
            # Send UP recovery alert only if they had a down alert time logged
            if old_doc.get("last_down_alert_time"):
                send_rocketchat_notification(
                    text=f"✅ **System UP:** Client `{client_id}` has recovered and is back online.",
                    color="#2ecc71"
                )
        else:
            # Client is already active, just update their last_seen timestamp
            clients_collection.update_one(
                {"_id": client_id},
                {"$set": {"last_seen": now_str}}
            )
    except Exception as e:
        print(f"[!] Error marking client active: {e}")

def load_clients() -> Dict[str, Any]:
    try:
        # Clean up any ghost clients with empty or whitespace-only keys
        clients_collection.delete_many({"_id": {"$in": ["", None]}})
        clients_collection.delete_many({"_id": {"$regex": "^\\s*$"}})

        # Load all documents
        cursor = clients_collection.find()
        data = {}
        for doc in cursor:
            cid = doc["_id"]
            
            # Compute active status dynamically for reads, but do NOT write back to database or alert
            last_seen_str = doc.get("last_seen")
            is_active = doc.get("active", False)
            if last_seen_str:
                try:
                    last_seen_dt = datetime.fromisoformat(last_seen_str)
                    if (datetime.now() - last_seen_dt).total_seconds() >= 30:
                        is_active = False
                except Exception:
                    pass
            
            info = {k: v for k, v in doc.items() if k != "_id"}
            info["active"] = is_active
            data[cid] = info
            
        return data
    except Exception as e:
        print(f"[!] Error loading clients from MongoDB: {e}")
        return {}

import asyncio

async def monitor_heartbeats_loop():
    while True:
        try:
            now = datetime.now()
            timeout_time = (now - timedelta(seconds=30)).isoformat()
            
            # Find clients that are active in DB but haven't checked in for 30s
            cursor = clients_collection.find({
                "active": True,
                "last_seen": {"$lt": timeout_time}
            })
            for doc in cursor:
                cid = doc["_id"]
                now_str = now.isoformat()
                
                # Atomically update to active = False
                # Ensures only one request/worker triggers the state transition and sends the alert
                old_doc = clients_collection.find_one_and_update(
                    {"_id": cid, "active": True, "last_seen": doc["last_seen"]},
                    {"$set": {"active": False, "last_down_alert_time": now_str}},
                    return_document=False
                )
                if old_doc:
                    send_rocketchat_notification(
                        text=f"🚨 **System DOWN:** Client `{cid}` has missed heartbeats for over 30 seconds.",
                        color="#e74c3c"
                    )
            
            # Hourly reminders for STILL DOWN clients
            reminder_time = (now - timedelta(hours=1)).isoformat()
            cursor_still_down = clients_collection.find({
                "active": False,
                "last_down_alert_time": {"$lt": reminder_time}
            })
            for doc in cursor_still_down:
                cid = doc["_id"]
                old_last_down = doc["last_down_alert_time"]
                now_str = now.isoformat()
                
                # Atomically update reminder timestamp
                old_doc = clients_collection.find_one_and_update(
                    {"_id": cid, "active": False, "last_down_alert_time": old_last_down},
                    {"$set": {"last_down_alert_time": now_str}},
                    return_document=False
                )
                if old_doc:
                    send_rocketchat_notification(
                        text=f"🚨 **System STILL DOWN:** Client `{cid}` remains offline (reminder sent every hour).",
                        color="#e74c3c"
                    )
        except Exception as e:
            print(f"[!] Error in heartbeat monitoring loop: {e}")
            
        await asyncio.sleep(10)

@app.on_event("startup")
async def startup_event():
    asyncio.create_task(monitor_heartbeats_loop())
    print("[*] Heartbeat monitoring background task initialized.")

class TelemetryPayload(BaseModel):
    cpu_percent: float
    memory_percent: float
    memory_total_gb: float
    memory_free_gb: float
    disks: List[Dict[str, Any]]
    gpus: List[Dict[str, Any]]

class CommandResultPayload(BaseModel):
    command: str
    returncode: int
    stdout: str
    stderr: str

class LoginPayload(BaseModel):
    username: str
    password: str

@app.post("/api/login")
async def login(payload: LoginPayload):
    if payload.username == "root" and payload.password == "seekright159@":
        token = generate_token("root")
        return {"token": token}
    raise HTTPException(status_code=401, detail="Invalid username or password")

@app.get("/api/get-commands")
async def get_whitelisted_commands(platform: str):
    """
    Endpoint for remote agents to fetch their OS-specific whitelisted commands.
    """
    commands_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "commands.json")
    try:
        with open(commands_file, "r") as f:
            data = json.load(f)
            return data.get(platform, {})
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to load commands: {e}")

@app.get("/api/get-command")
async def get_command(client_id: str):
    """
    Endpoint for clients to poll for pending commands.
    Client Agent hits this endpoint to ask: "Do I have any work to do?"
    """
    if not client_id or not client_id.strip():
        raise HTTPException(status_code=400, detail="client_id cannot be empty")
        
    now_str = datetime.now().isoformat()
    # Dynamic fleet auto-registration
    clients_collection.update_one(
        {"_id": client_id},
        {"$setOnInsert": {"pending_command": "none", "active": True, "last_seen": now_str}},
        upsert=True
    )
    
    # Get and reset command atomically
    updated_doc = clients_collection.find_one_and_update(
        {"_id": client_id},
        {"$set": {"pending_command": "none"}},
        return_document=False # Returns the state before update
    )
    
    cmd = "none"
    if updated_doc:
        cmd = updated_doc.get("pending_command", "none")
        
    if cmd != "none":
        append_to_logs("command_polled", client_id, {"command": cmd})
        
    mark_client_active(client_id)
    return {"command": cmd}

@app.post("/api/schedule-command")
async def schedule_command(client_id: str, command: str, current_user: str = Depends(get_current_user)):
    """
    Endpoint for your Dashboard/UI to schedule a new command.
    Supports a comma-separated list of client IDs for batch fleet updates.
    """
    target_ids = [cid.strip() for cid in client_id.split(",") if cid.strip()]
    if not target_ids:
        raise HTTPException(status_code=400, detail="No target client IDs specified")

    for cid in target_ids:
        clients_collection.update_one(
            {"_id": cid},
            {"$set": {"pending_command": command}},
            upsert=True
        )
        append_to_logs("command_scheduled", cid, {"command": command})
        
    return {"status": "success", "message": f"Command '{command}' scheduled for {', '.join(target_ids)}"}

@app.post("/api/telemetry")
async def receive_telemetry(client_id: str, payload: TelemetryPayload):
    """
    Endpoint for agents to push their live hardware telemetry.
    """
    if not client_id or not client_id.strip():
        raise HTTPException(status_code=400, detail="client_id cannot be empty")
        
    now_str = datetime.now().isoformat()
    # Auto-register client if we see it through telemetry first
    clients_collection.update_one(
        {"_id": client_id},
        {"$setOnInsert": {"pending_command": "none", "active": True, "last_seen": now_str}},
        upsert=True
    )
    
    mark_client_active(client_id)
    
    # Get current alert states from database to prevent duplicate alerts
    client_doc = clients_collection.find_one({"_id": client_id}) or {}
    storage_alerts = client_doc.get("active_storage_alerts", [])
    gpu_alert_sent = client_doc.get("gpu_alert_sent", False)
    
    storage_modified = False
    gpu_modified = False

    # Storage warning checker
    for disk in payload.disks:
        mount = disk.get("mount", "/")
        percent = disk.get("percent", 0.0)
        
        if percent >= 80.0:
            if mount not in storage_alerts:
                storage_alerts.append(mount)
                storage_modified = True
                send_rocketchat_notification(
                    text=f"⚠️ **Storage Warning:** Client `{client_id}` disk `{mount}` is at **{percent}%** capacity.",
                    color="#f39c12"
                )
        else:
            if mount in storage_alerts:
                storage_alerts.remove(mount)
                storage_modified = True
                send_rocketchat_notification(
                    text=f"✅ **Storage Recovered:** Client `{client_id}` disk `{mount}` has cleared warning state and is at **{percent}%**.",
                    color="#2ecc71"
                )

    # GPU Failure Monitor
    if len(payload.gpus) == 0:
        if not gpu_alert_sent:
            gpu_alert_sent = True
            gpu_modified = True
            send_rocketchat_notification(
                text=f"🚨 **GPU Failure:** Client `{client_id}` is not reporting any GPU data! (nvidia-smi has failed or is missing)",
                color="#e74c3c"
            )
    else:
        if gpu_alert_sent:
            gpu_alert_sent = False
            gpu_modified = True
            send_rocketchat_notification(
                text=f"✅ **GPU Recovered:** Client `{client_id}` is reporting GPU data successfully again.",
                color="#2ecc71"
            )

    update_fields = {
        "telemetry": payload.dict()
    }
    if storage_modified:
        update_fields["active_storage_alerts"] = storage_alerts
    if gpu_modified:
        update_fields["gpu_alert_sent"] = gpu_alert_sent
        
    clients_collection.update_one(
        {"_id": client_id},
        {"$set": update_fields}
    )

    # Log telemetry history event
    append_to_logs("telemetry", client_id, payload.dict())

    data = payload.dict()
    client_telemetry[client_id] = data
    
    # Log received stats
    print(f"\n[TELEMETRY RECEIVED] from {client_id}:")
    print(f"  CPU:  {data['cpu_percent']}%")
    print(f"  RAM:  {data['memory_percent']}% ({data['memory_free_gb']} GB Free / {data['memory_total_gb']} GB Total)")
    
    for disk in data['disks']:
        print(f"  Drive {disk['mount']}: {disk['percent']}% Used ({disk['free_gb']} GB Free / {disk['total_gb']} GB Total)")
        
    for gpu in data['gpus']:
        power_str = f"Power: {gpu.get('power_draw', 'N/A')} / {gpu.get('power_limit', 'N/A')}"
        fan_str = f"Fan: {gpu.get('fan_speed', 'N/A')}"
        print(f"  GPU [{gpu['name']}]: Core: {gpu['utilization']}, Temp: {gpu['temp']}, VRAM: {gpu['memory_used']}/{gpu['memory_total']}, {power_str}, {fan_str}")
        
    return {"status": "success"}

    data = payload.dict()
    client_telemetry[client_id] = data
    
    # Temporarily print the data to the console for the user to see!
    print(f"\n[TELEMETRY RECEIVED] from {client_id}:")
    print(f"  CPU:  {data['cpu_percent']}%")
    print(f"  RAM:  {data['memory_percent']}% ({data['memory_free_gb']} GB Free / {data['memory_total_gb']} GB Total)")
    
    for disk in data['disks']:
        print(f"  Drive {disk['mount']}: {disk['percent']}% Used ({disk['free_gb']} GB Free / {disk['total_gb']} GB Total)")
        
    for gpu in data['gpus']:
        power_str = f"Power: {gpu.get('power_draw', 'N/A')} / {gpu.get('power_limit', 'N/A')}"
        fan_str = f"Fan: {gpu.get('fan_speed', 'N/A')}"
        print(f"  GPU [{gpu['name']}]: Core: {gpu['utilization']}, Temp: {gpu['temp']}, VRAM: {gpu['memory_used']}/{gpu['memory_total']}, {power_str}, {fan_str}")
        
    return {"status": "success"}

@app.get("/api/telemetry")
async def get_all_telemetry(current_user: str = Depends(get_current_user)):
    """
    Endpoint to view the live dashboard data of all clients.
    """
    clients = load_clients()
    return {cid: info.get("telemetry", {}) for cid, info in clients.items() if info.get("telemetry")}

@app.get("/api/clients")
async def get_clients_api(current_user: str = Depends(get_current_user)):
    """
    Endpoint for the React UI to fetch the live active client registry with computed statuses.
    """
    return load_clients()

@app.get("/api/logs")
async def get_logs_history(current_user: str = Depends(get_current_user)):
    """
    Endpoint to retrieve logs history grouped by client_id.
    """
    try:
        cursor = logs_collection.find().sort("timestamp", 1)
        grouped_logs = {}
        for doc in cursor:
            cid = doc.get("client_id")
            if not cid:
                continue
            if cid not in grouped_logs:
                grouped_logs[cid] = []
            
            entry = {
                "timestamp": doc.get("timestamp"),
                "type": doc.get("type"),
                "detail": doc.get("detail")
            }
            grouped_logs[cid].append(entry)
        return grouped_logs
    except Exception as e:
        print(f"[!] Error loading logs from MongoDB: {e}")
        return {}

@app.post("/api/command-result")
async def receive_command_result(client_id: str, payload: CommandResultPayload):
    """
    Endpoint for remote agents to push execution results back to the central server.
    """
    if not client_id or not client_id.strip():
        raise HTTPException(status_code=400, detail="client_id cannot be empty")
        
    append_to_logs("command_result", client_id, payload.dict())
    return {"status": "success"}

@app.get("/api/get-raw-commands")
async def get_raw_commands(current_user: str = Depends(get_current_user)):
    """
    Endpoint for the UI to load the complete commands.json whitelist file.
    """
    commands_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "commands.json")
    try:
        with open(commands_file, "r") as f:
            return json.load(f)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to read commands: {e}")

@app.post("/api/save-commands")
async def save_commands(commands: Dict[str, Any], current_user: str = Depends(get_current_user)):
    """
    Endpoint for the UI to save changes back to commands.json.
    """
    commands_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "commands.json")
    try:
        with open(commands_file, "w") as f:
            json.dump(commands, f, indent=2)
        return {"status": "success", "message": "commands.json saved successfully"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to save commands: {e}")

@app.get("/metrics", response_class=PlainTextResponse)
async def get_prometheus_metrics():
    """
    Exposes the latest client telemetry in standard Prometheus text format.
    """
    clients = load_clients()
    metrics_lines = []
    
    for client_id, data in clients.items():
        # Active status (1.0 for active/online, 0.0 for inactive/offline)
        active_val = 1.0 if data.get("active", False) else 0.0
        metrics_lines.append(f'rmm_client_active{{client_id="{client_id}"}} {active_val}')
        
        # Telemetry stats
        telemetry = data.get("telemetry")
        if telemetry:
            cpu = telemetry.get("cpu_percent", 0.0)
            ram = telemetry.get("memory_percent", 0.0)
            ram_free = telemetry.get("memory_free_gb", 0.0)
            ram_total = telemetry.get("memory_total_gb", 0.0)
            
            metrics_lines.append(f'rmm_cpu_utilization{{client_id="{client_id}"}} {cpu}')
            metrics_lines.append(f'rmm_memory_utilization{{client_id="{client_id}"}} {ram}')
            metrics_lines.append(f'rmm_memory_free_bytes{{client_id="{client_id}"}} {ram_free * (1024**3)}')
            metrics_lines.append(f'rmm_memory_total_bytes{{client_id="{client_id}"}} {ram_total * (1024**3)}')
            
            # Disk mounts
            for disk in telemetry.get("disks", []):
                mount = disk.get("mount", "/")
                disk_percent = disk.get("percent", 0.0)
                metrics_lines.append(f'rmm_disk_utilization{{client_id="{client_id}",mount="{mount}"}} {disk_percent}')
            
            # GPU utilization and temps
            for gpu in telemetry.get("gpus", []):
                gpu_name = gpu.get("name", "GPU")
                
                # Utilization percent parsing
                gpu_util_str = str(gpu.get("utilization", "0"))
                gpu_util = float(gpu_util_str.replace("%", "").strip())
                
                # Temperature parsing
                gpu_temp_str = str(gpu.get("temp", "0"))
                gpu_temp = float(gpu_temp_str.replace("C", "").strip())
                
                metrics_lines.append(f'rmm_gpu_utilization{{client_id="{client_id}",gpu_name="{gpu_name}"}} {gpu_util}')
                metrics_lines.append(f'rmm_gpu_temperature{{client_id="{client_id}",gpu_name="{gpu_name}"}} {gpu_temp}')
                
    return "\n".join(metrics_lines) + "\n"

if __name__ == "__main__":
    print("[*] Starting FastAPI Central Server...")
    # uvicorn runs the FastAPI app on port 8000
    uvicorn.run("central_api_prototype:app", host="0.0.0.0", port=8000, reload=True)