"""
Marziel AI — Embedded LLM Inference

Provides an OpenAI-compatible /v1/chat/completions endpoint embedded
directly in the Flask server. Uses llama-cpp-python (plain, no [server]
extras needed) to load GGUF models.

Works on every platform and Python version:
  - macOS Apple Silicon — Metal GPU acceleration
  - macOS Intel — CPU
  - Linux/Windows NVIDIA — CUDA GPU
  - Linux/Windows AMD — ROCm (if compiled)
  - Any CPU — fallback

This avoids needing vLLM, numba, or llama_cpp.server which all have
Python version restrictions.
"""

import glob
import json
import logging
import os
import threading
import time

logger = logging.getLogger("marziel.llm_embedded")

# ── Global state ──
_model = None
_model_lock = threading.Lock()
_model_id = ""

MODEL_DIR = os.path.expanduser("~/.marziel/models/marziel-8b-custom")


def _get_n_gpu_layers() -> int:
    """Detect GPU and return number of layers to offload."""
    import platform

    # Apple Silicon — Metal is auto-detected by llama.cpp
    if platform.system() == "Darwin" and platform.machine() in ("arm64", "aarch64"):
        return 999

    # NVIDIA GPU
    try:
        import subprocess
        result = subprocess.run(["nvidia-smi"], capture_output=True, timeout=5)
        if result.returncode == 0:
            return 999
    except Exception:
        pass

    # CPU fallback
    return 0


def load_model(model_dir: str = None) -> bool:
    """Load a GGUF model into memory. Thread-safe, loads once."""
    global _model, _model_id

    if _model is not None:
        return True

    with _model_lock:
        if _model is not None:
            return True

        search_dir = model_dir or MODEL_DIR
        gguf_files = glob.glob(os.path.join(search_dir, "**/*.gguf"), recursive=True)
        if not gguf_files:
            logger.warning(f"No GGUF model found in {search_dir}")
            return False

        model_path = gguf_files[0]
        n_gpu = _get_n_gpu_layers()

        try:
            from llama_cpp import Llama
            logger.info(f"Loading model: {os.path.basename(model_path)} "
                       f"(GPU layers: {n_gpu})")
            _model = Llama(
                model_path=model_path,
                n_ctx=4096,
                n_gpu_layers=n_gpu,
                n_threads=os.cpu_count() or 4,
                verbose=False,
            )
            _model_id = os.path.basename(search_dir)
            logger.info(f"✅ Model loaded: {_model_id}")
            return True
        except Exception as e:
            logger.error(f"Failed to load model: {e}")
            return False


def chat_completion(messages: list, temperature: float = 0.7,
                    max_tokens: int = 2048, stop: list = None) -> dict:
    """OpenAI-compatible chat completion using loaded model."""
    if _model is None:
        return {"error": "Model not loaded"}

    try:
        result = _model.create_chat_completion(
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            stop=stop or [],
        )
        return result
    except Exception as e:
        logger.error(f"Inference error: {e}")
        return {"error": str(e)}


def register_routes(app):
    """Register OpenAI-compatible API routes on the Flask app."""
    from flask import request, jsonify

    @app.route("/v1/models", methods=["GET"])
    def v1_models():
        if _model is None:
            return jsonify({"data": []})
        return jsonify({
            "data": [{"id": _model_id, "object": "model",
                      "owned_by": "marziel"}]
        })

    @app.route("/v1/chat/completions", methods=["POST", "OPTIONS"])
    def v1_chat_completions():
        if request.method == "OPTIONS":
            return "", 200

        if _model is None:
            return jsonify({"error": "Model not loaded"}), 503

        data = request.get_json(force=True)
        messages = data.get("messages", [])
        temperature = data.get("temperature", 0.7)
        max_tokens = data.get("max_tokens", 2048)
        stop = data.get("stop", [])

        result = chat_completion(messages, temperature, max_tokens, stop)

        if "error" in result:
            return jsonify(result), 500

        return jsonify(result)

    logger.info("📡 Embedded /v1/chat/completions registered")
