#!/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-v6")
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 _find_omlx():
    """Find oMLX binary or Python with omlx installed."""
    # Check if omlx CLI is available
    omlx_bin = shutil.which("omlx")
    if omlx_bin:
        return omlx_bin, "bin"

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


def _start_mlx_engine():
    """Start oMLX inference server (Apple Silicon) with TurboQuant KV cache."""
    omlx_path, omlx_type = _find_omlx()

    if not omlx_path:
        _print("   📦 Installing oMLX (TurboQuant KV cache)...")
        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", "omlx", "mlx-lm"], check=True,
                       capture_output=True)
        omlx_path = os.path.join(venv_dir, "bin", "python3")
        omlx_type = "python"
        _print("   ✅ oMLX installed")

    # Setup SSD KV cache directory
    kv_cache_dir = os.path.join(MARZIEL_HOME, "kv_cache")
    os.makedirs(kv_cache_dir, exist_ok=True)

    _print(f"🚀 Starting Marziel AI engine (TurboQuant KV)...")
    _print(f"   Model: {MODEL_DIR}")
    _print(f"   Port: {LLM_PORT}")
    _print(f"   KV Cache: {kv_cache_dir}")

    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

    # Build oMLX command
    if omlx_type == "bin":
        cmd = [omlx_path, "serve",
               "--model", MODEL_DIR,
               "--port", str(LLM_PORT),
               "--host", "0.0.0.0"]
    else:
        cmd = [omlx_path, "-m", "omlx", "serve",
               "--model", MODEL_DIR,
               "--port", str(LLM_PORT),
               "--host", "0.0.0.0"]

    log_path = os.path.join(MARZIEL_HOME, "omlx.log")
    subprocess.Popen(
        cmd,
        stdout=subprocess.DEVNULL,
        stderr=open(log_path, "w"),
        start_new_session=True,
        env={**os.environ, "OMLX_KV_CACHE_DIR": kv_cache_dir},
    )

    for i in range(45):
        if _port_active(LLM_PORT):
            _print("   ✅ Engine running (TurboQuant KV cache active)")
            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 _ensure_model_downloaded():
    """Ensure the GGUF model is downloaded (for non-Apple-Silicon platforms)."""
    GGUF_FILE = "marziel-v6-Q4_K_M.gguf"
    gguf_path = os.path.join(MODEL_DIR, GGUF_FILE)

    if os.path.exists(gguf_path):
        return True

    _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!")
        return True
    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!")
        return True
    except Exception as e:
        _print(f"   ❌ Download failed: {e}")
        return False


def _ensure_llama_cpp_installed():
    """Ensure llama-cpp-python is installed (the engine backend)."""
    try:
        from llama_cpp import Llama  # noqa: F401
        return True
    except ImportError:
        _print("   📦 Installing llama-cpp-python...")
        try:
            env = os.environ.copy()
            # Detect GPU for build flags
            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:
                env["CMAKE_ARGS"] = "-DGGML_CUDA=on"
                _print("      (Native CUDA GPU acceleration enabled)")
            elif platform.system() == "Darwin" and platform.machine() in ("arm64", "aarch64"):
                _print("      (Apple Metal GPU acceleration enabled)")

            subprocess.run(
                [sys.executable, "-m", "pip", "install", "llama-cpp-python", "-q"],
                check=True, capture_output=True, timeout=600, env=env
            )
            _print("   ✅ llama-cpp-python installed")
            return True
        except Exception as e:
            _print(f"   ⚠️  llama-cpp-python install failed: {e}")
            return False


def ensure_ready():
    """Ensure model is downloaded and MarzielEngine can load it."""

    # Step 1: Apple Silicon → MLX (keep existing oMLX path)
    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 2: Everything else → MarzielEngine (in-process, no subprocess)
    _print("   🔍 Setting up Marziel Engine...")

    # Ensure GGUF model is downloaded
    if not _ensure_model_downloaded():
        _print("")
        _print("   ❌ Could not download model.")
        _print("   Please check your internet connection and try again.")
        _print("")
        return None

    # Ensure llama-cpp-python is installed
    if not _ensure_llama_cpp_installed():
        _print("")
        _print("   ❌ Could not install inference backend.")
        _print("")
        return None

    # Pre-load the engine (will be used in-process by server.py)
    try:
        from marziel.engine import get_engine
        engine = get_engine(model_dir=MODEL_DIR)
        _print(f"🚀 Starting Marziel Engine ({engine.hardware.backend})...")
        _print(f"   Model: {MODEL_DIR}")
        _print(f"   Hardware: {engine.hardware.gpu_name or 'CPU'}")
        if engine.load():
            _print(f"   ✅ Marziel Engine ready ({engine.hardware.backend})")
            os.environ["MARZIEL_ACTIVE_ENGINE"] = "marziel"
            return "marziel"
        else:
            _print(f"   ⚠️  Engine load deferred — will load on first request")
            os.environ["MARZIEL_ACTIVE_ENGINE"] = "marziel"
            return "marziel"
    except Exception as e:
        _print(f"   ⚠️  Engine pre-load failed: {e}")
        _print("   Engine will load on first request.")
        os.environ["MARZIEL_ACTIVE_ENGINE"] = "marziel"
        return "marziel"




# ══════════════════════════════════════════════════════════
# 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 v6 (already running)",
        "mlx": "Marziel v6 (Apple Silicon MLX)",
        "marziel": "Marziel v6 (Marziel Engine)",
    }

    _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("")

    # Auto-install semantic router for intent detection
    try:
        import semantic_router  # noqa: F401
    except ImportError:
        _print("📦 Installing semantic router...")
        try:
            subprocess.run(
                [sys.executable, "-m", "pip", "install", "semantic-router", "fastembed", "-q"],
                check=True, capture_output=True, timeout=120
            )
            _print("   ✅ Semantic router installed")
        except Exception:
            _print("   ⚠️  Semantic router install skipped (keyword fallback active)")

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

