#!/usr/bin/env python3
"""
⚡ MDA Local Hardware Model Bridge (v3.1)
Connects your local Mac or PC (Ollama, LM Studio, Apple Foundation Model)
to the MDA Intelligence Cloud Platform over a secure WebSocket tunnel.
Features macOS Menubar status indicator, auto-reconnect, and real-time streaming.

Usage:
    python3 mda_local_bridge.py
    (or CLI mode without menubar: python3 mda_local_bridge.py --no-gui)
"""

import asyncio
import argparse
import json
import os
import sys
import subprocess
import threading
import time
import re
import webbrowser

# Auto-install missing dependencies if needed
try:
    import httpx
    import websockets
except ImportError:
    print("\n📦 Installing connection dependencies (httpx, websockets)...")
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "--quiet", "httpx", "websockets"])
    except Exception:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "--quiet", "--break-system-packages", "httpx", "websockets"])
    import httpx
    import websockets

try:
    import rumps
    HAS_RUMPS = True
except ImportError:
    HAS_RUMPS = False

DEFAULT_HUB_URL = os.getenv("HUB_WS_URL", "ws://mda.54.175.45.61.nip.io/ws/daemon/mac-local")
DEFAULT_OLLAMA_URL = os.getenv("OLLAMA_URL", "http://localhost:11434")
DEFAULT_LMSTUDIO_URL = os.getenv("LMSTUDIO_URL", "http://localhost:1234/v1")
DEFAULT_AFM_URL = os.getenv("AFM_URL", "http://127.0.0.1:11535/v1")
DASHBOARD_URL = "http://mda.54.175.45.61.nip.io/"

APP_INSTANCE = None
IS_CONNECTED = False
ACTIVE_TASKS = 0
CURRENT_MODELS_LIST = []
STATUS_LOCK = threading.Lock()

DYNAMIC_ALIASES = {}
REVERSE_ALIASES = {}

def update_status(delta_tasks=0, connected=None, models=None):
    global ACTIVE_TASKS, IS_CONNECTED, CURRENT_MODELS_LIST
    with STATUS_LOCK:
        ACTIVE_TASKS += delta_tasks
        if connected is not None:
            IS_CONNECTED = connected
        if models is not None:
            CURRENT_MODELS_LIST = models
            
    if APP_INSTANCE and HAS_RUMPS:
        try:
            if not IS_CONNECTED:
                APP_INSTANCE.title = "🔴 MDA (Offline)"
                if "status_item" in APP_INSTANCE.menu:
                    APP_INSTANCE.menu["status_item"].title = "Status: Disconnected (Reconnecting...)"
            elif ACTIVE_TASKS > 0:
                APP_INSTANCE.title = f"⚡️ MDA ({ACTIVE_TASKS} Active)"
                if "status_item" in APP_INSTANCE.menu:
                    APP_INSTANCE.menu["status_item"].title = f"Status: Running {ACTIVE_TASKS} query..."
            else:
                APP_INSTANCE.title = "🟢 MDA (Bridge Ready)"
                if "status_item" in APP_INSTANCE.menu:
                    APP_INSTANCE.menu["status_item"].title = "Status: Connected to MDA Cloud"
        except Exception:
            pass

def get_alias(name: str) -> str:
    if name in REVERSE_ALIASES:
        return REVERSE_ALIASES[name]
    
    clean_name = name.split('/')[-1] if '/' in name else name
    base_name = clean_name.split(':')[0] if ':' in clean_name else clean_name
    first_char = base_name[0].upper() if base_name else 'X'
    digits = re.findall(r'\d+', base_name)
    if digits:
        alias = first_char + "".join(digits)[:2]
    else:
        consonants = re.sub(r'[^A-Za-z]', '', base_name)[1:]
        consonants = re.sub(r'[aeiouAEIOU]', '', consonants).upper()
        alias = first_char + consonants[:2]
        
    alias = alias[:4]
    counter = 1
    while alias in DYNAMIC_ALIASES and DYNAMIC_ALIASES[alias] != name:
        alias = alias[:2] + str(counter)
        counter += 1
        
    DYNAMIC_ALIASES[alias] = name
    REVERSE_ALIASES[name] = alias
    return alias

def print_banner(hub_url: str):
    banner = f"""\x1b[36m===============================================================
       ⚡ MDA LOCAL HARDWARE MODEL BRIDGE (v3.1)
===============================================================\x1b[0m
\x1b[32m-> Web Hub Target:\x1b[0m    {hub_url}
\x1b[32m-> Local Ollama:\x1b[0m      {DEFAULT_OLLAMA_URL}
\x1b[32m-> Local LM Studio:\x1b[0m   {DEFAULT_LMSTUDIO_URL}
\x1b[32m-> Apple FM (AFM):\x1b[0m    {DEFAULT_AFM_URL} (Port 11535)
\x1b[32m-> Menubar Icon:\x1b[0m      {'Active (macOS)' if HAS_RUMPS else 'CLI Mode'}
\x1b[36m===============================================================\x1b[0m"""
    print(banner, flush=True)

