#!/usr/bin/env python3
"""
Donjon GRC Platform - AI Setup Wizard

Configures the Donjon AI engine. AI analysis is a Pro feature; Community
Edition uses the built-in template backend (no LLM required). Donjon never
ships or hosts a model -- you point it at your own:

  * Ollama          -- local, air-gap friendly (this script can pull models)
  * Custom          -- ANY OpenAI-compatible base_url (llama.cpp, vLLM,
                       LM Studio, LiteLLM, OpenRouter, a private gateway...)
  * Generic         -- universal adapter for ANY HTTP endpoint via a
                       request_template ({prompt}/{system}) + response_path
  * Cloud           -- bring-your-own key (Anthropic, OpenAI, Gemini, ...)
  * Template        -- offline rule-based fallback, always available

Air-gap mode: set DONJON_OFFLINE=1 (env). The wizard then offers ONLY the
backends that perform no public egress (ollama / custom / generic / template)
and refuses any endpoint that is not loopback/RFC1918-private.

Usage:
    python bin/setup-ai.py                         # Interactive wizard
    python bin/setup-ai.py --models mistral:7b     # Pull specific Ollama model(s)
    python bin/setup-ai.py --skip-pull             # Skip downloads, configure only
    python bin/setup-ai.py --test-only             # Just test model availability
    python bin/setup-ai.py --backend custom \
        --api-url http://127.0.0.1:8000/v1/chat/completions  # Non-interactive BYO
    python bin/setup-ai.py --backend generic \
        --api-url http://127.0.0.1:9000/infer \
        --request-template '{"prompt":"{prompt}"}' \
        --response-path 'choices[0].text'           # Universal adapter
"""

import sys
import os
import json
import argparse
import urllib.request
import urllib.error
from pathlib import Path

# ---------------------------------------------------------------------------
# Resolve project paths
# ---------------------------------------------------------------------------
SCRIPT_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = SCRIPT_DIR.parent
DATA_DIR = PROJECT_ROOT / 'data'
AI_CONFIG_PATH = DATA_DIR / 'ai_config.json'

sys.path.insert(0, str(PROJECT_ROOT / 'lib'))

OLLAMA_BASE_URL = "http://localhost:11434"

# ---------------------------------------------------------------------------
# Backend metadata + airgap helpers (pure, import-light, unit-testable).
#
# We mirror the canonical AIEngine.get_supported_backends() at runtime so the
# wizard never drifts from the engine, but every helper below also accepts an
# injected backend list so they can be tested without importing the engine.
# ---------------------------------------------------------------------------

# Fallback used only if lib.ai_engine cannot be imported (keeps the wizard
# usable in a stripped checkout). Mirrors the engine's airgap/metered flags.
_FALLBACK_BACKENDS = [
    {'id': 'ollama', 'label': 'Ollama (Local)', 'airgap': True, 'metered': False,
     'requires': 'OLLAMA_URL or localhost:11434'},
    {'id': 'custom', 'label': 'Custom (OpenAI-compatible URL)', 'airgap': True, 'metered': False,
     'requires': 'api_url (+ api_key if the endpoint needs one)'},
    {'id': 'generic', 'label': 'Generic HTTP (any endpoint)', 'airgap': True, 'metered': False,
     'requires': 'api_url + request_template + response_path'},
    {'id': 'anthropic', 'label': 'Anthropic Claude (Direct)', 'airgap': False, 'metered': True,
     'requires': 'ANTHROPIC_API_KEY'},
    {'id': 'openai', 'label': 'OpenAI', 'airgap': False, 'metered': True,
     'requires': 'OPENAI_API_KEY'},
    {'id': 'gemini', 'label': 'Google Gemini', 'airgap': False, 'metered': True,
     'requires': 'GEMINI_API_KEY'},
    {'id': 'template', 'label': 'Template Fallback (No LLM)', 'airgap': True, 'metered': False,
     'requires': 'Nothing'},
]

# Backends whose configuration prompts for a base URL / endpoint.
_ENDPOINT_BACKENDS = frozenset({'ollama', 'custom', 'generic', 'openai', 'azure'})
# Backends that need a bring-your-own API key.
_KEY_BACKENDS = frozenset({
    'anthropic', 'openai', 'gemini', 'stepfun', 'mistral', 'cohere', 'xai',
    'deepseek', 'together', 'groq', 'fireworks', 'azure', 'openrouter',
})


def is_offline_mode(env=None) -> bool:
    """True when DONJON_OFFLINE forces air-gapped operation (no public egress).

    Mirrors lib.ai_engine._is_offline so the wizard and engine agree.
    """
    env = os.environ if env is None else env
    return str(env.get('DONJON_OFFLINE', '')).strip().lower() in ('1', 'true', 'yes', 'on')


def load_supported_backends():
    """Return the canonical backend list from AIEngine, or a static fallback.

    Kept side-effect free apart from the import attempt so callers/tests can
    rely on it returning a list of {id,label,airgap,metered,requires,...} dicts.
    """
    try:
        from ai_engine import AIEngine  # type: ignore
        backends = AIEngine.get_supported_backends()
        if backends:
            return backends
    except Exception:
        pass
    return list(_FALLBACK_BACKENDS)


