import json
import os
import re
import threading
import time
import urllib.request
import glob
import gzip
import ipaddress
from collections import Counter, defaultdict
from datetime import datetime, timedelta
from pathlib import Path

try:
    import docker
except ImportError:
    docker = None

GEO_CACHE = Path(os.environ.get("GEO_CACHE", "/cache/geo.json"))
HEADSCALE_CONTAINER = os.environ.get("HEADSCALE_CONTAINER", "headscale")

_client = None
_cache = {}


def _ttl(seconds):
    def deco(fn):
        def wrapper(*a, **kw):
            key = fn.__name__
            now = time.time()
            item = _cache.get(key)
            if item and now - item[0] < seconds:
                return item[1]
            val = fn(*a, **kw)
            _cache[key] = (now, val)
            return val
        return wrapper
    return deco


def _docker_client():
    global _client
    if docker is None:
        return None
    if _client is None:
        _client = docker.DockerClient(base_url="unix:///var/run/docker.sock", timeout=30)
    return _client


# Two timestamp formats appear in headscale logs (old + new)
TLS_ERR_RE = re.compile(
    r"(\d{4}/\d{2}/\d{2} \d{2}:\d{2}:\d{2}) .*TLS handshake error from (\d+\.\d+\.\d+\.\d+):\d+"
)
AUTH_ERR_RE = re.compile(
    r"(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}Z) .*client_address=(\d+\.\d+\.\d+\.\d+):\d+"
)

PRIVATE_PREFIXES = ("10.", "172.", "192.168.", "127.", "100.64.")


def _is_public(ip):
    return not ip.startswith(PRIVATE_PREFIXES)


def _load_geo_cache():
    if GEO_CACHE.exists():
        try:
            return json.loads(GEO_CACHE.read_text())
        except Exception:
            return {}
    return {}


def _save_geo_cache(cache):
    try:
        GEO_CACHE.parent.mkdir(parents=True, exist_ok=True)
        GEO_CACHE.write_text(json.dumps(cache))
    except Exception:
        pass


def geolocate(ips):
    cache = _load_geo_cache()
    missing = [ip for ip in ips if ip not in cache and _is_public(ip)]
    for i in range(0, len(missing), 100):
        batch = missing[i : i + 100]
        try:
            req = urllib.request.Request(
                "http://ip-api.com/batch",
                data=json.dumps(batch).encode(),
                headers={"Content-Type": "application/json"},
            )
            data = json.loads(urllib.request.urlopen(req, timeout=20).read())
            for ip, info in zip(batch, data):
                if info and info.get("status") == "success":
                    cache[ip] = {
                        "country": info.get("country"),
                        "cc": info.get("countryCode"),
                        "org": info.get("org") or info.get("isp"),
                        "as": info.get("as"),
                    }
        except Exception:
            pass
        time.sleep(1.1)
    if missing:
        _save_geo_cache(cache)
    return cache


def fetch_headscale_logs(hours=168, tail=300000):
    client = _docker_client()
    if client is None:
        return None
    try:
        c = client.containers.get(HEADSCALE_CONTAINER)
        since = datetime.utcnow() - timedelta(hours=hours)
        return c.logs(since=since, tail=tail, timestamps=False).decode(errors="replace")
    except Exception as e:
        return None


def parse_attempts(logs):
    attempts = defaultdict(lambda: {"count": 0, "first": None, "last": None})
    for line in logs.splitlines():
        m = TLS_ERR_RE.search(line) or AUTH_ERR_RE.search(line)
        if not m:
            continue
        ts_str, ip = m.group(1), m.group(2)
        a = attempts[ip]
        a["count"] += 1
        a["last"] = ts_str
        if a["first"] is None:
            a["first"] = ts_str
    return attempts


def _flagged_countries():
    raw = os.environ.get("FLAGGED_COUNTRIES", "China,Russia,Iran,North Korea")
    return {c.strip() for c in raw.split(",") if c.strip()}


def send_webhook(url, text):
    try:
        req = urllib.request.Request(
            url,
            data=json.dumps({"content": text}).encode(),
            headers={"Content-Type": "application/json"},
        )
        urllib.request.urlopen(req, timeout=10)
        return True
    except Exception as e:
        print("webhook error:", e)
        return False


