import argparse
import json
import os
import subprocess
import sys
import time
import urllib.request
from datetime import datetime
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import HTMLResponse, StreamingResponse
from pydantic import BaseModel
from typing import List
import psutil
import uvicorn

# --- ARGUMENT PARSING ---
parser = argparse.ArgumentParser(description="LLM Edge Web Chat Bridge")
parser.add_argument("--host", type=str, default="0.0.0.0", help="Web Chat Server Host")
parser.add_argument("--port", type=int, default=8080, help="Web Chat Server Port")
parser.add_argument("--aaf-server", type=str, default="http://127.0.0.1:8000", help="AAF Connector Base URL")
parser.add_argument("--verbose", action="store_true", default=False, help="Enable verbose terminal logging")
args, _ = parser.parse_known_args()

MODEL_NAME = "Qwen2.5-7B-Instruct"
CONFIG_PATH = "/usr/share/eiq/aaf-connector/server_config.json"
SERVICE_NAME = "eiq-aaf-connector.service"

# --- SERVICE CONTROL & MODEL MANAGERS ---
def ensure_model_enabled(target_model: str, config_path: str) -> bool:
    if not os.path.exists(config_path):
        return False
    try:
        with open(config_path, "r") as f:
            config = json.load(f)

        config_changed = False
        available_models = config.get("available_models", [])

        if isinstance(available_models, list):
            for model_info in available_models:
                model_id = model_info.get("name") or model_info.get("id")
                should_be_enabled = (model_id == target_model)
                if model_info.get("enabled") != should_be_enabled:
                    model_info["enabled"] = should_be_enabled
                    config_changed = True

        if config_changed:
            print(f"[Config] Enabling '{target_model}' in server_config.json...")
            with open(config_path, "w") as f:
                json.dump(config, f, indent=4)
            return True
    except Exception as e:
        print(f"[Config Error] {e}")
    return False

def is_service_running(service_name: str) -> bool:
    try:
        res = subprocess.run(
            ["systemctl", "is-active", service_name],
            stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
        )
        return res.stdout.strip() == "active"
    except Exception:
        return False

def restart_aaf_service(service_name: str):
    print(f"[Service] Restarting '{service_name}'...")
    try:
        subprocess.run(["systemctl", "restart", service_name], check=True)
    except subprocess.CalledProcessError as e:
        print(f"[Error] Failed to restart service: {e}")
        sys.exit(1)

def wait_for_model_ready(target_model: str, server_base_url: str, timeout: int = 300) -> bool:
    models_url = f"{server_base_url}/v1/models"
    print(f"[Service] Waiting for '{target_model}' to finish loading on NPU...", end="", flush=True)
    start_time = time.time()

    while time.time() - start_time < timeout:
        try:
            req = urllib.request.Request(models_url, headers={"Accept": "application/json"})
            with urllib.request.urlopen(req, timeout=2.0) as resp:
                if resp.status == 200:
                    data = json.loads(resp.read().decode('utf-8'))
                    loaded_models = [m.get("id") or m.get("name") for m in data.get("data", [])]
                    if target_model in loaded_models or len(loaded_models) > 0:
                        print("\n[Service] AAF Connector is ready and LLM endpoint is active!")
                        return True
        except Exception:
            pass

        print(".", end="", flush=True)
        time.sleep(2.0)

    print("\n[Error] Timed out waiting for AAF Connector model initialization.")
    return False

def prepare_aaf_backend():
    config_updated = ensure_model_enabled(MODEL_NAME, CONFIG_PATH)
    service_active = is_service_running(SERVICE_NAME)

    if config_updated or not service_active:
        if not service_active:
            print(f"[Service] '{SERVICE_NAME}' is not active. Starting service...")
        restart_aaf_service(SERVICE_NAME)

    if not wait_for_model_ready(MODEL_NAME, args.aaf_server, timeout=300):
        print(f"[Error] Unable to reach target model '{MODEL_NAME}' on AAF Server ({args.aaf_server}).")
        sys.exit(1)

# --- FASTAPI WEB APPLICATION ---
app = FastAPI(title="Gateworks AI LLM Chat")

class ChatMessage(BaseModel):
    role: str
    content: str

class ChatPayload(BaseModel):
    messages: List[ChatMessage]

def get_system_thermals():
    zones = []
    try:
        for zone in sorted(os.listdir("/sys/class/thermal/")):
            if zone.startswith("thermal_zone0"):
                with open(f"/sys/class/thermal/{zone}/temp", "r") as f:
                    z_temp = int(f.read().strip()) / 1000.0
                zones.append(z_temp)
    except Exception:
        pass
    return zones

