#!/usr/bin/env python3
"""
Cambium Fiber OLT Zabbix Template Installer.

- Ensures a Zabbix-writable cache directory for cambium_olt_ssh_json.py
- Installs cambium_olt_ssh_json.py into Zabbix ExternalScripts directory
- Downloads template YAML
- Ensures template group exists
- Imports template via Zabbix API

Supports running a single step with --step <name>, or all steps if omitted.

Debian/Ubuntu only for auto-install of OS packages.
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
import textwrap
import urllib.request
import urllib.error
from dataclasses import dataclass
from typing import Callable, Dict, List, Optional, Set


BASE_URL = "https://joshaven.com/resources/tools/cambium-fiber-olt-zabbix-template"
SCRIPT_NAME = "cambium_olt_ssh_json.py"
TEMPLATE_NAME = "zbx_cambium_fiber_olt_1.3.0.yaml"
DEFAULT_EXT_DIR = "/usr/lib/zabbix/externalscripts"
DEFAULT_ZBX_URL = "http://localhost/"
DEFAULT_CACHE_DIR = "/var/cache/cambium-olt"


# ---------- pretty printing ----------
def green(s: str) -> str: return f"\033[0;32m{s}\033[0m"
def yellow(s: str) -> str: return f"\033[0;33m{s}\033[0m"
def red(s: str) -> str: return f"\033[0;31m{s}\033[0m"

def ok(msg: str) -> None: print(green(f"✓ {msg}"))
def step(msg: str) -> None: print(yellow(f"… {msg}"))
def fail(msg: str, code: int = 1) -> None:
    print(red(f"✗ {msg}"), file=sys.stderr)
    sys.exit(code)


# ---------- utilities ----------
def need_cmd(cmd: str) -> None:
    if shutil.which(cmd) is None:
        fail(f"Missing dependency: {cmd}")

def has_cmd(cmd: str) -> bool:
    return shutil.which(cmd) is not None

def is_root() -> bool:
    return os.geteuid() == 0

def apt_install(packages: List[str]) -> None:
    if not has_cmd("apt-get"):
        fail(
            "apt-get not found. This installer currently supports Debian/Ubuntu only.\n"
            f"Please install manually: {' '.join(packages)}"
        )
    if not is_root():
        fail("Auto-install requires root. Re-run with sudo.")

    step(f"Installing missing packages via apt-get: {' '.join(packages)}")
    subprocess.run(["apt-get", "update", "-y"], check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    res = subprocess.run(["apt-get", "install", "-y"] + packages)
    if res.returncode != 0:
        fail(f"apt-get failed installing: {' '.join(packages)}")
    ok(f"Installed: {' '.join(packages)}")

def read_secret(prompt: str) -> str:
    import getpass
    return getpass.getpass(prompt)

def detect_externalscripts_dir(explicit: Optional[str] = None) -> str:
    if explicit and os.path.isdir(explicit):
        return explicit

    env_dir = os.environ.get("EXTERNALSCRIPTS")
    if env_dir and os.path.isdir(env_dir):
        return env_dir

    for conf in ("/etc/zabbix/zabbix_server.conf", "/etc/zabbix/zabbix_proxy.conf"):
        if os.path.isfile(conf):
            try:
                with open(conf, "r", encoding="utf-8", errors="ignore") as f:
                    for line in f:
                        if line.startswith("ExternalScripts="):
                            d = line.split("=", 1)[1].strip()
                            if d:
                                return d
            except OSError:
                pass

    return DEFAULT_EXT_DIR

def download_file(url: str, dest: str) -> None:
    try:
        with urllib.request.urlopen(url) as resp:
            data = resp.read()
    except urllib.error.URLError as e:
        fail(f"Failed to download {url}: {e}")
    if not data:
        fail(f"Downloaded file is empty: {dest}")
    os.makedirs(os.path.dirname(dest), exist_ok=True)
    with open(dest, "wb") as f:
        f.write(data)

def ensure_api_endpoint(base: str) -> str:
    base = (base or "").rstrip("/")
    if not base:
        base = "http://localhost"

    candidates = []

    if base.endswith("/api_jsonrpc.php"):
        candidates.append(base)
        base_root = base[:-len("/api_jsonrpc.php")]
        candidates.append(base_root + "/api_jsonrpc.php")
        candidates.append(base_root + "/zabbix/api_jsonrpc.php")
    else:
        candidates.append(base + "/api_jsonrpc.php")
        candidates.append(base + "/zabbix/api_jsonrpc.php")

    probe_payload = json.dumps({
        "jsonrpc": "2.0",
        "method": "apiinfo.version",
        "params": {},
        "id": 1
    }).encode("utf-8")

    for url in candidates:
        try:
            req = urllib.request.Request(
                url,
                data=probe_payload,
                headers={"Content-Type": "application/json-rpc"},
                method="POST",
            )
            with urllib.request.urlopen(req, timeout=5) as resp:
                raw = resp.read().decode("utf-8", errors="replace")
            out = json.loads(raw)
            if "result" in out and out["result"]:
                return url
        except Exception:
            continue

    fail(
        "Could not find a working Zabbix API endpoint.\n"
        "Tried:\n  - " + "\n  - ".join(candidates) + "\n"
        "If your Zabbix UI is on a custom path, re-run with --zbx-url "
        "set to your normal browser URL."
    )


# ---------- zabbix api ----------
@dataclass
class ZabbixClient:
    api_url: str
    token: str
    timeout: int = 15

    def call(self, method: str, params: dict) -> dict:
        payload = {
            "jsonrpc": "2.0",
            "method": method,
            "params": params,
            "auth": self.token,
            "id": 1,
        }
        data = json.dumps(payload).encode("utf-8")
        req = urllib.request.Request(
            self.api_url,
            data=data,
            headers={"Content-Type": "application/json-rpc"},
            method="POST",
        )
        try:
            with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                raw = resp.read().decode("utf-8", errors="replace")
        except urllib.error.HTTPError as e:
            body = e.read().decode("utf-8", errors="replace")
            fail(f"HTTP error calling Zabbix API ({e.code}): {body[:400]}")
        except urllib.error.URLError as e:
            fail(f"Could not reach Zabbix API at {self.api_url}: {e}")

        try:
            out = json.loads(raw)
        except json.JSONDecodeError:
            snippet = raw.strip().splitlines()[0][:200]
            fail(
                "Non-JSON response from API endpoint. "
                "Are you pointing at api_jsonrpc.php?\n"
                f"Got: {snippet}"
            )

        if "error" in out:
            err = out["error"]
            code = err.get("code")
            data_msg = err.get("data") or err.get("message") or str(err)
            fail(f"Zabbix API error {code}: {data_msg}")

        return out.get("result", {})

    def version(self) -> str:
        payload = {"jsonrpc": "2.0", "method": "apiinfo.version", "params": {}, "id": 1}
        data = json.dumps(payload).encode("utf-8")
        req = urllib.request.Request(
            self.api_url,
            data=data,
            headers={"Content-Type": "application/json-rpc"},
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=self.timeout) as resp:
            raw = resp.read().decode("utf-8", errors="replace")
        out = json.loads(raw)
        return out.get("result", "")


# ---------- steps ----------
class Context:
    def __init__(self, args: argparse.Namespace):
        self.args = args
        self.ext_dir: Optional[str] = None
        self.script_path: Optional[str] = None
        self.tmp_template: Optional[str] = None
        self.zbx_url: Optional[str] = None
        self.api_url: Optional[str] = None
        self.token: Optional[str] = None
        self.client: Optional[ZabbixClient] = None
        self.template_id: Optional[str] = None  # Store template ID after import

# ---------- helper functions ----------
def get_olt_name(script_path: str, host: str, password: str) -> Optional[str]:
    """Get OLT system name via SSH."""
    try:
        result = subprocess.run(
            [script_path, host, password, "System.Name"],
            capture_output=True,
            text=True,
            timeout=30  # Increased timeout to handle cold cache
        )
        if result.returncode == 0:
            name = result.stdout.strip()
            # Don't accept "0" (default for missing data) or empty as valid names
            if name and not name.startswith("error:") and name != "0":
                return name
    except Exception:
        pass
    return None

def modify_template_password(template_path: str, password: str) -> None:
    """Modify template YAML to set {$OLT.PASS} macro value."""
    with open(template_path, "r", encoding="utf-8") as f:
        content = f.read()

    # Replace the macro value
    import re
    pattern = r"(- macro: '\{\$OLT\.PASS\}'\s+value:) .+"
    replacement = f"\\1 {password}"
    content = re.sub(pattern, replacement, content, flags=re.MULTILINE)

    with open(template_path, "w", encoding="utf-8") as f:
        f.write(content)

# ---------- steps ----------
def step_deps_api(ctx: Context) -> None:
    step("Checking API prerequisites")
    ok("API prerequisites OK")

def step_deps_ssh(ctx: Context) -> None:
    step("Checking SSH prerequisites")

    if not has_cmd("ssh"):
        apt_install(["openssh-client"])
    else:
        ok("ssh present")

    if not has_cmd("sshpass"):
        apt_install(["sshpass"])
    else:
        ok("sshpass present")

def step_ensure_cache_dir(ctx: Context) -> None:
    step("Ensuring cache directory exists")
    cache_dir = ctx.args.cache_dir or os.environ.get("OLT_CACHE_DIR") or DEFAULT_CACHE_DIR

    if not is_root():
        fail("Creating cache directory requires root. Re-run with sudo.")

    os.makedirs(cache_dir, exist_ok=True)

    try:
        import pwd, grp
        u = pwd.getpwnam("zabbix").pw_uid
        g = grp.getgrnam("zabbix").gr_gid
        os.chown(cache_dir, u, g)
        os.chmod(cache_dir, 0o775)
    except Exception as e:
        fail(f"Failed to set cache directory ownership: {e}")

    # Verify zabbix user can write by testing as that user
    test_file = os.path.join(cache_dir, ".write_test")
    try:
        import pwd
        zabbix_uid = pwd.getpwnam("zabbix").pw_uid
        # Create test file as zabbix user
        subprocess.run(
            ["sudo", "-u", "zabbix", "touch", test_file],
            check=True,
            capture_output=True
        )
        os.remove(test_file)
    except (subprocess.CalledProcessError, Exception) as e:
        fail(f"Cache directory not writable by zabbix user: {cache_dir}")

    ok(f"Cache directory ready: {cache_dir}")

def step_detect_externalscripts(ctx: Context) -> None:
    step("Detecting ExternalScripts directory")
    ctx.ext_dir = detect_externalscripts_dir(ctx.args.ext_dir)
    os.makedirs(ctx.ext_dir, exist_ok=True)
    if not os.access(ctx.ext_dir, os.W_OK):
        fail(f"ExternalScripts directory not writable: {ctx.ext_dir}")
    ok(f"ExternalScripts directory: {ctx.ext_dir}")

def step_install_getter(ctx: Context) -> None:
    if not ctx.ext_dir:
        step_detect_externalscripts(ctx)
    step("Downloading backend script")
    ctx.script_path = os.path.join(ctx.ext_dir, SCRIPT_NAME)
    download_file(f"{BASE_URL}/{SCRIPT_NAME}", ctx.script_path)
    os.chmod(ctx.script_path, 0o755)

    try:
        import pwd, grp
        u = pwd.getpwnam("zabbix").pw_uid
        g = grp.getgrnam("zabbix").gr_gid
        os.chown(ctx.script_path, u, g)
    except Exception:
        pass

    ok(f"Installed backend script: {ctx.script_path}")

def step_download_template(ctx: Context) -> None:
    step("Downloading template")
    ctx.tmp_template = os.path.join("/tmp", TEMPLATE_NAME)
    download_file(f"{BASE_URL}/{TEMPLATE_NAME}", ctx.tmp_template)
    ok(f"Downloaded template: {ctx.tmp_template}")

def _load_api_creds(ctx: Context) -> None:
    if ctx.client:
        return
    ctx.zbx_url = ctx.args.zbx_url or os.environ.get("ZBX_URL") or DEFAULT_ZBX_URL
    ctx.token = ctx.args.api_token or os.environ.get("ZBX_TOKEN")
    if not ctx.token:
        ctx.token = read_secret("Zabbix API token: ")
    ctx.api_url = ensure_api_endpoint(ctx.zbx_url)
    ctx.client = ZabbixClient(ctx.api_url, ctx.token)

def step_auth_check(ctx: Context) -> None:
    step("Validating Zabbix API access")
    _load_api_creds(ctx)
    ver = ctx.client.version()
    if not ver:
        fail("Could not read Zabbix API version. Check URL/token.")
    ok(f"Connected to Zabbix API (version {ver})")

def step_ensure_group(ctx: Context) -> None:
    _load_api_creds(ctx)

    step("Ensuring template group exists")
    group_name = ctx.args.group_name

    res = ctx.client.call("templategroup.get", {
        "filter": {"name": [group_name]},
        "output": ["groupid", "name", "uuid"]
    })

    if isinstance(res, list) and res:
        ok(f"Template group exists: {group_name} (id {res[0]['groupid']})")
        return

    ctx.client.call("templategroup.create", {"name": group_name})
    ok(f"Created template group: {group_name}")

def step_import_template(ctx: Context) -> None:
    step_ensure_group(ctx)
    if not ctx.tmp_template:
        step_download_template(ctx)

    # Modify template password if specified
    if ctx.args.olt_password:
        step("Setting template macro {$OLT.PASS}")
        modify_template_password(ctx.tmp_template, ctx.args.olt_password)
        ok(f"Template macro updated")

    step("Importing template via Zabbix API")
    with open(ctx.tmp_template, "r", encoding="utf-8") as f:
        source = f.read()

    params = {
        "format": "yaml",
        "source": source,
        "rules": {
            "templates": {"createMissing": True, "updateExisting": True},
            "items": {"createMissing": True, "updateExisting": True},
            "discoveryRules": {"createMissing": True, "updateExisting": True},
            "triggers": {"createMissing": True, "updateExisting": True},
            "graphs": {"createMissing": True, "updateExisting": True},
            "valueMaps": {"createMissing": True, "updateExisting": True},
        },
    }

    ctx.client.call("configuration.import", params)
    ok("Template imported successfully")

    # Store template ID for later use
    res = ctx.client.call("template.get", {
        "filter": {"host": ["Cambium Fiber OLT by SSH v1.3.0"]},
        "output": ["templateid"]
    })
    if isinstance(res, list) and res:
        ctx.template_id = res[0]["templateid"]

def step_flush_template(ctx: Context) -> None:
    _load_api_creds(ctx)

    step("Finding template")
    res = ctx.client.call("template.get", {
        "filter": {"host": ["Cambium Fiber OLT by SSH v1.3.0"]},
        "output": ["templateid", "name"]
    })

    if not res or not isinstance(res, list):
        ok("Template not found (nothing to flush)")
        return

    template_id = res[0]["templateid"]

    # Find all hosts using this template
    step("Finding hosts using template")
    hosts = ctx.client.call("host.get", {
        "templateids": [template_id],
        "output": ["hostid", "host"]
    })

    if hosts and isinstance(hosts, list):
        step(f"Removing {len(hosts)} host(s) using template")
        for host in hosts:
            ctx.client.call("host.delete", [host["hostid"]])
            ok(f"Removed host: {host['host']}")

    step("Removing template")
    ctx.client.call("template.delete", [template_id])
    ok("Template removed")

def step_add_hosts(ctx: Context) -> None:
    if not ctx.args.add_hosts:
        return

    _load_api_creds(ctx)

    if not ctx.script_path:
        step_detect_externalscripts(ctx)
        step_install_getter(ctx)

    if not ctx.template_id:
        # Get template ID
        res = ctx.client.call("template.get", {
            "filter": {"host": ["Cambium Fiber OLT by SSH v1.3.0"]},
            "output": ["templateid"]
        })
        if not res or not isinstance(res, list):
            fail("Template not found. Import template first.")
        ctx.template_id = res[0]["templateid"]

    # Get or create host group
    group_res = ctx.client.call("hostgroup.get", {
        "filter": {"name": ["Cambium Fiber OLT"]},
        "output": ["groupid"]
    })
    if group_res and isinstance(group_res, list):
        group_id = group_res[0]["groupid"]
    else:
        group_res = ctx.client.call("hostgroup.create", {"name": "Cambium Fiber OLT"})
        group_id = group_res["groupids"][0]

    # Parse host IPs
    host_ips = [ip.strip() for ip in ctx.args.add_hosts.split(",")]
    password = ctx.args.olt_password or "FiberDemo!"

    for ip in host_ips:
        if not ip:
            continue

        step(f"Processing host {ip}")

        # Get OLT name
        olt_name = get_olt_name(ctx.script_path, ip, password)
        if not olt_name:
            print(yellow(f"  ⚠ Could not get name from {ip}, using IP as name"))
            olt_name = f"OLT-{ip.replace('.', '-')}"

        # Check if host already exists
        existing = ctx.client.call("host.get", {
            "filter": {"host": [olt_name]},
            "output": ["hostid", "host"]
        })

        if existing and isinstance(existing, list):
            print(yellow(f"  ⚠ Host '{olt_name}' already exists (skipping)"))
            continue

        # Create host
        try:
            ctx.client.call("host.create", {
                "host": olt_name,
                "interfaces": [{
                    "type": 1,  # Agent interface
                    "main": 1,
                    "useip": 1,
                    "ip": ip,
                    "dns": "",
                    "port": "10050"
                }],
                "groups": [{"groupid": group_id}],
                "templates": [{"templateid": ctx.template_id}],
                "macros": [
                    {
                        "macro": "{$OLT.PASS}",
                        "value": password,
                        "type": 1  # Secret text
                    }
                ]
            })
            ok(f"Created host: {olt_name} ({ip})")
        except Exception as e:
            print(red(f"  ✗ Failed to create host {olt_name}: {e})"))

def step_postcheck(ctx: Context) -> None:
    step("Post-check")
    ok("Install complete")

    if ctx.args.add_hosts:
        print("\nHosts have been created and will start collecting data within 60 seconds.")
    else:
        print("Next steps:")
        print("  1. Add your OLT host in Zabbix.")
        print("  2. Link template: Cambium Fiber OLT by SSH v1.3.0.")
        print("  3. Set macros: {$OLT.PASS} (and optional {$OLT_CACHE_TTL}).")


STEPS: Dict[str, Callable[[Context], None]] = {
    "deps_api": step_deps_api,
    "deps_ssh": step_deps_ssh,
    "ensure_cache_dir": step_ensure_cache_dir,
    "detect_externalscripts": step_detect_externalscripts,
    "install_getter": step_install_getter,
    "download_template": step_download_template,
    "auth_check": step_auth_check,
    "ensure_group": step_ensure_group,
    "import_template": step_import_template,
    "flush_template": step_flush_template,
    "add_hosts": step_add_hosts,
    "postcheck": step_postcheck,
}

DEPS: Dict[str, List[str]] = {
    "deps_api": [],
    "deps_ssh": [],
    "ensure_cache_dir": ["deps_ssh"],
    "detect_externalscripts": ["deps_ssh"],
    "install_getter": ["detect_externalscripts"],
    "download_template": ["deps_api"],
    "auth_check": ["deps_api"],
    "ensure_group": ["auth_check"],
    "import_template": ["ensure_group", "download_template"],
    "flush_template": ["auth_check"],
    "add_hosts": ["import_template", "install_getter"],
    "postcheck": ["deps_ssh"],
}

def run_step_with_deps(step_name: str, ctx: Context, ran: Set[str]) -> None:
    for prereq in DEPS.get(step_name, []):
        if prereq not in ran:
            run_step_with_deps(prereq, ctx, ran)
    if step_name not in ran:
        STEPS[step_name](ctx)
        ran.add(step_name)

def parse_args(argv: List[str]) -> argparse.Namespace:
    p = argparse.ArgumentParser(
        prog="install.py",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description="Cambium Fiber OLT Zabbix Template Installer",
        epilog=textwrap.dedent("""\
        Examples:
          Full install:
            sudo ./install.py --api-token <TOKEN>

          Install with OLT password and auto-create hosts:
            sudo ./install.py --api-token <TOKEN> --olt-password "FiberDemo!" --add-hosts "192.168.50.10, 192.168.50.11"

          Flush everything and reinstall:
            sudo ./install.py --api-token <TOKEN> --flush --olt-password "FiberDemo!" --add-hosts "192.168.50.10"

          Only install SSH prerequisites:
            sudo ./install.py --step deps_ssh

          Only check API auth:
            sudo ./install.py --api-token <TOKEN> --step auth_check
        """)
    )
    p.add_argument("--api-token", dest="api_token",
                   help="Zabbix API token. Can also be set via ZBX_TOKEN env var.")
    p.add_argument("--zbx-url", dest="zbx_url",
                   help="Zabbix frontend URL (defaults to http://localhost/).")
    p.add_argument("--olt-password", dest="olt_password",
                   help="OLT admin password. Sets {$OLT.PASS} macro in template.")
    p.add_argument("--add-hosts", dest="add_hosts",
                   help="Comma-separated list of OLT IP addresses to add as hosts.")
    p.add_argument("--flush", action="store_true",
                   help="Remove all hosts using template and the template itself before install.")
    p.add_argument("--ext-dir", dest="ext_dir",
                   help="Override ExternalScripts directory (optional).")
    p.add_argument("--cache-dir", dest="cache_dir",
                   help=f"Cache directory for OLT JSON snapshots. Default: {DEFAULT_CACHE_DIR}")
    p.add_argument("--group-name", dest="group_name",
                   default="Templates/Network devices/Cambium Fiber",
                   help="Template group name to ensure/create. Default: Templates/Network devices/Cambium Fiber.")
    p.add_argument("--step", choices=sorted(STEPS.keys()),
                   help="Run a single step (with prerequisites). If omitted, runs full install.")
    return p.parse_args(argv)

def main(argv: List[str]) -> None:
    args = parse_args(argv)
    ctx = Context(args)

    print("Cambium Fiber OLT Zabbix Template Installer (Python)")

    ran: Set[str] = set()
    if args.step:
        run_step_with_deps(args.step, ctx, ran)
        return

    # Build pipeline based on arguments
    pipeline = []

    if args.flush:
        pipeline.extend([
            "auth_check",
            "flush_template",
        ])

    if not args.flush or args.add_hosts:
        # Normal install or install after flush
        pipeline.extend([
            "deps_ssh",
            "ensure_cache_dir",
            "detect_externalscripts",
            "install_getter",
            "download_template",
            "ensure_group",
            "import_template",
            "postcheck",
        ])

    if args.add_hosts:
        pipeline.append("add_hosts")

    for s in pipeline:
        run_step_with_deps(s, ctx, ran)

if __name__ == "__main__":
    main(sys.argv[1:])