#!/usr/bin/env python3
"""
Marziel CLI — Complete AI system.

`marziel serve` does everything:
  1. Downloads model from Hugging Face (first run only)
  2. Starts inference engine (MLX on Apple Silicon, vLLM elsewhere)
  3. Starts web server with UI

Zero manual steps.
"""

import argparse
import json
import os
import platform
import shutil
import subprocess
import sys
import time
import urllib.request

# ── Config ──
HF_REPO = os.environ.get("MARZIEL_HF_REPO", "efops/marziel-8b-custom")
MARZIEL_HOME = os.path.expanduser("~/.marziel")
MODEL_DIR = os.path.join(MARZIEL_HOME, "models", "marziel-8b-custom")
LLM_PORT = int(os.environ.get("LLM_PORT", "8000"))


def _print(msg):
    print(f"  {msg}")


# ══════════════════════════════════════════════════════════
# MODEL MANAGEMENT
# ══════════════════════════════════════════════════════════

def _model_exists():
    """Check if model is already downloaded."""
    return (
        os.path.isdir(MODEL_DIR)
        and os.path.exists(os.path.join(MODEL_DIR, "config.json"))
        and os.path.exists(os.path.join(MODEL_DIR, "model.safetensors"))
    )


def _download_model():
    """Download the Marziel model from Hugging Face (public model)."""
    # Try to find a HuggingFace token (optional, speeds up downloads)
    hf_token = None

    # 1. License file (Enterprise/Pro users may have one)
    license_file = os.path.expanduser("~/.marziel/license.json")
    try:
        if os.path.exists(license_file):
            import json as _json
            lic = _json.loads(open(license_file).read())
            hf_token = lic.get("hf_token", "") or None
    except Exception:
        pass

    # 2. Standard HuggingFace cached token (from `huggingface-cli login`)
    if not hf_token:
        hf_token_file = os.path.expanduser("~/.cache/huggingface/token")
        try:
            if os.path.exists(hf_token_file):
                hf_token = open(hf_token_file).read().strip() or None
        except Exception:
            pass

    # 3. Environment variable
    if not hf_token:
        hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")

    _print(f"📥 Downloading Marziel AI model...")
    _print(f"   Source: huggingface.co/{HF_REPO}")
    _print(f"   Destination: {MODEL_DIR}")
    _print(f"   Size: ~4.2 GB (one-time download)")
    _print(f"   Your model stays 100% local after download.")
    _print("")

    os.makedirs(MODEL_DIR, exist_ok=True)

    try:
        from huggingface_hub import snapshot_download

        snapshot_download(
            repo_id=HF_REPO,
            local_dir=MODEL_DIR,
            local_dir_use_symlinks=False,
            token=hf_token if hf_token else None,
        )
        _print("   ✅ Model downloaded successfully!")
        return True

    except Exception as e:
        error_msg = str(e)
        _print(f"   ❌ Download failed: {e}")
        _print("   Check your internet connection and try again.")
        return False


# ══════════════════════════════════════════════════════════
# ENGINE MANAGEMENT
# ══════════════════════════════════════════════════════════

def _port_active(port):
    """Check if an LLM server is running on a port."""
    try:
        req = urllib.request.Request(f"http://localhost:{port}/v1/models")
        resp = urllib.request.urlopen(req, timeout=3)
        return resp.status == 200
    except Exception:
        return False


def _find_mlx_python():
    """Find a Python with mlx_lm installed."""
    candidates = [
        os.path.expanduser("~/.venv-mlx/bin/python3"),
        os.path.join(MARZIEL_HOME, "venv/bin/python3"),
        shutil.which("python3"),
    ]
    for py in candidates:
        if py and os.path.exists(py):
            try:
                result = subprocess.run(
                    [py, "-c", "import mlx_lm; print('ok')"],
                    capture_output=True, text=True, timeout=5
                )
                if result.returncode == 0:
                    return py
            except Exception:
                continue
    return None