async def scan_local_models(ollama_url: str, lmstudio_url: str, afm_url: str):
    models = ["AFM"]
    mapping = {
        "AFM": "Apple Foundation Model (Apple Silicon On-Device)"
    }

    # 1. Scan Ollama
    try:
        async with httpx.AsyncClient(timeout=2.0) as client:
            res = await client.get(f"{ollama_url.rstrip('/')}/api/tags")
            if res.status_code == 200:
                data = res.json()
                for m in data.get("models", []):
                    full_name = m.get("name", "")
                    alias = get_alias(full_name)
                    if alias not in models:
                        models.append(alias)
                    mapping[alias] = f"Ollama: {full_name}"
    except Exception:
        pass

    # 2. Scan LM Studio
    try:
        async with httpx.AsyncClient(timeout=2.0) as client:
            res = await client.get(f"{lmstudio_url.rstrip('/')}/models")
            if res.status_code == 200:
                data = res.json()
                for m in data.get("data", []):
                    model_id = m.get("id", "")
                    alias = get_alias(model_id)
                    if alias not in models:
                        models.append(alias)
                    mapping[alias] = f"LM Studio: {model_id}"
    except Exception:
        pass

    return models, mapping

async def safe_send(ws, payload_dict: dict) -> bool:
    try:
        await ws.send(json.dumps(payload_dict))
        return True
    except Exception:
        return False

async def stream_ollama(model_name: str, history: list, base_url: str, task_id: str, websocket):
    url = f"{base_url.rstrip('/')}/api/chat"
    payload = {
        "model": model_name,
        "messages": history,
        "stream": True
    }
    print(f"\x1b[34m[Ollama]\x1b[0m Executing query with '{model_name}' (Task: {task_id[:8]})...", flush=True)
    token_count = 0
    try:
        async with httpx.AsyncClient(timeout=180.0) as client:
            async with client.stream("POST", url, json=payload) as response:
                if response.status_code != 200:
                    err = f"[Ollama Error: HTTP {response.status_code}]"
                    await safe_send(websocket, {"type": "stream_chunk", "task_id": task_id, "chunk": err})
                    return

                async for line in response.aiter_lines():
                    if not line:
                        continue
                    try:
                        data = json.loads(line)
                        if "message" in data and "content" in data["message"]:
                            chunk = data["message"]["content"]
                            token_count += 1
                            sent = await safe_send(websocket, {
                                "type": "stream_chunk",
                                "task_id": task_id,
                                "chunk": chunk
                            })
                            if not sent:
                                return
                    except Exception:
                        pass
        print(f"\x1b[32m[SUCCESS]\x1b[0m Streamed {token_count} tokens from '{model_name}'.", flush=True)
        await safe_send(websocket, {"type": "stream_end", "task_id": task_id})
    except Exception as e:
        print(f"\x1b[31m[-] Ollama Error: {str(e)}\x1b[0m", flush=True)
        await safe_send(websocket, {
            "type": "stream_chunk",
            "task_id": task_id,
            "chunk": f"\n[Ollama Error: {str(e)}]\n"
        })
        await safe_send(websocket, {"type": "stream_end", "task_id": task_id})

async def stream_afm_or_openai(model_name: str, history: list, base_url: str, task_id: str, websocket) -> bool:
    url = f"{base_url.rstrip('/')}/chat/completions"
    payload = {
        "model": model_name,
        "messages": history,
        "stream": True
    }
    print(f"\x1b[35m[Apple/AFM]\x1b[0m Executing on-device query with '{model_name}' at {base_url} (Task: {task_id[:8]})...", flush=True)
    token_count = 0
    try:
        async with httpx.AsyncClient(timeout=180.0) as client:
            async with client.stream("POST", url, json=payload) as response:
                if response.status_code != 200:
                    err_text = await response.aread()
                    err = f"[Model Error {response.status_code}: {err_text.decode('utf-8', errors='ignore')}]"
                    await safe_send(websocket, {"type": "stream_chunk", "task_id": task_id, "chunk": err})
                    await safe_send(websocket, {"type": "stream_end", "task_id": task_id})
                    return False

                async for line in response.aiter_lines():
                    if not line:
                        continue
                    if line.startswith("data: "):
                        raw = line[6:].strip()
                        if raw == "[DONE]":
                            break
                        try:
                            parsed = json.loads(raw)
                            choices = parsed.get("choices", [])
                            if choices:
                                delta = choices[0].get("delta", {})
                                chunk = delta.get("content", "")
                                if chunk:
                                    token_count += 1
                                    sent = await safe_send(websocket, {
                                        "type": "stream_chunk",
                                        "task_id": task_id,
                                        "chunk": chunk
                                    })
                                    if not sent:
                                        return False
                        except Exception:
                            pass
        print(f"\x1b[32m[SUCCESS]\x1b[0m Streamed {token_count} tokens from '{model_name}'.", flush=True)
        await safe_send(websocket, {"type": "stream_end", "task_id": task_id})
        return True
    except Exception as e:
        print(f"\x1b[33m[!] AFM endpoint ({base_url}) unreachable: {e}\x1b[0m", flush=True)
        return False