@app.get("/api/telemetry")
def get_telemetry():
    cpu_usage = psutil.cpu_percent()
    sys_ram = psutil.virtual_memory().percent
    thermals = get_system_thermals()
    temp_str = "/".join([f"{t:.1f}C" for t in thermals]) if thermals else "N/A"
    return {
        "cpu": cpu_usage,
        "ram": sys_ram,
        "temps": temp_str
    }

@app.post("/api/chat")
def chat_stream(payload: ChatPayload, request: Request):
    client_ip = request.client.host
    start_time = time.time()
    
    aaf_payload = {
        "model": MODEL_NAME,
        "messages": [msg.dict() for msg in payload.messages],
        "temperature": 0.7,
        "stream": True
    }

    print(f"[{datetime.now().strftime('%H:%M:%S')}] INFERENCE REQ from {client_ip} | Messages: {len(payload.messages)}")

    def stream_generator():
        target_endpoint = f"{args.aaf_server}/v1/chat/completions"
        data_bytes = json.dumps(aaf_payload).encode('utf-8')

        req = urllib.request.Request(
            target_endpoint,
            data=data_bytes,
            headers={"Content-Type": "application/json", "Accept": "application/json"},
            method="POST"
        )

        try:
            with urllib.request.urlopen(req, timeout=300.0) as response:
                while True:
                    line_bytes = response.readline()
                    if not line_bytes:
                        break
                    line_str = line_bytes.decode('utf-8', errors='ignore').strip()
                    if line_str:
                        yield f"{line_str}\n\n"
        except urllib.error.HTTPError as http_err:
            err_body = http_err.read().decode('utf-8', errors='ignore')
            yield f"data: {{\"error\": \"AAF Server Engine error {http_err.code}: {err_body}\"}}\n\n"
        except Exception as e:
            yield f"data: {{\"error\": \"Pipeline fault: {str(e)}\"}}\n\n"

    return StreamingResponse(stream_generator(), media_type="text/event-stream")