def check_alerts():
    webhook = os.environ.get("WEBHOOK_URL", "")
    if not webhook:
        return {"status": "disabled"}
    threshold = int(os.environ.get("ALERT_THRESHOLD", "50"))
    state_file = Path("/cache/alerts_state.json")

    geo_cache = _load_geo_cache()
    known = set(geo_cache.keys())
    state = json.loads(state_file.read_text()) if state_file.exists() else {}

    logs = fetch_headscale_logs(24, 200000)
    if logs is None:
        return {"status": "error", "msg": "headscale unreachable"}
    attempts = parse_attempts(logs)

    geo = geolocate(list(attempts.keys()))
    fired = []

    for ip, a in attempts.items():
        if not _is_public(ip) or ip == "73.201.12.182":
            continue
        g = geo.get(ip, {})
        entry = state.setdefault(ip, {})
        if ip not in known and not entry.get("seen"):
            entry["seen"] = True
            send_webhook(webhook, f"🆕 New IP probing headscale: {ip} — {g.get('country','?')} / {g.get('org','?')} — {a['count']} attempts")
            fired.append({"ip": ip, "type": "new"})
        if a["count"] >= threshold and not entry.get("thresh"):
            entry["thresh"] = True
            send_webhook(webhook, f"⚠️ High attempts: {ip} — {g.get('country','?')} / {g.get('org','?')} — {a['count']} attempts")
            fired.append({"ip": ip, "type": "threshold"})

    state_file.parent.mkdir(parents=True, exist_ok=True)
    state_file.write_text(json.dumps(state))
    return {"status": "ok", "fired": fired}


def alert_loop():
    interval = int(os.environ.get("ALERT_INTERVAL", "300"))
    while True:
        try:
            check_alerts()
        except Exception as e:
            print("alert loop error:", e)
        time.sleep(interval)


def security_summary(hours=168, tail=300000):
    logs = fetch_headscale_logs(hours, tail)
    if logs is None:
        return {"error": "headscale container not reachable"}
    attempts = parse_attempts(logs)
    geo = geolocate(list(attempts.keys()))

    total_attempts = sum(a["count"] for a in attempts.values())
    public = [(ip, a) for ip, a in attempts.items() if _is_public(ip)]

    rows = []
    for ip, a in public:
        g = geo.get(ip, {})
        rows.append(
            {
                "ip": ip,
                "country": g.get("country", "?"),
                "cc": g.get("cc", ""),
                "org": g.get("org", "?"),
                "as": g.get("as", "?"),
                "count": a["count"],
                "first": a["first"],
                "last": a["last"],
            }
        )
    rows.sort(key=lambda r: r["count"], reverse=True)

    flagged = _flagged_countries()
    for r in rows:
        r["flagged"] = r["country"] in flagged
    flagged_attempts = sum(r["count"] for r in rows if r["flagged"])
    flagged_ips = sum(1 for r in rows if r["flagged"])

    country_counts = Counter(r["country"] for r in rows)
    org_counts = Counter(r["org"] for r in rows if r["org"] and r["org"] != "?")

    return {
        "window_hours": hours,
        "total_attempts": total_attempts,
        "unique_ips": len(rows),
        "countries": [{"country": c, "count": n} for c, n in country_counts.most_common()],
        "top_orgs": [{"org": o, "count": n} for o, n in org_counts.most_common(15)],
        "flagged": {"countries": sorted(flagged), "attempts": flagged_attempts, "ips": flagged_ips},
        "attempts": rows[:300],
    }


@_ttl(60)
def ip_attempts(ip, hours=168):
    logs = fetch_headscale_logs(hours, 300000)
    if logs is None:
        return {"ip": ip, "count": 0, "attempts": [], "error": "headscale unreachable"}
    attempts = []
    for line in logs.splitlines():
        m = TLS_ERR_RE.search(line) or AUTH_ERR_RE.search(line)
        if not m:
            continue
        ts, src = m.group(1), m.group(2)
        if src != ip:
            continue
        evt = "TLS handshake failure" if "TLS" in line else ("auth error" if "client_address" in line else "connection attempt")
        attempts.append({"time": ts, "event": evt})
    attempts.reverse()
    return {"ip": ip, "count": len(attempts), "attempts": attempts[:500]}