def _start_mlx_engine():
    """Start MLX inference server (Apple Silicon)."""
    python = _find_mlx_python()

    if not python:
        _print("   ⚠️  mlx_lm not found. Installing...")
        venv_dir = os.path.join(MARZIEL_HOME, "venv")
        if not os.path.exists(venv_dir):
            subprocess.run([sys.executable, "-m", "venv", venv_dir], check=True)
        pip = os.path.join(venv_dir, "bin", "pip")
        subprocess.run([pip, "install", "mlx-lm"], check=True)
        python = os.path.join(venv_dir, "bin", "python3")

    _print(f"🚀 Starting Marziel AI engine...")
    _print(f"   Model: marziel-8b-custom")
    _print(f"   Port: {LLM_PORT}")

    try:
        subprocess.run(
            f"lsof -ti:{LLM_PORT} | xargs kill -9 2>/dev/null || true",
            shell=True, capture_output=True
        )
        time.sleep(1)
    except Exception:
        pass

    subprocess.Popen(
        [python, "-m", "mlx_lm.server",
         "--model", MODEL_DIR,
         "--port", str(LLM_PORT),
         "--host", "0.0.0.0"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        start_new_session=True,
    )

    for i in range(30):
        if _port_active(LLM_PORT):
            _print("   ✅ Engine running")
            return True
        time.sleep(1)
        if i % 5 == 4:
            _print(f"   ⏳ Loading model... ({i+1}s)")

    _print("   ⚠️  Engine started — model still loading, will be ready shortly")
    return True


def _start_vllm():
    """Start vLLM inference server (non-Apple-Silicon)."""
    GGUF_FILE = "marziel-8b-custom.gguf"
    gguf_path = os.path.join(MODEL_DIR, GGUF_FILE)

    # Step 1: Ensure GGUF model is downloaded
    if not os.path.exists(gguf_path):
        _print("📥 Downloading Marziel AI model (GGUF)...")
        _print(f"   Source: huggingface.co/{HF_REPO}")
        _print(f"   Size: ~4.6 GB (one-time download)")
        _print(f"   Your model stays 100% local after download.")
        _print("")
        os.makedirs(MODEL_DIR, exist_ok=True)
        try:
            from huggingface_hub import hf_hub_download
            hf_hub_download(
                repo_id=HF_REPO,
                filename=GGUF_FILE,
                local_dir=MODEL_DIR,
                local_dir_use_symlinks=False,
            )
            _print("   ✅ Model downloaded!")
        except ImportError:
            _print("   📦 Installing huggingface_hub...")
            subprocess.run(
                [sys.executable, "-m", "pip", "install", "huggingface_hub", "-q"],
                check=True, capture_output=True, timeout=120
            )
            from huggingface_hub import hf_hub_download
            hf_hub_download(
                repo_id=HF_REPO,
                filename=GGUF_FILE,
                local_dir=MODEL_DIR,
                local_dir_use_symlinks=False,
            )
            _print("   ✅ Model downloaded!")
        except Exception as e:
            _print(f"   ❌ Download failed: {e}")
            return False

    # Step 2: Ensure vLLM is installed
    try:
        import vllm  # noqa: F401
    except ImportError:
        # Detect GPU to decide which variant to install
        has_gpu = False
        try:
            result = subprocess.run(["nvidia-smi"], capture_output=True, timeout=5)
            has_gpu = result.returncode == 0
        except Exception:
            pass

        if has_gpu:
            _print("   📦 Installing vLLM (GPU)...")
            try:
                subprocess.run(
                    [sys.executable, "-m", "pip", "install", "vllm", "-q"],
                    check=True, capture_output=True, timeout=600
                )
            except Exception as e:
                _print(f"   ⚠️  vLLM install failed: {e}")
                return False
        else:
            _print("   📦 Installing vLLM (CPU)...")
            try:
                import json as _json
                import urllib.request as _urlreq
                try:
                    resp = _urlreq.urlopen(
                        "https://api.github.com/repos/vllm-project/vllm/releases/latest",
                        timeout=10
                    )
                    vllm_ver = _json.loads(resp.read().decode()).get("tag_name", "v0.17.1").lstrip("v")
                except Exception:
                    vllm_ver = "0.17.1"

                import platform as _plat
                arch = _plat.machine()
                wheel_arch = "aarch64" if arch in ("aarch64", "arm64") else "x86_64"

                wheel_url = (
                    f"https://github.com/vllm-project/vllm/releases/download/"
                    f"v{vllm_ver}/vllm-{vllm_ver}+cpu-cp38-abi3-manylinux_2_35_{wheel_arch}.whl"
                )
                _print(f"   📥 Wheel: vllm-{vllm_ver}+cpu ({wheel_arch})")
                subprocess.run(
                    [sys.executable, "-m", "pip", "install", wheel_url,
                     "--extra-index-url", "https://download.pytorch.org/whl/cpu", "-q"],
                    check=True, capture_output=True, timeout=600
                )
                # Install TCMalloc for CPU performance
                subprocess.run(
                    ["sudo", "apt-get", "install", "-y", "-qq", "libtcmalloc-minimal4"],
                    capture_output=True, timeout=30
                )
            except Exception as e:
                _print(f"   ⚠️  vLLM CPU install failed: {e}")
                return False

    # Detect GPU for runtime config
    has_gpu = False
    try:
        result = subprocess.run(["nvidia-smi"], capture_output=True, timeout=5)
        has_gpu = result.returncode == 0
    except Exception:
        pass

    gpu_label = " (GPU)" if has_gpu else " (CPU)"

    # Step 4: Start vLLM server
    # CPU: use HF repo (safetensors format) — vLLM CPU lacks GGML dequantize ops
    # GPU: use local GGUF file for faster loading
    model_ref = HF_REPO if not has_gpu else gguf_path
    _print(f"🚀 Starting vLLM engine...")
    _print(f"   Model: {model_ref}")
    _print(f"   Port: {LLM_PORT}")
    _print(f"   Mode: {gpu_label.strip()}")

    try:
        subprocess.run(
            f"lsof -ti:{LLM_PORT} | xargs kill -9 2>/dev/null || true",
            shell=True, capture_output=True
        )
        time.sleep(1)
    except Exception:
        pass

    cmd = [
        sys.executable, "-m", "vllm.entrypoints.openai.api_server",
        "--model", model_ref,
        "--port", str(LLM_PORT),
        "--host", "0.0.0.0",
        "--max-model-len", "2048",
        "--enforce-eager",
        "--enable-prefix-caching",
        "--trust-remote-code",
    ]

    # GPU needs separate tokenizer flag (for GGUF models)
    if has_gpu:
        cmd.extend(["--tokenizer", HF_REPO])

    env = os.environ.copy()

    if has_gpu:
        cmd.extend(["--swap-space", "2"])
    else:
        # CPU-specific: bfloat16 recommended by vLLM docs, KV cache space, reserve CPU
        cmd.extend(["--dtype", "bfloat16"])
        env["VLLM_CPU_KVCACHE_SPACE"] = "4"
        env["VLLM_CPU_NUM_OF_RESERVED_CPU"] = "1"
        # Set LD_PRELOAD for TCMalloc + Intel OpenMP if available
        preloads = []
        for lib in ["/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4",
                     "/usr/lib/aarch64-linux-gnu/libtcmalloc_minimal.so.4"]:
            if os.path.exists(lib):
                preloads.append(lib)
        # Intel OpenMP from vLLM CPU package
        try:
            import site
            for sp in site.getsitepackages():
                iomp = os.path.join(sp, "vllm", "libiomp5.so")
                if os.path.exists(iomp):
                    preloads.append(iomp)
        except Exception:
            pass
        if preloads:
            env["LD_PRELOAD"] = ":".join(preloads)

    subprocess.Popen(
        cmd,
        stdout=subprocess.DEVNULL,
        stderr=open(os.path.expanduser("~/.marziel/vllm.log"), "w"),
        start_new_session=True,
        env=env,
    )

    timeout = 180 if not has_gpu else 60
    for i in range(timeout):
        if _port_active(LLM_PORT):
            _print(f"   ✅ vLLM engine running{gpu_label}")
            return True
        time.sleep(1)
        if i % 10 == 9:
            _print(f"   ⏳ Loading model... ({i+1}s)")

    _print("   ⚠️  vLLM started — model still loading, will be ready shortly")
    return True


def ensure_ready():
    """Ensure model is downloaded and engine is running."""

    # Step 1: engine already running?
    if _port_active(LLM_PORT):
        _print(f"✅ AI engine already running (port {LLM_PORT})")
        return "active"

    # Step 2: Apple Silicon → MLX
    if platform.machine() in ("arm64", "aarch64") and platform.system() == "Darwin":
        if not _model_exists():
            if not _download_model():
                _print("❌ Cannot start without the Marziel model.")
                return None
        _start_mlx_engine()
        return "mlx"

    # Step 3: Everything else → vLLM
    _print("   🔍 Setting up vLLM inference engine...")
    if _start_vllm():
        return "vllm"

    _print("")
    _print("   ❌ Could not start inference engine.")
    _print("   Please check your internet connection and try again.")
    _print("")
    return None

# ══════════════════════════════════════════════════════════
# CLI COMMANDS
# ══════════════════════════════════════════════════════════

def cmd_serve(args):
    """Start the complete Marziel AI system."""
    port = args.port or int(os.environ.get("MARZIEL_PORT", "8001"))
    host = args.host or os.environ.get("MARZIEL_HOST", "0.0.0.0")

    print()
    print("  ╔══════════════════════════════════════════════════╗")
    print("  ║  🧠 Marziel AI — Private Intelligence Engine    ║")
    print("  ╚══════════════════════════════════════════════════╝")
    print()

    # Auto-setup everything
    engine = ensure_ready()
    _print("")

    if engine is None:
        _print("   Exiting. Fix the issue above and try again.")
        sys.exit(1)

    engine_labels = {
        "active": "Marziel 8B Custom (already running)",
        "mlx": "Marziel 8B Custom (Apple Silicon MLX)",
        "vllm": "Marziel 8B Custom (vLLM)",
    }

    _print("━━━ System Ready ━━━")
    _print("")
    _print(f"🌐 Dashboard: http://localhost:{port}")
    _print(f"🔌 API:       http://localhost:{port}/llm/chat")
    _print(f"🛡️  Sentinel:  Active")
    _print(f"🧬 Neurons:   8 synapse types")
    _print(f"🧠 LLM:       {engine_labels.get(engine, engine)}")
    _print("")

    from marziel.server import create_app
    app = create_app()

    if args.debug:
        app.run(host=host, port=port, debug=True)
    else:
        try:
            from gunicorn.app.base import BaseApplication

            class MarzielApp(BaseApplication):
                def __init__(self, app, options=None):
                    self.options = options or {}
                    self.application = app
                    super().__init__()

                def load_config(self):
                    for key, value in self.options.items():
                        if key in self.cfg.settings and value is not None:
                            self.cfg.set(key.lower(), value)

                def load(self):
                    return self.application

            options = {
                "bind": f"{host}:{port}",
                "workers": 2,
                "timeout": 120,
                "accesslog": "-",
            }
            MarzielApp(app, options).run()
        except ImportError:
            app.run(host=host, port=port, debug=False)


def cmd_status(args):
    """Check all system components."""
    port = args.port or 8001
    from marziel import __version__
    _print(f"Marziel AI v{__version__}")
    _print("")

    # License status removed

    try:
        req = urllib.request.Request(f"http://localhost:{port}/health")
        urllib.request.urlopen(req, timeout=3)
        _print(f"✅ Server (port {port})")
    except Exception:
        _print(f"❌ Server (port {port})")

    if _port_active(LLM_PORT):
        _print(f"✅ AI Engine (port {LLM_PORT})")
    else:
        _print(f"⏸️  AI Engine (idle — auto-restarts on next request)")

    if _model_exists():
        _print(f"✅ Model: marziel-8b-custom ({MODEL_DIR})")
    else:
        _print(f"❌ Model not downloaded — run 'marziel serve'")

    if _ollama_running():
        _print("✅ Ollama (port 11434)")
    elif shutil.which("ollama"):
        _print("⏸️  Ollama installed (not running)")

    # Resource info
    _print("")
    _print("━━━ Resources ━━━")
    try:
        if platform.system() == "Darwin":
            result = subprocess.run(
                ["sysctl", "-n", "hw.memsize"],
                capture_output=True, text=True, timeout=3
            )
            total_gb = round(int(result.stdout.strip()) / 1024**3, 1)
            _print(f"💾 System RAM: {total_gb} GB")
        # Engine RAM
        if _port_active(LLM_PORT):
            result = subprocess.run(
                f"lsof -ti:{LLM_PORT} | head -1",
                shell=True, capture_output=True, text=True, timeout=3
            )
            pid = result.stdout.strip()
            if pid:
                result = subprocess.run(
                    ["ps", "-o", "rss=", "-p", pid],
                    capture_output=True, text=True, timeout=3
                )
                rss_mb = round(int(result.stdout.strip()) / 1024)
                _print(f"🧠 Engine RAM: {rss_mb} MB")
        idle_timeout = int(os.environ.get("MARZIEL_IDLE_TIMEOUT", "600"))
        _print(f"♻️  Idle shutdown: {idle_timeout // 60} min (saves RAM when not in use)")
    except Exception:
        pass



def cmd_version(args):
    """Show version information."""
    from marziel import __version__
    _print(f"Marziel AI v{__version__}")
    _print(f"Python {sys.version.split()[0]}")
    _print(f"Platform: {sys.platform} ({platform.machine()})")
    _print(f"Model repo: {HF_REPO}")
    _print(f"Model dir: {MODEL_DIR}")
    _print(f"Model ready: {'Yes' if _model_exists() else 'No'}")


def main():
    parser = argparse.ArgumentParser(
        prog="marziel",
        description="🧠 Marziel AI — Private Intelligence, Locally Powered",
    )
    sub = parser.add_subparsers(dest="command")

    p_serve = sub.add_parser("serve", help="Start Marziel (auto-downloads model)")
    p_serve.add_argument("--port", type=int, default=None, help="Server port (default: 8001)")
    p_serve.add_argument("--host", default=None, help="Host (default: 0.0.0.0)")
    p_serve.add_argument("--debug", action="store_true", help="Debug mode")

    p_status = sub.add_parser("status", help="Check system status")
    p_status.add_argument("--port", type=int, default=None)

    sub.add_parser("version", help="Show version info")

    args = parser.parse_args()

    if args.command == "serve":
        cmd_serve(args)
    elif args.command == "status":
        cmd_status(args)
    elif args.command == "version":
        cmd_version(args)
    else:
        parser.print_help()
        print("\n  Quick start:")
        print("    pip install marziel")
        print("    marziel serve")
        print()


if __name__ == "__main__":
    main()