@app.get("/", response_class=HTMLResponse)
def serve_ui():
    html_content = r"""
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <title>Gateworks AI</title>
        <script src="https://cdn.tailwindcss.com"></script>
    </head>
    <body class="bg-gray-900 text-gray-100 min-h-screen p-6 font-mono">
        <div class="max-w-5xl mx-auto grid grid-cols-1 md:grid-cols-4 gap-6">
            <!-- SIDEBAR -->
            <div class="md:col-span-1 space-y-4 bg-gray-800 p-4 rounded-lg border border-gray-700 h-fit">
                <h2 class="text-lg font-bold text-indigo-400">Gateworks AI</h2>
                <div class="text-xs space-y-2 border-t border-gray-700 pt-3">
                    <div><b>Model:</b> <span class="text-indigo-300">__MODEL_NAME__</span></div>
                    <div><b>Status:</b> <span id="statusBadge" class="text-green-400 font-bold">🟢 READY</span></div>
                </div>
                
                <div class="text-xs space-y-1 border-t border-gray-700 pt-3">
                    <div class="font-bold text-gray-400">💻 Telemetry</div>
                    <div id="telemetryData" class="text-gray-300">Loading metrics...</div>
                </div>

                <div class="text-xs space-y-1 border-t border-gray-700 pt-3">
                    <div class="font-bold text-gray-400">⚡ Performance</div>
                    <div id="perfData" class="text-indigo-300">N/A</div>
                </div>

                <button id="clearChatBtn" class="w-full bg-red-900/40 hover:bg-red-800 border border-red-700 text-red-200 text-xs py-1.5 px-3 rounded mt-4">
                    Clear Conversation
                </button>
            </div>

            <!-- MAIN CHAT -->
            <div class="md:col-span-3 flex flex-col h-[650px] bg-gray-800 rounded-lg border border-gray-700 p-4">
                <div id="chatHistory" class="flex-1 overflow-y-auto space-y-3 p-3 bg-gray-900 rounded border border-gray-750">
                    <div class="text-gray-500 text-xs italic">System initialized. Type a message to start...</div>
                </div>

                <div class="mt-4 flex space-x-2">
                    <textarea id="userInput" rows="2" class="flex-1 bg-gray-900 border border-gray-700 rounded p-2.5 text-sm focus:outline-none focus:border-indigo-500 resize-none" placeholder="Ask the i.MX LLM..."></textarea>
                    <button id="sendBtn" class="bg-indigo-600 hover:bg-indigo-500 text-white px-5 rounded font-medium text-sm">Send</button>
                </div>
            </div>
        </div>

        <script>
            const chatHistory = document.getElementById('chatHistory');
            const userInput = document.getElementById('userInput');
            const sendBtn = document.getElementById('sendBtn');
            const clearChatBtn = document.getElementById('clearChatBtn');
            const telemetryData = document.getElementById('telemetryData');
            const perfData = document.getElementById('perfData');
            const statusBadge = document.getElementById('statusBadge');

            let messages = [];

            async function updateTelemetry() {
                try {
                    const res = await fetch('/api/telemetry');
                    const data = await res.json();
                    telemetryData.innerHTML = `CPU: ${data.cpu}% | RAM: ${data.ram}%<br>Temp: ${data.temps}`;
                } catch(e) {}
            }
            setInterval(updateTelemetry, 3000);
            updateTelemetry();

            function appendMessage(role, text) {
                const msgDiv = document.createElement('div');
                msgDiv.className = `p-3 rounded text-sm border ${role === 'user' ? 'bg-indigo-950/50 border-indigo-800 ml-8' : 'bg-gray-800 border-gray-700 mr-8'}`;
                
                const label = document.createElement('div');
                label.className = `text-[10px] font-bold uppercase mb-1 ${role === 'user' ? 'text-indigo-400' : 'text-gray-400'}`;
                label.textContent = role === 'user' ? 'You' : 'AI Assistant';

                const body = document.createElement('div');
                body.className = "whitespace-pre-wrap leading-relaxed";
                body.textContent = text;

                msgDiv.appendChild(label);
                msgDiv.appendChild(body);
                chatHistory.appendChild(msgDiv);
                chatHistory.scrollTop = chatHistory.scrollHeight;
                return body;
            }

            async function sendMessage() {
                const text = userInput.value.trim();
                if (!text) return;

                userInput.value = "";
                sendBtn.disabled = true;
                statusBadge.textContent = "⚠️ BUSY";
                statusBadge.className = "text-yellow-400 font-bold";

                appendMessage('user', text);
                messages.push({ role: 'user', content: text });

                const aiNode = appendMessage('assistant', "Thinking...");
                let fullReply = "";
                let tokenCount = 0;
                const startTime = performance.now();

                try {
                    const res = await fetch('/api/chat', {
                        method: 'POST',
                        headers: { 'Content-Type': 'application/json' },
                        body: JSON.stringify({ messages })
                    });

                    aiNode.textContent = "";
                    const reader = res.body.getReader();
                    const decoder = new TextDecoder();
                    let buffer = "";

                    while (true) {
                        const { done, value } = await reader.read();
                        if (done) break;

                        buffer += decoder.decode(value, { stream: true });
                        const lines = buffer.split('\n');
                        buffer = lines.pop();

                        for (const line of lines) {
                            const trimmed = line.trim();
                            if (!trimmed || !trimmed.startsWith('data: ')) continue;

                            const dataStr = trimmed.slice(5).trim();
                            if (dataStr === '[DONE]') continue;

                            try {
                                const chunk = JSON.parse(dataStr);
                                const token = chunk.choices?.[0]?.delta?.content || "";
                                if (token) {
                                    fullReply += token;
                                    tokenCount++;
                                    aiNode.textContent = fullReply;
                                    chatHistory.scrollTop = chatHistory.scrollHeight;
                                }
                            } catch(e) {}
                        }
                    }

                    messages.push({ role: 'assistant', content: fullReply });
                    const duration = (performance.now() - startTime) / 1000;
                    const tps = (tokenCount / duration).toFixed(1);
                    perfData.textContent = `${tokenCount} tokens @ ${tps} t/s`;

                } catch(err) {
                    aiNode.textContent = `[Error]: ${err.message}`;
                } finally {
                    sendBtn.disabled = false;
                    statusBadge.textContent = "🟢 READY";
                    statusBadge.className = "text-green-400 font-bold";
                }
            }

            sendBtn.addEventListener('click', sendMessage);
            userInput.addEventListener('keydown', (e) => {
                if (e.key === 'Enter' && !e.shiftKey) {
                    e.preventDefault();
                    sendMessage();
                }
            });

            clearChatBtn.addEventListener('click', () => {
                messages = [];
                chatHistory.innerHTML = `<div class="text-gray-500 text-xs italic">Conversation wiped. Ready...</div>`;
                perfData.textContent = "N/A";
            });
        </script>
    </body>
    </html>
    """
    return HTMLResponse(content=html_content.replace("__MODEL_NAME__", MODEL_NAME))

if __name__ == "__main__":
    prepare_aaf_backend()
    uvicorn.run(app, host=args.host, port=args.port)