def filter_backends_for_mode(backends, offline: bool):
    """Filter a backend list for the current mode.

    Offline (air-gapped): keep ONLY backends flagged airgap=True (ollama,
    custom, generic, template) -- everything that egresses to a public cloud
    is removed so it is never even offered. Online: return all backends.
    """
    if not offline:
        return list(backends)
    return [b for b in backends if b.get('airgap')]


def backend_ids(backends):
    """Return just the ordered list of backend id strings."""
    return [b.get('id') for b in backends if b.get('id')]


def backend_choice_is_allowed(backend_id: str, backends, offline: bool) -> bool:
    """True if *backend_id* is selectable given the (filtered) mode.

    Used to validate a non-interactive --backend against air-gap policy.
    """
    allowed = backend_ids(filter_backends_for_mode(backends, offline))
    return backend_id in allowed


def backend_requires_endpoint(backend_id: str) -> bool:
    """True if this backend prompts for a base URL / endpoint."""
    return backend_id in _ENDPOINT_BACKENDS


def backend_requires_key(backend_id: str) -> bool:
    """True if this backend needs a bring-your-own API key."""
    return backend_id in _KEY_BACKENDS


def endpoint_is_airgap_safe(url: str) -> bool:
    """True if *url* is loopback / RFC1918-private / clearly-internal.

    Delegates to the engine's DNS-free check so the wizard enforces exactly the
    same rule the engine enforces at call time. Falls back to a conservative
    local-only heuristic if the engine import is unavailable. Empty/None URLs
    are treated as unsafe (cannot prove they are local).
    """
    if not url:
        return False
    try:
        from ai_engine import _is_private_endpoint  # type: ignore
        return bool(_is_private_endpoint(url))
    except Exception:
        from urllib.parse import urlparse
        try:
            host = (urlparse(url).hostname or '').strip().lower()
        except Exception:
            return False
        if not host:
            return False
        if host == 'localhost' or host.startswith(('127.', '10.', '192.168.')):
            return True
        if host.endswith(('.local', '.internal', '.lan', '.home', '.corp', '.intra')):
            return True
        if host.startswith('172.'):
            try:
                second = int(host.split('.')[1])
                return 16 <= second <= 31
            except (IndexError, ValueError):
                return False
        return False


def build_backend_config(backend_id, *, model=None, api_url=None, api_key=None,
                         request_template=None, response_path=None, headers=None,
                         ollama_url=None):
    """Assemble a persistable AI config dict for *backend_id*.

    Pure: no I/O, no env reads. Only includes keys relevant to the backend so
    the saved ai_config.json stays clean and round-trips through
    AIEngine._apply_config. Returns the dict (caller persists it).
    """
    cfg = {'backend': backend_id}
    if backend_id == 'ollama':
        cfg['ollama_url'] = ollama_url or api_url or OLLAMA_BASE_URL
        if model:
            cfg['ollama_model'] = model
            cfg['model'] = model
    elif backend_id == 'custom':
        cfg['api_url'] = api_url or ''
        if api_key:
            cfg['api_key'] = api_key
        cfg['model'] = model or 'default'
    elif backend_id == 'generic':
        cfg['api_url'] = api_url or ''
        if api_key:
            cfg['api_key'] = api_key
        if request_template:
            cfg['request_template'] = request_template
        if response_path:
            cfg['response_path'] = response_path
        if headers:
            cfg['headers'] = headers
        cfg['model'] = model or 'generic'
    elif backend_id == 'template':
        pass  # template needs nothing
    else:
        # Cloud / OpenAI-compatible BYO-key backends.
        if api_key:
            cfg['api_key'] = api_key
        if api_url:
            cfg['api_url'] = api_url
        if model:
            cfg['model'] = model
    return cfg


def validate_backend_config(cfg, offline: bool):
    """Validate a config dict. Returns (ok: bool, error: str|None).

    Enforces required fields per backend and the air-gap endpoint rule. Pure.
    """
    backend = cfg.get('backend', '')
    if not backend:
        return False, "no backend selected"

    if backend in ('custom', 'generic') and not cfg.get('api_url'):
        return False, f"backend '{backend}' requires an api_url"
    if backend == 'generic' and not (cfg.get('request_template') and cfg.get('response_path')):
        # Both default sensibly in the engine, but the wizard should capture them.
        if not cfg.get('request_template'):
            return False, "generic backend requires a request_template (with {prompt})"
        if not cfg.get('response_path'):
            return False, "generic backend requires a response_path (e.g. choices[0].message.content)"
    if backend == 'generic' and cfg.get('request_template') and '{prompt}' not in cfg['request_template']:
        return False, "request_template must contain the {prompt} placeholder"

    if offline:
        # Only airgap-safe backends are permitted, and their endpoint must be local.
        if backend not in ('ollama', 'custom', 'generic', 'template'):
            return False, (f"DONJON_OFFLINE: backend '{backend}' performs public egress "
                           f"and is not allowed in air-gapped mode")
        endpoint = cfg.get('ollama_url') or cfg.get('api_url')
        if backend in ('ollama', 'custom', 'generic'):
            if not endpoint_is_airgap_safe(endpoint or ''):
                return False, (f"DONJON_OFFLINE: endpoint '{endpoint}' is not loopback/"
                               f"private; air-gapped mode forbids public endpoints")
    return True, None

