"""Infer in-progress imaging sessions from server-side logs.

The /imaging/status endpoint only sees clients that have reached the WinPE
startnet.cmd push (stage_index 2). A bay that gets stuck earlier (no DHCP,
TFTP fail, boot.wim download stall) is invisible to the operator.

This module tails dnsmasq leases + Apache access log + per-host Samba logs
and synthesizes a "session" record per active MAC, with a coarse inferred
stage based on which boot assets the client has actually fetched. Output
shape matches services/imaging_status so list_sessions() can merge both.

Inferred sessions carry source="inferred" and never overwrite a real
client-pushed session for the same serial (correlated by MAC when present).
"""
from __future__ import annotations

import os
import re
import time
from datetime import datetime, timezone
from typing import Optional

import config

# Active window: an inferred session is shown if any evidence is newer than
# this many seconds. Past that we assume the bay is idle / done / off.
INFERRED_ACTIVE_WINDOW_S = 90 * 60  # 90 min

# Tail size caps so a giant log doesn't pull the whole file into memory on
# each dashboard refresh.
APACHE_TAIL_BYTES = 512 * 1024  # 512 KB
SAMBA_TAIL_BYTES = 64 * 1024   # 64 KB per file
SYSLOG_TAIL_BYTES = 256 * 1024  # 256 KB

# Path prefixes Apache serves, mapped to coarse imaging stage signals.
# Order matters: we take the highest-numbered match (latest in the boot chain).
_APACHE_STAGE_HITS = [
    ("/menu.ipxe",            ("stage_0_menu",      0)),
    ("/win11/boot/",          ("stage_0_boot_pre", 0)),
    ("/win11/efi/",           ("stage_0_boot_pre", 0)),
    ("/win11/sources/boot.wim", ("stage_1_wim_get", 1)),
    ("/win11/sources/",       ("stage_1_wim_get",  1)),
]

# TFTP bootloader fetches arrive via dnsmasq, not Apache. Mapped same way.
_TFTP_STAGE_HITS = [
    ("undionly.kpxe",   ("stage_0_tftp_bios", 0)),
    ("ipxe.efi",        ("stage_0_tftp_uefi", 0)),
]

# Apache combined log: <ip> - - [DD/Mon/YYYY:HH:MM:SS +ZZZZ] "GET /path HTTP/x.y" status bytes "ref" "ua"
_APACHE_RE = re.compile(
    r'^(?P<ip>\S+)\s+\S+\s+\S+\s+'
    r'\[(?P<ts>[^\]]+)\]\s+'
    r'"(?P<method>\S+)\s+(?P<path>\S+)\s+\S+"\s+'
    r'(?P<status>\d+)\s+(?P<bytes>\S+)'
)
_APACHE_TS_FMT = "%d/%b/%Y:%H:%M:%S %z"

# dnsmasq syslog: "<MMM DD HH:MM:SS> host dnsmasq-tftp[pid]: sent /tftp/path to ip"
# or "dnsmasq-dhcp[pid]: DHCPACK(...) ip mac hostname"
_SYSLOG_DNSMASQ_RE = re.compile(
    r'^(?P<mon>\w{3})\s+(?P<day>\d+)\s+(?P<time>\d{2}:\d{2}:\d{2})\s+'
    r'\S+\s+dnsmasq-(?P<kind>\w+)\[\d+\]:\s+(?P<msg>.+)$'
)


def _tail_bytes(path: str, max_bytes: int) -> str:
    """Read last max_bytes of a text log, return decoded string. Returns
    empty string if file is missing / unreadable. Skips partial first line."""
    try:
        size = os.path.getsize(path)
    except OSError:
        return ""
    start = max(0, size - max_bytes)
    try:
        with open(path, "rb") as f:
            f.seek(start)
            data = f.read()
    except OSError:
        return ""
    text = data.decode("utf-8", errors="replace")
    if start > 0:
        # Drop partial first line.
        nl = text.find("\n")
        if nl >= 0:
            text = text[nl + 1:]
    return text


def _now_iso() -> str:
    return datetime.now().astimezone().isoformat(timespec="seconds")


def _epoch_to_iso(epoch: float) -> str:
    return datetime.fromtimestamp(epoch).astimezone().isoformat(timespec="seconds")


def parse_leases() -> list[dict]:
    """dnsmasq.leases format: <expires_epoch> <mac> <ip> <hostname> <client_id>."""
    out = []
    try:
        with open(config.DNSMASQ_LEASES, "r") as f:
            for line in f:
                parts = line.split()
                if len(parts) < 4:
                    continue
                try:
                    expires = int(parts[0])
                except ValueError:
                    continue
                out.append({
                    "expires": expires,
                    "mac": parts[1].lower(),
                    "ip": parts[2],
                    "hostname": parts[3] if parts[3] != "*" else "",
                })
    except OSError:
        pass
    return out


