"""Talks to the agents (SSH), GeoIP and Loki.

Errors are returned as {"ok": False, "error_key": ..., "error_args": {...}} so the
web layer can translate them into the language of each request. Errors coming from
an agent arrive as plain text in "error" and are passed through unchanged.
"""
import base64
import gzip
import json
import logging
import os
import re
import shutil
import subprocess
import tempfile
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from datetime import date, datetime, timedelta

import maxminddb

log = logging.getLogger("webui")


def fail(key, **args):
    return {"ok": False, "error_key": key, "error_args": args}


# ---------------------------------------------------------------------
# GeoIP (optional)
# ---------------------------------------------------------------------
class Geo:
    def __init__(self, settings):
        self.settings = settings
        self._reader = None
        self._lock = threading.Lock()

    @property
    def available(self):
        return os.path.exists(self.settings.geoip_db)

    def reload(self):
        with self._lock:
            self._reader = None

    def lookup(self, ip):
        with self._lock:
            if self._reader is None:
                try:
                    self._reader = maxminddb.open_database(self.settings.geoip_db)
                except Exception:
                    self._reader = False
            reader = self._reader
        if not reader:
            return {}
        try:
            rec = reader.get(ip) or {}
        except ValueError:
            return {}
        country = rec.get("country", {})
        names = country.get("names", {})
        city = rec.get("city", {}).get("names", {})
        location = rec.get("location", {})
        # The browser picks the name matching its UI language, English otherwise
        return {
            "cc": country.get("iso_code", ""),
            "country": names.get("en") or "",
            "country_names": names,
            "city": city.get("en") or "",
            "city_names": city,
            "lat": location.get("latitude"),
            "lon": location.get("longitude"),
        }

    def download_loop(self):
        """Keeps the free DB-IP Lite database current (monthly release)."""
        while True:
            try:
                self._download_if_old()
            except Exception as exc:
                log.warning("GeoIP download failed: %s", exc)
            time.sleep(86400)

    def _download_if_old(self):
        path = self.settings.geoip_db
        if os.path.exists(path) and time.time() - os.path.getmtime(path) < 32 * 86400:
            return
        os.makedirs(os.path.dirname(path), exist_ok=True)
        first = date.today().replace(day=1)
        for month in (first, (first - timedelta(days=1)).replace(day=1)):
            url = f"https://download.db-ip.com/free/dbip-city-lite-{month:%Y-%m}.mmdb.gz"
            try:
                with urllib.request.urlopen(url, timeout=300) as resp, \
                        tempfile.NamedTemporaryFile(dir=os.path.dirname(path), delete=False) as tmp:
                    with gzip.GzipFile(fileobj=resp) as gz:
                        shutil.copyfileobj(gz, tmp)
            except urllib.error.HTTPError as exc:
                if exc.code == 404:
                    continue
                raise
            os.replace(tmp.name, path)
            log.info("GeoIP database updated from %s", url)
            self.reload()
            return