# Default models to pull (practical for most hardware: 8-16GB VRAM)
DEFAULT_MODELS = [
    'mistral:7b',
    'codellama:7b',
    'llama3.2:3b',
]

# Step 3.5 Flash: 196B MoE, 11B active params, 111.5GB GGUF Q4
# Requires 120GB+ VRAM (DGX-Spark, Mac Studio M4 Max 128GB, multi-GPU)
# For most users, use the cloud API instead: export STEPFUN_API_KEY=...
STEP35_GGUF_URL = "https://huggingface.co/stepfun-ai/Step-3.5-Flash-GGUF-Q4_K_S"
STEP35_GGUF_SIZE_GB = 111.5

# Custom Modelfile name for the security-tuned model
CUSTOM_MODEL_NAME = 'donjon-security'


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def banner():
    print()
    print("  ============================================================")
    print("   Donjon GRC Platform - AI Model Setup")
    print("  ============================================================")
    print()


def ollama_is_running() -> bool:
    """Check if Ollama is reachable at the default endpoint."""
    try:
        req = urllib.request.Request(
            f"{OLLAMA_BASE_URL}/api/tags",
            headers={'User-Agent': 'DonjonGRC/1.0'},
            method='GET',
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            data = json.loads(resp.read().decode('utf-8'))
            return 'models' in data
    except Exception:
        return False


def get_installed_models() -> list:
    """Return a list of model names currently available in Ollama."""
    try:
        req = urllib.request.Request(
            f"{OLLAMA_BASE_URL}/api/tags",
            headers={'User-Agent': 'DonjonGRC/1.0'},
            method='GET',
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            data = json.loads(resp.read().decode('utf-8'))
            return [m.get('name', '') for m in data.get('models', [])]
    except Exception:
        return []


def pull_model(model_name: str) -> bool:
    """
    Pull a model via Ollama's /api/pull endpoint.

    Streams the response and prints progress.
    """
    print(f"\n  Pulling model: {model_name}")
    print(f"  This may take several minutes depending on model size and connection speed...")

    url = f"{OLLAMA_BASE_URL}/api/pull"
    payload = json.dumps({'name': model_name, 'stream': True}).encode('utf-8')

    req = urllib.request.Request(
        url,
        data=payload,
        headers={
            'Content-Type': 'application/json',
            'User-Agent': 'DonjonGRC/1.0',
        },
        method='POST',
    )

    try:
        with urllib.request.urlopen(req, timeout=1800) as resp:
            last_status = ''
            for line in resp:
                if not line.strip():
                    continue
                try:
                    obj = json.loads(line.decode('utf-8'))
                    status = obj.get('status', '')

                    # Progress display
                    if 'total' in obj and 'completed' in obj:
                        total = obj['total']
                        completed = obj['completed']
                        if total > 0:
                            pct = completed * 100 // total
                            mb_done = completed / (1024 * 1024)
                            mb_total = total / (1024 * 1024)
                            print(
                                f"\r  {status}: {pct}% "
                                f"({mb_done:.0f}/{mb_total:.0f} MB)",
                                end='', flush=True,
                            )
                    elif status and status != last_status:
                        print(f"\n  {status}", end='', flush=True)
                        last_status = status

                    if obj.get('error'):
                        print(f"\n  ERROR: {obj['error']}")
                        return False

                except json.JSONDecodeError:
                    continue

        print(f"\n  Model {model_name} pulled successfully.")
        return True

    except urllib.error.URLError as exc:
        print(f"\n  FAILED to pull {model_name}: {exc}")
        return False
    except Exception as exc:
        print(f"\n  FAILED to pull {model_name}: {exc}")
        return False


def create_custom_modelfile(base_model: str) -> bool:
    """
    Create a custom Ollama model with a security-focused system prompt
    using the Modelfile template from ai_prompts.
    """
    print(f"\n  Creating custom model '{CUSTOM_MODEL_NAME}' from {base_model}...")

    # Import the template
    try:
        from ai_prompts import MODELFILE_TEMPLATE
    except ImportError:
        print("  WARNING: Could not import ai_prompts.MODELFILE_TEMPLATE")
        print("  Using inline fallback Modelfile.")
        MODELFILE_TEMPLATE = (
            "FROM {base_model}\n"
            "SYSTEM \"You are a senior cybersecurity analyst for the Donjon "
            "GRC Platform. Analyze vulnerabilities, prioritize findings, and "
            "provide actionable remediation guidance. Respond in structured JSON "
            "when requested.\"\n"
            "PARAMETER temperature 0.3\n"
            "PARAMETER num_predict 2048\n"
        )

    modelfile_content = MODELFILE_TEMPLATE.format(base_model=base_model)

    # POST to /api/create
    url = f"{OLLAMA_BASE_URL}/api/create"
    payload = json.dumps({
        'name': CUSTOM_MODEL_NAME,
        'modelfile': modelfile_content,
        'stream': True,
    }).encode('utf-8')

    req = urllib.request.Request(
        url,
        data=payload,
        headers={
            'Content-Type': 'application/json',
            'User-Agent': 'DonjonGRC/1.0',
        },
        method='POST',
    )

    try:
        with urllib.request.urlopen(req, timeout=300) as resp:
            for line in resp:
                if not line.strip():
                    continue
                try:
                    obj = json.loads(line.decode('utf-8'))
                    status = obj.get('status', '')
                    if status:
                        print(f"  {status}")
                    if obj.get('error'):
                        print(f"  ERROR: {obj['error']}")
                        return False
                except json.JSONDecodeError:
                    continue

        print(f"  Custom model '{CUSTOM_MODEL_NAME}' created successfully.")
        return True

    except Exception as exc:
        print(f"  FAILED to create custom model: {exc}")
        return False


def test_model(model_name: str) -> bool:
    """Send a quick test prompt to verify the model works."""
    print(f"  Testing model: {model_name} ...", end=' ', flush=True)

    url = f"{OLLAMA_BASE_URL}/api/generate"
    payload = json.dumps({
        'model': model_name,
        'prompt': 'Respond with exactly: MODEL_OK',
        'stream': False,
        'options': {
            'temperature': 0.0,
            'num_predict': 20,
        },
    }).encode('utf-8')

    req = urllib.request.Request(
        url,
        data=payload,
        headers={
            'Content-Type': 'application/json',
            'User-Agent': 'DonjonGRC/1.0',
        },
        method='POST',
    )

    try:
        with urllib.request.urlopen(req, timeout=120) as resp:
            result = json.loads(resp.read().decode('utf-8'))
            response_text = result.get('response', '').strip()
            if response_text:
                print(f"OK (response: {response_text[:60]})")
                return True
            else:
                print("EMPTY RESPONSE")
                return False
    except Exception as exc:
        print(f"FAILED ({exc})")
        return False


def save_config(models_available: list, custom_model: str, test_results: dict):
    """Write AI configuration to data/ai_config.json."""
    DATA_DIR.mkdir(parents=True, exist_ok=True)

    config = {
        'version': '1.0',
        'ollama_url': OLLAMA_BASE_URL,
        'models_available': models_available,
        'custom_model': custom_model,
        'preferred_model': custom_model if custom_model in models_available else (
            models_available[0] if models_available else None
        ),
        'test_results': test_results,
        'setup_timestamp': __import__('datetime').datetime.now(__import__('datetime').timezone.utc).isoformat(),
        'setup_host': os.environ.get('COMPUTERNAME', os.environ.get('HOSTNAME', 'unknown')),
    }

    with open(AI_CONFIG_PATH, 'w', encoding='utf-8') as f:
        json.dump(config, f, indent=2)

    print(f"\n  Configuration saved: {AI_CONFIG_PATH}")
    return config


def get_system_memory_gb() -> float:
    """Return total system RAM in GB (best-effort, cross-platform)."""
    try:
        import psutil
        return psutil.virtual_memory().total / (1024 ** 3)
    except ImportError:
        pass
    # Fallback: read /proc/meminfo on Linux
    try:
        with open('/proc/meminfo') as f:
            for line in f:
                if line.startswith('MemTotal'):
                    kb = int(line.split()[1])
                    return kb / (1024 ** 2)
    except (OSError, ValueError):
        pass
    # Fallback: Windows wmic
    try:
        import subprocess
        out = subprocess.check_output(
            ['wmic', 'ComputerSystem', 'get', 'TotalPhysicalMemory'],
            text=True, timeout=10,
        )
        for line in out.strip().split('\n'):
            line = line.strip()
            if line.isdigit():
                return int(line) / (1024 ** 3)
    except Exception:
        pass
    return 0.0


def download_step35_gguf() -> bool:
    """
    Download Step 3.5 Flash GGUF Q4_K_S from HuggingFace and register
    it with Ollama as 'step-3.5-flash'.

    The model is 111.5 GB and requires 120GB+ VRAM for inference.
    For most users, using the cloud API (STEPFUN_API_KEY) is recommended.
    """
    ram_gb = get_system_memory_gb()
    print(f"\n  System RAM detected: {ram_gb:.0f} GB")

    if ram_gb < 100:
        print(f"  WARNING: Step 3.5 Flash needs ~120 GB memory for inference.")
        print(f"  Your system has {ram_gb:.0f} GB. The download will proceed but")
        print(f"  inference may fail or be extremely slow with heavy swapping.")
        print(f"  Consider using the cloud API instead: export STEPFUN_API_KEY=...")
        print()

    models_dir = DATA_DIR / 'models'
    models_dir.mkdir(parents=True, exist_ok=True)
    gguf_path = models_dir / 'step3.5_flash_Q4_K_S.gguf'

    if gguf_path.exists():
        size_gb = gguf_path.stat().st_size / (1024 ** 3)
        if size_gb > 100:
            print(f"  Step 3.5 Flash GGUF already downloaded ({size_gb:.1f} GB)")
            print(f"  Path: {gguf_path}")
            return _register_step35_with_ollama(gguf_path)
        else:
            print(f"  Incomplete download found ({size_gb:.1f} GB), re-downloading...")

    # Download from HuggingFace using urllib (stdlib)
    hf_url = (
        "https://huggingface.co/stepfun-ai/Step-3.5-Flash-GGUF-Q4_K_S/"
        "resolve/main/step3.5_flash_Q4_K_S.gguf"
    )

    print(f"  Downloading Step 3.5 Flash GGUF Q4_K_S (~111.5 GB)...")
    print(f"  Source: {hf_url}")
    print(f"  Destination: {gguf_path}")
    print(f"  This will take a long time. Use Ctrl+C to cancel.")
    print()

    try:
        req = urllib.request.Request(hf_url, headers={'User-Agent': 'DonjonGRC/1.0'})
        with urllib.request.urlopen(req, timeout=30) as resp:
            total = int(resp.headers.get('Content-Length', 0))
            downloaded = 0
            chunk_size = 8 * 1024 * 1024  # 8MB chunks

            with open(gguf_path, 'wb') as f:
                while True:
                    chunk = resp.read(chunk_size)
                    if not chunk:
                        break
                    f.write(chunk)
                    downloaded += len(chunk)
                    if total > 0:
                        pct = downloaded * 100 // total
                        gb_done = downloaded / (1024 ** 3)
                        gb_total = total / (1024 ** 3)
                        print(
                            f"\r  Progress: {pct}% ({gb_done:.1f}/{gb_total:.1f} GB)",
                            end='', flush=True,
                        )
                    else:
                        gb_done = downloaded / (1024 ** 3)
                        print(f"\r  Downloaded: {gb_done:.1f} GB", end='', flush=True)

        print(f"\n  Download complete: {gguf_path}")
        return _register_step35_with_ollama(gguf_path)

    except KeyboardInterrupt:
        print(f"\n  Download cancelled by user.")
        if gguf_path.exists():
            print(f"  Partial download kept at: {gguf_path}")
            print(f"  Re-run with --step35 to resume (if server supports range requests).")
        return False
    except Exception as exc:
        print(f"\n  Download failed: {exc}")
        return False


def _register_step35_with_ollama(gguf_path) -> bool:
    """Create an Ollama model from the downloaded GGUF file."""
    print(f"\n  Registering Step 3.5 Flash with Ollama...")

    modelfile_content = (
        f"FROM {gguf_path}\n"
        "SYSTEM \"You are a senior cybersecurity analyst for the Donjon "
        "GRC Platform. Analyze vulnerabilities, prioritize findings by risk, "
        "and provide actionable remediation guidance. When asked for structured "
        "output, respond in JSON format.\"\n"
        "PARAMETER temperature 0.3\n"
        "PARAMETER num_predict 2048\n"
        "PARAMETER num_ctx 16384\n"
    )

    url = f"{OLLAMA_BASE_URL}/api/create"
    payload = json.dumps({
        'name': 'step-3.5-flash',
        'modelfile': modelfile_content,
        'stream': True,
    }).encode('utf-8')

    req = urllib.request.Request(
        url,
        data=payload,
        headers={
            'Content-Type': 'application/json',
            'User-Agent': 'DonjonGRC/1.0',
        },
        method='POST',
    )

    try:
        with urllib.request.urlopen(req, timeout=600) as resp:
            for line in resp:
                if not line.strip():
                    continue
                try:
                    obj = json.loads(line.decode('utf-8'))
                    status = obj.get('status', '')
                    if status:
                        print(f"  {status}")
                    if obj.get('error'):
                        print(f"  ERROR: {obj['error']}")
                        return False
                except json.JSONDecodeError:
                    continue

        print(f"  Model 'step-3.5-flash' registered with Ollama successfully.")
        return True

    except Exception as exc:
        print(f"  Failed to register with Ollama: {exc}")
        print(f"  You can still use the GGUF file directly with llama.cpp:")
        print(f"    ./llama-cli -m {gguf_path} -c 16384 -b 2048 -ub 2048 -fa on")
        return False


def print_install_instructions():
    """Print Ollama installation instructions for the user."""
    print("  Ollama is NOT running at http://localhost:11434")
    print()
    print("  To install Ollama:")
    print()
    print("    Windows:")
    print("      1. Download from https://ollama.com/download/windows")
    print("      2. Run the installer")
    print("      3. Ollama will start automatically as a background service")
    print()
    print("    Linux:")
    print("      curl -fsSL https://ollama.com/install.sh | sh")
    print()
    print("    macOS:")
    print("      1. Download from https://ollama.com/download/mac")
    print("      2. Move Ollama.app to Applications and launch it")
    print()
    print("  After installation, run this script again:")
    print("      python bin/setup-ai.py")
    print()
    print("  For air-gapped deployment:")
    print("      1. Run this script on a connected machine to download models")
    print("      2. Copy the Ollama model storage directory to the target:")
    print("         - Windows: %USERPROFILE%\\.ollama\\models")
    print("         - Linux/macOS: ~/.ollama/models")
    print("      3. Install Ollama on the target (installer does not need internet)")
    print()


# ---------------------------------------------------------------------------
# Interactive BYO / backend wizard
# ---------------------------------------------------------------------------

def _prompt(label, default=''):
    """Read one line from stdin with an optional default. Returns a string."""
    suffix = f" [{default}]" if default else ""
    try:
        val = input(f"  {label}{suffix}: ").strip()
    except EOFError:
        val = ''
    return val or default


def print_backend_menu(backends, offline: bool):
    """Print the numbered backend menu for the current mode."""
    print()
    print("  Select an AI backend (AI analysis is a Pro feature):")
    if offline:
        print("  *** DONJON_OFFLINE=1 -- air-gapped: local backends only, no cloud. ***")
    print()
    for i, b in enumerate(backends, start=1):
        tags = []
        tags.append("air-gap OK" if b.get('airgap') else "cloud / egress")
        if b.get('metered'):
            tags.append("metered (your key)")
        print(f"    {i:>2}. {b.get('label', b.get('id')):<34} ({', '.join(tags)})")
        req = b.get('requires')
        if req and req != 'Nothing':
            print(f"        needs: {req}")
    print()


def run_wizard(offline: bool):
    """Interactive backend chooser. Returns a config dict (or None to abort).

    Drives the menu entirely from AIEngine.get_supported_backends() so it never
    drifts from the engine, and refuses non-local endpoints when offline.
    """
    all_backends = load_supported_backends()
    backends = filter_backends_for_mode(all_backends, offline)
    print_backend_menu(backends, offline)

    raw = _prompt("Backend number", '1')
    try:
        idx = int(raw) - 1
        chosen = backends[idx]
    except (ValueError, IndexError):
        print("  Invalid selection.")
        return None
    backend_id = chosen['id']
    print(f"\n  -> {chosen.get('label', backend_id)}")

    model = api_url = api_key = req_tmpl = resp_path = ollama_url = None

    if backend_id == 'ollama':
        ollama_url = _prompt("Ollama base URL", OLLAMA_BASE_URL)
        model = _prompt("Model name (e.g. qwen2.5-coder:14b)", '') or None

    elif backend_id == 'custom':
        print("  Custom = ANY OpenAI-compatible /chat/completions endpoint")
        print("  (llama.cpp --api, vLLM, LM Studio, LiteLLM, OpenRouter, a gateway...).")
        api_url = _prompt("Base URL (full /v1/chat/completions URL)", '')
        api_key = _prompt("API key (blank if the endpoint needs none)", '') or None
        model = _prompt("Model name", 'default') or 'default'

    elif backend_id == 'generic':
        print("  Generic = universal adapter for ANY HTTP inference endpoint.")
        print("  request_template is JSON with {prompt} (and optional {system}).")
        print("  response_path extracts the text, e.g. choices[0].message.content")
        api_url = _prompt("Endpoint URL", '')
        req_tmpl = _prompt("Request template JSON", '{"prompt": "{prompt}"}') or '{"prompt": "{prompt}"}'
        resp_path = _prompt("Response path", 'response') or 'response'
        api_key = _prompt("API key / bearer token (optional)", '') or None
        model = _prompt("Model name (optional)", 'generic') or 'generic'

    elif backend_id == 'template':
        pass  # nothing to configure

    else:
        # Cloud / BYO-key backend.
        if backend_requires_key(backend_id):
            env_hint = chosen.get('requires', '')
            print(f"  Bring your own key. You can also set it via env: {env_hint}")
            api_key = _prompt("API key (blank to use the environment variable)", '') or None
        if backend_id in ('openai', 'azure'):
            api_url = _prompt("Custom base URL (blank for provider default)", '') or None
        model = _prompt("Model name (blank for provider default)", '') or None
        if chosen.get('metered'):
            print("  Note: this backend is billed per-token against YOUR account/key.")

    cfg = build_backend_config(
        backend_id, model=model, api_url=api_url, api_key=api_key,
        request_template=req_tmpl, response_path=resp_path, ollama_url=ollama_url,
    )
    ok, err = validate_backend_config(cfg, offline)
    if not ok:
        print(f"\n  Configuration rejected: {err}")
        return None
    return cfg


def save_backend_config(cfg) -> dict:
    """Persist a wizard/CLI backend config to data/ai_config.json.

    Merges over any existing file so unrelated keys (e.g. setup metadata) are
    preserved, then writes back the selected backend.
    """
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    existing = {}
    try:
        if AI_CONFIG_PATH.exists():
            with open(AI_CONFIG_PATH, 'r', encoding='utf-8') as f:
                existing = json.load(f)
    except Exception:
        existing = {}
    existing.update(cfg)
    existing.setdefault('version', '1.0')
    with open(AI_CONFIG_PATH, 'w', encoding='utf-8') as f:
        json.dump(existing, f, indent=2)
    return existing


def configure_backend_noninteractive(args, offline: bool) -> int:
    """Configure a backend straight from CLI flags (no prompts).

    Returns a process exit code. Used when --backend is supplied.
    """
    all_backends = load_supported_backends()
    if not backend_choice_is_allowed(args.backend, all_backends, offline):
        if offline:
            print(f"  DONJON_OFFLINE: backend '{args.backend}' performs public egress "
                  f"and is not allowed in air-gapped mode.")
        else:
            print(f"  Unknown backend '{args.backend}'. "
                  f"Valid: {', '.join(backend_ids(all_backends))}")
        return 1

    headers = None
    if getattr(args, 'header', None):
        headers = {}
        for h in args.header:
            if ':' in h:
                k, v = h.split(':', 1)
                headers[k.strip()] = v.strip()

    cfg = build_backend_config(
        args.backend,
        model=getattr(args, 'model_name', None),
        api_url=getattr(args, 'api_url', None),
        api_key=getattr(args, 'api_key', None),
        request_template=getattr(args, 'request_template', None),
        response_path=getattr(args, 'response_path', None),
        headers=headers,
    )
    ok, err = validate_backend_config(cfg, offline)
    if not ok:
        print(f"  Configuration rejected: {err}")
        return 1
    save_backend_config(cfg)
    masked = dict(cfg)
    if masked.get('api_key'):
        masked['api_key'] = '****'
    print(f"  Saved AI config: {AI_CONFIG_PATH}")
    print(f"  {json.dumps(masked)}")
    return 0


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description='Configure the Donjon AI engine (Ollama / BYO key / custom / '
                    'generic endpoint). AI analysis is a Pro feature.')

    parser.add_argument(
        '--models', nargs='+', metavar='MODEL',
        default=None,
        help='Specific Ollama model(s) to pull (e.g. mistral:7b codellama:7b). '
             'Defaults to: ' + ', '.join(DEFAULT_MODELS),
    )
    parser.add_argument(
        '--skip-pull', action='store_true',
        help='Skip model downloads, only configure and test existing models',
    )
    parser.add_argument(
        '--test-only', action='store_true',
        help='Only test model availability, do not download or configure',
    )
    parser.add_argument(
        '--step35', action='store_true',
        help='Download Step 3.5 Flash GGUF (111.5GB, needs 120GB+ VRAM). '
             'For most users, set STEPFUN_API_KEY instead.',
    )
    # -- Wizard / bring-your-own backend ------------------------------------
    parser.add_argument(
        '--wizard', action='store_true',
        help='Launch the interactive backend chooser (Ollama, custom, generic, '
             'BYO cloud key, template). Air-gap aware.',
    )
    parser.add_argument(
        '--list-backends', action='store_true',
        help='List the AI backends available in the current mode and exit.',
    )
    parser.add_argument(
        '--backend', metavar='ID', default=None,
        help='Configure this backend non-interactively (e.g. custom, generic, '
             'ollama, anthropic, openai, gemini, template).',
    )
    parser.add_argument(
        '--api-url', dest='api_url', metavar='URL', default=None,
        help='Endpoint / base URL for custom, generic, ollama, openai backends.',
    )
    parser.add_argument(
        '--api-key', dest='api_key', metavar='KEY', default=None,
        help='Bring-your-own API key for cloud / custom / generic backends.',
    )
    parser.add_argument(
        '--model-name', dest='model_name', metavar='NAME', default=None,
        help='Model name to record for the selected backend.',
    )
    parser.add_argument(
        '--request-template', dest='request_template', metavar='JSON', default=None,
        help="generic backend: JSON request body with {prompt}/{system} placeholders.",
    )
    parser.add_argument(
        '--response-path', dest='response_path', metavar='PATH', default=None,
        help="generic backend: dotted path to the answer (e.g. choices[0].message.content).",
    )
    parser.add_argument(
        '--header', action='append', metavar='K:V', default=None,
        help='generic backend: extra HTTP header (repeatable), e.g. "X-Api-Key: abc".',
    )

    args = parser.parse_args()
    models_to_pull = args.models or DEFAULT_MODELS
    offline = is_offline_mode()

    banner()

    # -------------------------------------------------------------------
    # Backend listing (no Ollama needed)
    # -------------------------------------------------------------------
    if args.list_backends:
        backends = filter_backends_for_mode(load_supported_backends(), offline)
        print_backend_menu(backends, offline)
        sys.exit(0)

    # -------------------------------------------------------------------
    # Non-interactive BYO / custom / generic configuration (no Ollama needed)
    # -------------------------------------------------------------------
    if args.backend:
        sys.exit(configure_backend_noninteractive(args, offline))

    # -------------------------------------------------------------------
    # Interactive wizard (no Ollama needed unless Ollama is chosen)
    # -------------------------------------------------------------------
    if args.wizard:
        cfg = run_wizard(offline)
        if not cfg:
            sys.exit(1)
        save_backend_config(cfg)
        masked = dict(cfg)
        if masked.get('api_key'):
            masked['api_key'] = '****'
        print(f"\n  Saved AI config: {AI_CONFIG_PATH}")
        print(f"  {json.dumps(masked, indent=2)}")
        print()
        print("  The Donjon AI engine will pick this up on next launch.")
        if cfg.get('backend') == 'ollama':
            print("  Tip: run 'python bin/setup-ai.py' to pull/test Ollama models.")
        sys.exit(0)

    # -------------------------------------------------------------------
    # Step 1: Check Ollama availability
    # -------------------------------------------------------------------
    print("[1/5] Checking Ollama availability...")
    if not ollama_is_running():
        print_install_instructions()
        print("  Don't want to run Ollama? Point Donjon at a model you already have:")
        print("    python bin/setup-ai.py --wizard          # guided BYO setup")
        print("    python bin/setup-ai.py --backend custom \\")
        print("        --api-url http://127.0.0.1:8000/v1/chat/completions")
        if not offline:
            print("    python bin/setup-ai.py --backend anthropic --api-key sk-...")
        print()
        sys.exit(1)

    installed = get_installed_models()
    print(f"  Ollama is running. {len(installed)} model(s) installed:")
    for m in installed:
        print(f"    - {m}")

    # -------------------------------------------------------------------
    # Test-only mode
    # -------------------------------------------------------------------
    if args.test_only:
        print()
        print("[TEST] Testing all installed models...")
        test_results = {}
        for m in installed:
            ok = test_model(m)
            test_results[m] = 'pass' if ok else 'fail'
        passed = sum(1 for v in test_results.values() if v == 'pass')
        print(f"\n  Results: {passed}/{len(test_results)} models responding")
        save_config(installed, CUSTOM_MODEL_NAME, test_results)
        sys.exit(0 if passed > 0 else 1)

    # -------------------------------------------------------------------
    # Step 1b: Download Step 3.5 Flash GGUF (optional)
    # -------------------------------------------------------------------
    if args.step35:
        print(f"\n[1b] Downloading Step 3.5 Flash (111.5 GB GGUF)...")
        step35_ok = download_step35_gguf()
        if step35_ok:
            installed = get_installed_models()
            print(f"  Updated model list: {len(installed)} model(s)")
        else:
            print("  Step 3.5 Flash download/registration failed.")
            print("  Continuing with standard models...")

    # -------------------------------------------------------------------
    # Step 2: Pull models
    # -------------------------------------------------------------------
    if not args.skip_pull:
        print(f"\n[2/5] Pulling {len(models_to_pull)} model(s)...")
        pull_results = {}
        for model in models_to_pull:
            # Check if already installed
            if model in installed:
                print(f"\n  Model {model} is already installed, skipping pull.")
                pull_results[model] = 'already_installed'
                continue
            ok = pull_model(model)
            pull_results[model] = 'success' if ok else 'failed'

        # Refresh installed list
        installed = get_installed_models()
        success_count = sum(1 for v in pull_results.values() if v != 'failed')
        print(f"\n  Pull results: {success_count}/{len(models_to_pull)} successful")
    else:
        print("\n[2/5] Skipping model downloads (--skip-pull)")

    # -------------------------------------------------------------------
    # Step 3: Create custom Modelfile
    # -------------------------------------------------------------------
    print(f"\n[3/5] Creating custom security-tuned model...")

    # Pick the best base model for the custom Modelfile
    base_model = None
    preferred_bases = ['mistral:7b', 'llama3.2:3b', 'llama3:latest']
    for pb in preferred_bases:
        if pb in installed:
            base_model = pb
            break
    # Fall back to any model tagged '7b' or the first available
    if not base_model:
        for m in installed:
            if '7b' in m:
                base_model = m
                break
    if not base_model and installed:
        base_model = installed[0]

    if base_model:
        custom_ok = create_custom_modelfile(base_model)
        if custom_ok:
            installed = get_installed_models()
    else:
        print("  No base model available for custom Modelfile creation.")
        custom_ok = False

    # -------------------------------------------------------------------
    # Step 4: Test models
    # -------------------------------------------------------------------
    print(f"\n[4/5] Testing model availability...")
    test_results = {}
    for m in installed:
        ok = test_model(m)
        test_results[m] = 'pass' if ok else 'fail'

    passed = sum(1 for v in test_results.values() if v == 'pass')
    print(f"\n  Test results: {passed}/{len(test_results)} models responding")

    # -------------------------------------------------------------------
    # Step 5: Save configuration
    # -------------------------------------------------------------------
    print(f"\n[5/5] Saving configuration...")
    config = save_config(installed, CUSTOM_MODEL_NAME, test_results)

    # -------------------------------------------------------------------
    # Summary
    # -------------------------------------------------------------------
    print()
    print("  " + "=" * 58)
    print("  Setup Complete")
    print("  " + "=" * 58)
    print()
    print(f"  Backend:          Ollama ({OLLAMA_BASE_URL})")
    print(f"  Models installed: {len(installed)}")
    print(f"  Models passing:   {passed}")
    print(f"  Preferred model:  {config.get('preferred_model', 'none')}")
    print(f"  Config file:      {AI_CONFIG_PATH}")
    print()

    if passed == 0:
        print("  WARNING: No models passed testing. AI features will use template fallback.")
        print()
        sys.exit(1)

    print("  The AI engine will auto-detect these models at runtime.")
    print("  For air-gapped transfer, copy the Ollama models directory to the target.")
    print()
    print("  Other ways to point Donjon at a model:")
    print("    python bin/setup-ai.py --wizard            # guided backend chooser")
    print("    python bin/setup-ai.py --backend custom \\  # any OpenAI-compatible URL")
    print("        --api-url http://127.0.0.1:8000/v1/chat/completions")
    print("    python bin/setup-ai.py --backend generic \\ # any HTTP endpoint")
    print("        --api-url http://127.0.0.1:9000/infer \\")
    print("        --request-template '{\"prompt\":\"{prompt}\"}' --response-path response")
    print()

    if offline:
        print("  DONJON_OFFLINE=1 is set: cloud backends are disabled. AI stays fully")
        print("  local (Ollama / custom / generic against a loopback or private endpoint).")
        print()
    else:
        print("  Bring-your-own cloud key (AI is a Pro feature; billed to your account):")
        print("    export STEPFUN_API_KEY=...    # Step 3.5 Flash cloud API")
        print("    export ANTHROPIC_API_KEY=...  # Anthropic Claude")
        print("    export GEMINI_API_KEY=...     # Google Gemini 2.5 Flash (fast, free tier)")
        print("    export OPENAI_API_KEY=...     # OpenAI GPT-4")
        print()
        print("  To download Step 3.5 Flash for offline use (111.5 GB, needs 120GB+ VRAM):")
        print("    python bin/setup-ai.py --step35")
        print()


if __name__ == '__main__':
    main()
