#!/usr/bin/env python3

"""
NEMS AI Copilot CLI
===================
Author: Robbie Ferguson - NEMS Linux
Description: On-device terminal assistant for NEMS configuration,
             troubleshooting, and documentation-assisted guidance.
"""

import os
import sys
import re

# Behavioral Variables
PERSONALITY = "professional but personable"

# System Paths
MODEL_PATH = "/usr/local/share/nems/models/nems-ai.gguf"
DOCS_DIR = "/usr/local/share/nems/docs"  # Provided by 'nems-docs' package


def detect_hardware_settings():
    """Dynamically profile host hardware to tune llama.cpp parameters without external deps."""
    # 1. Core count detection
    logical_cores = os.cpu_count() or 4

    # 2. Total RAM detection (Linux-native)
    try:
        mem_bytes = os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES')
        total_ram_gb = mem_bytes / (1024 ** 3)
    except Exception:
        total_ram_gb = 4.0  # Fallback assumption for standard Pi 4/5

    # Optimal thread count calculation:
    # - On 4-core Pi 5: Use all 4 cores.
    # - On high-core VMs: Cap at 8 threads to prevent thread-contention slowdowns.
    if logical_cores <= 4:
        threads = logical_cores
    else:
        threads = min(8, max(4, logical_cores - 1))

    # Scale context window if host memory allows
    n_ctx = 4096 if total_ram_gb >= 8.0 else 2048

    return {
        "n_threads": threads,
        "n_ctx": n_ctx,
        "n_batch": 512  # Accelerates prompt/context prefill phase
    }


# Capture User Query (CLI argument or piped stdin)
if len(sys.argv) > 1:
    user_query = " ".join(sys.argv[1:])
elif not sys.stdin.isatty():
    user_query = sys.stdin.read().strip()
else:
    print("NEMS AI Copilot")
    print("Usage:   nems-ai \"<your question here>\"")
    print("Example: nems-ai \"How do I configure check_dhtxx in NConf?\"")
    sys.exit(0)

# Verify Model File Existence
if not os.path.exists(MODEL_PATH):
    print("Error: NEMS AI model file not found.", file=sys.stderr)
    print(f"Expected path: {MODEL_PATH}", file=sys.stderr)
    print("Run 'sudo apt install --reinstall nems-ai' to resolve.", file=sys.stderr)
    sys.exit(1)

# Check for llama_cpp Library
try:
    from llama_cpp import Llama
except ImportError:
    print("Error: Python llama_cpp library missing.", file=sys.stderr)
    print("Ensure llama-cpp-python is installed for the system python3 interpreter.", file=sys.stderr)
    sys.exit(1)

# Lightweight Sphinx Document Retriever (Local RAG)
def get_relevant_nems_docs(query, docs_path):
    if not os.path.exists(docs_path):
        return ""

    stop_words = {
        "how", "do", "i", "a", "the", "in", "on", "to", "for", "of",
        "and", "is", "can", "you", "what", "with", "from", "my", "me",
        "check", "who", "where", "when", "why", "which", "tell", "about"
    }

    clean_query = re.sub(r'[^\w\s]', '', query.lower())
    query_words = {word for word in clean_query.split() if word not in stop_words and len(word) > 1}

    if not query_words:
        return ""

    matched_chunks = []

    for root, dirs, files in os.walk(docs_path):
        dirs[:] = [d for d in dirs if not d.startswith('.') and d not in ['_build', '_static', '_templates', 'build']]

        for file in files:
            if file.endswith(('.rst', '.md', '.txt', '.html')) and file != 'conf.py':
                filepath = os.path.join(root, file)
                try:
                    with open(filepath, "r", encoding="utf-8", errors="ignore") as f:
                        content = f.read()
                        content_lower = content.lower()

                        score = sum(content_lower.count(word) for word in query_words)

                        if score > 0:
                            first_idx = min(content_lower.find(word) for word in query_words if word in content_lower)
                            start = max(0, first_idx - 200)
                            end = min(len(content), first_idx + 2300)
                            snippet = content[start:end]

                            matched_chunks.append((score, snippet))
                except Exception:
                    pass

    matched_chunks.sort(key=lambda x: x[0], reverse=True)

    if matched_chunks:
        top_match = matched_chunks[0][1]
        return f"\n--- RELEVANT NEMS DOCUMENTATION ---\n{top_match}\n-----------------------------------\n"

    return ""

# Scan local docs installed by nems-docs package
nems_doc_context = get_relevant_nems_docs(user_query, DOCS_DIR)

# Build System Prompt with Injected Documentation Context
SYSTEM_PROMPT = f"""You are NEMS AI, the built-in system administrator copilot for NEMS Linux (Nagios Enterprise Monitoring Server).
Rules:
1. Provide concise, accurate Linux CLI solutions and Nagios plugin guidance.
2. NEMS check plugins live in /usr/lib/nagios/plugins/ and /usr/local/bin/.
3. In NConf, Nagios macros are formatted like $HOSTADDRESS$, $ARG1$, $ARG2$, $USER1$.
4. Use the provided NEMS documentation snippet if relevant to accurately answer questions about NEMS features, scripts, or plugins.
5. When technical or configuration questions are asked, provide a direct answer with exact NConf command lines or Linux CLI syntax.
6. Your personality is {PERSONALITY}.
7. Plain text output only. Do not include markdown, ascii or other non-text characters in response.
8. Do not use conversational filler, introductions, or greetings.
9. If you don't know the answer to the query, tell me. Do not "guess" at the answer or reply with something unrelated..

{nems_doc_context}"""

prompt = f"<|start_header_id|>system<|end_header_id|>\n\n{SYSTEM_PROMPT}<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n{user_query}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"

# Load Model & Run Streaming On-Device Inference
hw_config = detect_hardware_settings()

try:
    llm = Llama(
        model_path=MODEL_PATH,
        n_ctx=hw_config["n_ctx"],
        n_threads=hw_config["n_threads"],
        n_batch=hw_config["n_batch"],
        verbose=False
    )

    response_stream = llm(
        prompt,
        max_tokens=300,
        temperature=0.1,
        stop=["<|eot_id|>"],
        stream=True  # Streaming output active
    )

    # Stream output token-by-token directly to terminal
    for chunk in response_stream:
        token = chunk["choices"][0]["text"]
        sys.stdout.write(token)
        sys.stdout.flush()

    print()  # Output trailing newline

except Exception as e:
    print(f"\nError during AI inference: {e}", file=sys.stderr)
    sys.exit(1)