# ---------------------------------------------------------------------
# Agents via SSH
# ---------------------------------------------------------------------
class Agents:
    def __init__(self, settings, geo):
        self.settings = settings
        self.geo = geo
        self._cache = {"ts": 0.0, "data": None}
        self._lock = threading.Lock()

    @property
    def hosts(self):
        return self.settings.hosts

    def ssh(self, name, *args, timeout=None):
        s = self.settings
        h = self.hosts[name]
        # -F none: ignore /etc/ssh/ssh_config (Debian enables GSSAPI there, which is never needed here)
        cmd = [
            "ssh", "-F", "none", "-i", s.ssh_key, "-p", str(h.port),
            "-o", "BatchMode=yes", "-o", "IdentitiesOnly=yes", "-o", "GSSAPIAuthentication=no",
            "-o", "StrictHostKeyChecking=yes", "-o", f"UserKnownHostsFile={s.known_hosts}",
            "-o", "ConnectTimeout=8", "-o", "ServerAliveInterval=5",
            "-o", "LogLevel=ERROR",
            "-l", h.user, "--", h.address, " ".join(args),
        ]
        try:
            r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout or s.ssh_timeout)
        except subprocess.TimeoutExpired:
            return fail("err_timeout", host=name)
        try:
            return json.loads(r.stdout.strip().splitlines()[-1])
        except (ValueError, IndexError):
            detail = (r.stderr or r.stdout).strip()[:300]
            if detail:
                return fail("err_ssh", host=name, detail=detail)
            return fail("err_exit", host=name, code=r.returncode)

    def run_many(self, names, *args):
        if not names:
            return {}
        with ThreadPoolExecutor(max_workers=min(len(names), 16)) as pool:
            return dict(zip(names, pool.map(lambda n: self.ssh(n, *args), names)))

    def bans(self, force=False):
        with self._lock:
            if not force and self._cache["data"] and time.time() - self._cache["ts"] < self.settings.cache_seconds:
                return self._cache["data"]
        results = self.run_many(list(self.hosts), "list")

        rows, hosts = [], {}
        for name, res in results.items():
            if not res.get("ok"):
                hosts[name] = {**{k: v for k, v in res.items() if k != "jails"}, "ok": False, "jails": {}}
                continue
            hosts[name] = {"ok": True, "jails": {}}
            for jail, info in res.get("jails", {}).items():
                hosts[name]["jails"][jail] = {"current": len(info["bans"]), "total": info.get("total")}
                for b in info["bans"]:
                    rows.append({"host": name, "jail": jail, **b, **self.geo.lookup(b["ip"])})
        rows.sort(key=lambda r: r["start"], reverse=True)
        data = {"generated": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "hosts": hosts, "bans": rows}
        with self._lock:
            self._cache.update(ts=time.time(), data=data)
        return data


# ---------------------------------------------------------------------
# Loki (optional)
# ---------------------------------------------------------------------
def loki_history(settings, ip):
    s = settings
    if not s.loki_url:
        return fail("err_loki_disabled"), 404
    esc = re.escape(ip).replace("\\", "\\\\")
    # |= is a fast pre-filter, the regex then only checks the hits
    query = (f'{s.loki_selector} |= "{ip}" '
             f'|~ "(^|[^0-9a-fA-F.:]){esc}([^0-9a-fA-F.:]|$)"')
    now_ns = int(time.time() * 1e9)
    params = urllib.parse.urlencode({
        "query": query, "limit": s.loki_limit, "direction": "backward",
        "start": now_ns - s.loki_days * 86400 * 10**9, "end": now_ns,
    })
    req = urllib.request.Request(f"{s.loki_url}/loki/api/v1/query_range?{params}")
    if s.loki_user:
        token = base64.b64encode(f"{s.loki_user}:{s.loki_password}".encode()).decode()
        req.add_header("Authorization", f"Basic {token}")
    if s.loki_tenant:
        req.add_header("X-Scope-OrgID", s.loki_tenant)
    try:
        with urllib.request.urlopen(req, timeout=30) as r:
            res = json.load(r)
    except urllib.error.HTTPError as exc:
        detail = exc.read()[:200].decode(errors="replace")
        return fail("err_loki_http", code=exc.code, detail=detail), 502
    except TimeoutError:
        return fail("err_loki_timeout"), 504
    except Exception as exc:
        return fail("err_loki_unreachable", error=str(exc)), 502
    lines = []
    for stream in res.get("data", {}).get("result", []):
        labels = stream.get("stream", {})
        for ts, line in stream.get("values", []):
            lines.append({
                "ts": datetime.fromtimestamp(int(ts) / 1e9).strftime("%Y-%m-%d %H:%M:%S"),
                "host": labels.get("host", ""), "job": labels.get("job", ""), "line": line[:400],
            })
    lines.sort(key=lambda x: x["ts"], reverse=True)
    return {"ok": True, "lines": lines[:s.loki_limit]}, 200