def parse_apache_hits() -> dict:
    """Return {ip: {"last_ts": epoch, "max_stage": int, "stage_label": str,
    "paths": [recent_paths]}}."""
    out: dict = {}
    text = _tail_bytes(config.APACHE_ACCESS_LOG, APACHE_TAIL_BYTES)
    for line in text.splitlines():
        m = _APACHE_RE.match(line)
        if not m:
            continue
        ip = m.group("ip")
        path = m.group("path").lower()
        try:
            ts = datetime.strptime(m.group("ts"), _APACHE_TS_FMT).timestamp()
        except ValueError:
            continue
        # Match against known stage prefixes; keep the highest-numbered.
        for prefix, (label, stage) in _APACHE_STAGE_HITS:
            if path.startswith(prefix):
                rec = out.setdefault(ip, {
                    "last_ts": 0.0, "max_stage": -1,
                    "stage_label": "", "paths": [],
                })
                if ts > rec["last_ts"]:
                    rec["last_ts"] = ts
                if stage > rec["max_stage"]:
                    rec["max_stage"] = stage
                    rec["stage_label"] = label
                rec["paths"].append(path)
                # Cap path list so a busy client doesn't explode the dict.
                rec["paths"] = rec["paths"][-10:]
                break
    return out


def parse_syslog_tftp() -> dict:
    """Return {mac_or_ip: {"last_ts": epoch, "max_stage": int,
    "stage_label": str, "evidence": [str]}} from dnsmasq syslog lines."""
    out: dict = {}
    text = _tail_bytes(config.DNSMASQ_SYSLOG, SYSLOG_TAIL_BYTES)
    year = datetime.now().year
    for line in text.splitlines():
        m = _SYSLOG_DNSMASQ_RE.match(line)
        if not m:
            continue
        kind = m.group("kind")
        msg = m.group("msg")
        # Best-effort ts parse. Syslog format has no year; assume current.
        try:
            ts = datetime.strptime(
                f"{year} {m.group('mon')} {m.group('day')} {m.group('time')}",
                "%Y %b %d %H:%M:%S"
            ).timestamp()
        except ValueError:
            continue

        if kind == "tftp":
            # "sent /srv/tftp/ipxe.efi to 172.16.9.42"
            tm = re.search(r'(?:sent|read|file)\s+(?P<file>\S+)\s+to\s+(?P<ip>\S+)', msg)
            if not tm:
                continue
            key = tm.group("ip")
            fname = os.path.basename(tm.group("file")).lower()
            for needle, (label, stage) in _TFTP_STAGE_HITS:
                if needle in fname:
                    rec = out.setdefault(key, {
                        "last_ts": 0.0, "max_stage": -1,
                        "stage_label": "", "evidence": [],
                    })
                    if ts > rec["last_ts"]:
                        rec["last_ts"] = ts
                    if stage > rec["max_stage"]:
                        rec["max_stage"] = stage
                        rec["stage_label"] = label
                    rec["evidence"].append(f"tftp:{fname}")
                    rec["evidence"] = rec["evidence"][-5:]
                    break
        elif kind == "dhcp":
            # DHCPACK lines confirm a lease handshake completed; useful as a
            # very-early "this MAC is alive" signal even before TFTP fires.
            dm = re.search(
                r'DHCP(?:ACK|OFFER|REQUEST)\([^)]+\)\s+'
                r'(?P<ip>\d+\.\d+\.\d+\.\d+)\s+'
                r'(?P<mac>[0-9a-f:]{17})',
                msg
            )
            if not dm:
                continue
            key = dm.group("ip")
            rec = out.setdefault(key, {
                "last_ts": 0.0, "max_stage": -1,
                "stage_label": "", "evidence": [],
            })
            if ts > rec["last_ts"]:
                rec["last_ts"] = ts
            if rec["max_stage"] < 0:
                rec["max_stage"] = 0
                rec["stage_label"] = "stage_0_dhcp"
            rec["evidence"].append(f"dhcp:{dm.group('mac')}")
            rec["evidence"] = rec["evidence"][-5:]
    return out