def registered_nodes():
    client = _docker_client()
    if client is None:
        return []
    try:
        c = client.containers.get(HEADSCALE_CONTAINER)
        exit_code, output = c.exec_run(["headscale", "nodes", "list", "-o", "json"])
        if exit_code != 0:
            return []
        return json.loads(output.decode(errors="replace"))
    except Exception:
        return []


HOSTLOGS_DIR = Path(os.environ.get("HOSTLOGS_DIR", "/hostlogs"))
ACCEPTED_RE = re.compile(r"Accepted (?:password|publickey|keyboard-interactive) for (\S+) from ([0-9.]+)")
FAILED_RE = re.compile(r"Failed password for (?:invalid user )?(\S+) from ([0-9.]+)")


@_ttl(30)
def ssh_logins():
    accepted = Counter()
    failed = Counter()
    for path in sorted(glob.glob(str(HOSTLOGS_DIR / "auth.log*"))):
        try:
            if path.endswith(".gz"):
                content = gzip.open(path, "rt", errors="replace").read()
            else:
                content = open(path, errors="replace").read()
        except Exception:
            continue
        for line in content.splitlines():
            m = ACCEPTED_RE.search(line)
            if m:
                accepted[m.group(2)] += 1
                continue
            m = FAILED_RE.search(line)
            if m:
                failed[m.group(2)] += 1
    return {
        "accepted": [{"ip": ip, "count": c} for ip, c in accepted.most_common(20)],
        "failed": [{"ip": ip, "count": c} for ip, c in failed.most_common(20)],
        "total_accepted": sum(accepted.values()),
        "total_failed": sum(failed.values()),
    }


@_ttl(120)
def nextcloud_status():
    client = _docker_client()
    if client is None:
        return {"error": "docker not available"}
    try:
        c = client.containers.get("nextcloud")
    except Exception:
        return {"error": "nextcloud container not found"}
    result = {"users": [], "failed_logins": [], "total_failed": 0}
    try:
        _, out = c.exec_run(
            ["sh", "-c", "grep -iE 'Login failed' /var/www/html/data/nextcloud.log 2>/dev/null | tail -100"],
            stdout=True,
        )
        for line in out.decode(errors="replace").splitlines():
            try:
                d = json.loads(line)
                result["failed_logins"].append({
                    "time": d.get("time"),
                    "ip": d.get("remoteAddr"),
                    "user": d.get("user"),
                    "message": d.get("message"),
                })
            except Exception:
                pass
        result["total_failed"] = len(result["failed_logins"])
    except Exception:
        pass
    try:
        _, out = c.exec_run(
            ["sh", "-c", "cd /var/www/html && php occ user:list --no-warnings 2>/dev/null"],
            user="www-data", stdout=True,
        )
        for line in out.decode(errors="replace").splitlines():
            m = re.search(r"-\s+(\S+):", line)
            if m:
                result["users"].append(m.group(1))
    except Exception:
        pass
    return result


def _is_local(ip):
    try:
        a = ipaddress.ip_address(ip)
        return a.is_private or a.is_loopback or a.is_link_local or ip == "73.201.12.182"
    except ValueError:
        return False


def compromise_status():
    ssh = ssh_logins()
    nodes = registered_nodes()
    nc = nextcloud_status()
    warnings = []

    unknown_ssh = [a["ip"] for a in ssh["accepted"] if not _is_local(a["ip"])]
    if unknown_ssh:
        warnings.append(f"SSH logins from unknown IPs: {', '.join(unknown_ssh)}")
    if ssh["total_failed"] > 100:
        warnings.append(f"High SSH failed-password attempts: {ssh['total_failed']}")
    if nc.get("total_failed", 0) > 20:
        warnings.append(f"Nextcloud failed logins: {nc['total_failed']}")

    return {
        "verdict": "review" if warnings else "clean",
        "warnings": warnings,
        "ssh": ssh,
        "headscale_nodes": nodes,
        "nextcloud": nc,
    }