async def handle_afm_execution(history: list, args, task_id: str, websocket):
    """
    Resilient Multi-Metric Apple Foundation Model Handler:
    1. Probes primary AFM URL (default: 11535)
    2. Probes Apple MLX standard port (8080)
    3. Probes LM Studio Apple Metal port (1234)
    4. Routes to local Apple Silicon Metal GPU via Ollama if available
    5. Falls back to clear diagnostic instructions
    """
    candidate_urls = [
        args.afm_url,
        "http://127.0.0.1:8080/v1",
        "http://localhost:8080/v1",
        args.lmstudio_url
    ]
    # Remove duplicates preserving order
    seen = set()
    candidate_urls = [u for u in candidate_urls if u and not (u in seen or seen.add(u))]

    for target_url in candidate_urls:
        success = await stream_afm_or_openai("apple-on-device", history, target_url, task_id, websocket)
        if success:
            return

    # Fallback Metric: Run on Apple Silicon Metal GPU via local Ollama
    try:
        async with httpx.AsyncClient(timeout=1.5) as client:
            r = await client.get(f"{args.ollama_url.rstrip('/')}/api/tags")
            if r.status_code == 200:
                models_data = r.json().get("models", [])
                if models_data:
                    # Pick best available Apple Silicon model
                    chosen_model = models_data[0].get("name", "llama3")
                    print(f"\x1b[36m[Apple Silicon Metal Fallback]\x1b[0m Routing AFM query to local Metal model '{chosen_model}'...", flush=True)
                    await safe_send(websocket, {
                        "type": "stream_chunk",
                        "task_id": task_id,
                        "chunk": f"> *[⚡️ Apple Silicon Acceleration: Executed on local Metal GPU via {chosen_model}]*\n\n"
                    })
                    await stream_ollama(chosen_model, history, args.ollama_url, task_id, websocket)
                    return
    except Exception:
        pass

    # Diagnostic error if no Apple Silicon backend is active
    error_msg = (
        "\n⚠️ **Apple Foundation Model Backend Unavailable**\n\n"
        "No local Apple Foundation Model / MLX server responded on port `11535` or `8080`.\n\n"
        "**To enable Apple MLX on your Mac:**\n"
        "```bash\n"
        "pip install mlx-lm\n"
        "python3 -m mlx_lm.server --model mlx-community/Qwen2.5-7B-Instruct-4bit --port 11535\n"
        "```\n"
    )
    await safe_send(websocket, {
        "type": "stream_chunk",
        "task_id": task_id,
        "chunk": error_msg
    })
    await safe_send(websocket, {"type": "stream_end", "task_id": task_id})

async def handle_execution(data: dict, websocket, args):
    update_status(delta_tasks=1)
    try:
        task_id = data.get("task_id", "")
        model = data.get("model", "llama3")
        prompt = data.get("prompt", "")
        history = data.get("history", [{"role": "user", "content": prompt}])

        actual_model = DYNAMIC_ALIASES.get(model, model)

        if model.upper() == "AFM" or actual_model == "AFM" or "apple" in model.lower():
            await handle_afm_execution(history, args, task_id, websocket)
        else:
            # Check Ollama
            ollama_ok = False
            try:
                async with httpx.AsyncClient(timeout=1.5) as client:
                    r = await client.get(f"{args.ollama_url.rstrip('/')}/api/tags")
                    if r.status_code == 200:
                        ollama_ok = True
            except Exception:
                ollama_ok = False

            if ollama_ok:
                await stream_ollama(actual_model, history, args.ollama_url, task_id, websocket)
            else:
                success = await stream_afm_or_openai(actual_model, history, args.lmstudio_url, task_id, websocket)
                if not success:
                    await safe_send(websocket, {
                        "type": "stream_chunk",
                        "task_id": task_id,
                        "chunk": f"\n[Model Error: Model '{actual_model}' not responding on local bridge]\n"
                    })
                    await safe_send(websocket, {"type": "stream_end", "task_id": task_id})
    finally:
        update_status(delta_tasks=-1)