def parse_samba_hits() -> dict:
    """Return {ip: {"last_ts": epoch, "shares": set[str]}} for clients that
    have opened a Samba session in the recent tail. Samba writes per-host log
    files: /var/log/samba/log.<ip>. Mere existence + recent mtime = client
    has connected. Share names from inside the file are best-effort."""
    out: dict = {}
    log_dir = config.SAMBA_LOG_DIR
    if not os.path.isdir(log_dir):
        return out
    try:
        entries = os.listdir(log_dir)
    except OSError:
        return out
    for name in entries:
        if not name.startswith("log."):
            continue
        ip = name[len("log."):]
        # Skip the daemon's own log files (log.smbd, log.nmbd, log.winbindd).
        if ip in ("smbd", "nmbd", "winbindd") or "." not in ip:
            continue
        path = os.path.join(log_dir, name)
        try:
            st = os.stat(path)
        except OSError:
            continue
        rec = {"last_ts": st.st_mtime, "shares": set()}
        text = _tail_bytes(path, SAMBA_TAIL_BYTES)
        for share in ("winpeapps", "clonezilla", "enrollment", "image-upload", "blancco-reports"):
            if f"[{share}]" in text or f"connect to service {share}" in text:
                rec["shares"].add(share)
        out[ip] = rec
    return out


def infer_sessions(
    active_window_s: int = INFERRED_ACTIVE_WINDOW_S,
) -> list[dict]:
    """Build inferred-session records from all log sources. Keyed by MAC when
    available, else IP. Only returns sessions with evidence in the active
    window."""
    leases = parse_leases()
    apache = parse_apache_hits()
    samba = parse_samba_hits()
    tftp = parse_syslog_tftp()

    # IP -> lease (MAC, hostname).
    by_ip = {lease["ip"]: lease for lease in leases}

    cutoff = time.time() - active_window_s
    sessions: dict[str, dict] = {}

    def _ensure(ip: str) -> dict:
        lease = by_ip.get(ip, {})
        key = lease.get("mac") or ip
        rec = sessions.get(key)
        if rec is None:
            rec = {
                "source": "inferred",
                "serial": None,
                "mac": lease.get("mac", ""),
                "ip": ip,
                "hostname_target": lease.get("hostname", ""),
                "started_at": None,
                "last_updated": None,
                "stage_index": 0,
                "stage_total": 8,
                "current_stage": "",
                "status": "in_progress",
                "evidence": [],
                "log_tail": [],
            }
            sessions[key] = rec
        return rec

    def _bump_stage(rec: dict, stage: int, label: str):
        if stage > rec["stage_index"]:
            rec["stage_index"] = stage
            rec["current_stage"] = label

    def _touch(rec: dict, ts: float):
        iso = _epoch_to_iso(ts)
        if rec["started_at"] is None or iso < rec["started_at"]:
            rec["started_at"] = iso
        if rec["last_updated"] is None or iso > rec["last_updated"]:
            rec["last_updated"] = iso

    for ip, hit in tftp.items():
        if hit["last_ts"] < cutoff:
            continue
        rec = _ensure(ip)
        _touch(rec, hit["last_ts"])
        _bump_stage(rec, hit["max_stage"], hit["stage_label"])
        rec["evidence"].extend(hit["evidence"])

    for ip, hit in apache.items():
        if hit["last_ts"] < cutoff:
            continue
        rec = _ensure(ip)
        _touch(rec, hit["last_ts"])
        _bump_stage(rec, hit["max_stage"], hit["stage_label"])
        for p in hit["paths"]:
            rec["evidence"].append(f"http:{p}")

    for ip, hit in samba.items():
        if hit["last_ts"] < cutoff:
            continue
        rec = _ensure(ip)
        _touch(rec, hit["last_ts"])
        # WinPE-stage SMB activity = boot.wim applied + apps stage starting.
        _bump_stage(rec, 2, "stage_2_smb_winpe")
        for share in sorted(hit["shares"]):
            rec["evidence"].append(f"smb:{share}")

    # Trim evidence + bake log_tail for display.
    for rec in sessions.values():
        rec["evidence"] = rec["evidence"][-15:]
        rec["log_tail"] = list(rec["evidence"])

    return list(sessions.values())


def merge_with_client_sessions(
    client_sessions: list[dict],
    inferred: Optional[list[dict]] = None,
) -> list[dict]:
    """Return a single list. Real (client-pushed) sessions win over inferred
    ones for the same MAC. Inferred sessions are dropped when a real session
    with the same MAC already exists, since the real one has strictly more
    information."""
    if inferred is None:
        inferred = infer_sessions()

    real_macs = {(s.get("mac") or "").lower() for s in client_sessions if s.get("mac")}
    real_ips = {s.get("ip") for s in client_sessions if s.get("ip")}

    # Tag real sessions so the template can style them differently from
    # inferred ones (and not break older sessions that lack the field).
    for s in client_sessions:
        s.setdefault("source", "client")

    merged = list(client_sessions)
    for inf in inferred:
        mac = (inf.get("mac") or "").lower()
        ip = inf.get("ip")
        if mac and mac in real_macs:
            continue
        if ip and ip in real_ips:
            continue
        merged.append(inf)

    merged.sort(key=lambda s: s.get("last_updated") or "", reverse=True)
    return merged