async def connect_and_listen(args):
    reconnect_delay = 3
    while True:
        try:
            print(f"\n[*] Connecting to MDA Web Hub at {args.hub_url}...", flush=True)
            async with websockets.connect(args.hub_url) as websocket:
                update_status(connected=True)
                print(f"\x1b[32m[+] Connected to MDA Cloud successfully! (ID: {args.daemon_id})\x1b[0m", flush=True)
                
                models, mapping = await scan_local_models(args.ollama_url, args.lmstudio_url, args.afm_url)
                update_status(models=list(mapping.values()))
                print(f"[*] Discovered Local Hardware Models: {list(mapping.values())}", flush=True)
                
                await safe_send(websocket, {
                    "type": "models_update",
                    "models": models,
                    "mapping": mapping
                })
                print("[*] Ready! Live queries from the web dashboard will execute on this machine.\n", flush=True)

                while True:
                    msg_text = await websocket.recv()
                    try:
                        data = json.loads(msg_text)
                    except json.JSONDecodeError:
                        continue

                    msg_type = data.get("type")
                    if msg_type == "execute_local":
                        asyncio.create_task(handle_execution(data, websocket, args))
                    elif msg_type == "ping":
                        await safe_send(websocket, {"type": "pong"})

        except (websockets.exceptions.ConnectionClosed, websockets.exceptions.WebSocketException) as e:
            update_status(connected=False)
            print(f"[-] Disconnected from Web Hub: {e}", flush=True)
        except Exception as e:
            update_status(connected=False)
            print(f"[-] Connection Error: {e}", flush=True)

        print(f"[*] Reconnecting in {reconnect_delay}s... (Press Ctrl+C to stop)", flush=True)
        await asyncio.sleep(reconnect_delay)

class MDAMenubarApp(rumps.App):
    def __init__(self, args):
        super(MDAMenubarApp, self).__init__("⚡️ MDA Bridge", quit_button=None)
        self.args = args
        self.menu = [
            rumps.MenuItem("Status: Connecting...", callback=None, key=None),
            None,
            rumps.MenuItem("🌐 Open Web Dashboard", callback=self.open_dashboard),
            rumps.MenuItem("🔄 Scan Local Models", callback=self.rescan_models),
            None,
            rumps.MenuItem("Quit MDA Bridge", callback=self.quit_app)
        ]
        # Store status reference
        self.menu["Status: Connecting..."].name = "status_item"

    def open_dashboard(self, _):
        webbrowser.open(DASHBOARD_URL)

    def rescan_models(self, _):
        rumps.notification("MDA Bridge", "Scanning Hardware", "Checking local Ollama and Apple Intelligence...")

    def quit_app(self, _):
        rumps.quit_application()
        os._exit(0)

def parse_args():
    parser = argparse.ArgumentParser(description="MDA Local Hardware Model Bridge")
    parser.add_argument("--hub-url", "-u", default=DEFAULT_HUB_URL, help="WebSocket URL of MDA Web Hub")
    parser.add_argument("--daemon-id", "-d", default="mac-local", help="Unique ID for this local daemon")
    parser.add_argument("--ollama-url", default=DEFAULT_OLLAMA_URL, help="Ollama server URL")
    parser.add_argument("--lmstudio-url", default=DEFAULT_LMSTUDIO_URL, help="LM Studio server URL")
    parser.add_argument("--afm-url", default=DEFAULT_AFM_URL, help="Apple Foundation Model URL")
    parser.add_argument("--no-gui", action="store_true", help="Run without macOS menubar GUI (pure CLI)")
    return parser.parse_args()

def main():
    global APP_INSTANCE
    args = parse_args()
    print_banner(args.hub_url)

    # Start asyncio loop in a daemon background thread
    def run_loop():
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            loop.run_until_complete(connect_and_listen(args))
        except Exception as e:
            print(f"[-] Daemon background loop error: {e}")

    bg_thread = threading.Thread(target=run_loop, daemon=True)
    bg_thread.start()

    # If running on macOS and GUI not disabled, launch Menubar app
    if HAS_RUMPS and not args.no_gui and sys.platform == "darwin":
        try:
            APP_INSTANCE = MDAMenubarApp(args)
            APP_INSTANCE.run()
        except Exception as e:
            print(f"[!] Menubar app notice: {e}, continuing in CLI mode.")
            bg_thread.join()
    else:
        try:
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            print("\n[!] Bridge stopped by user.")

if __name__ == "__main__":
    main()