"""
Togen Print Windows Agent  (DVI-667 — Print System P1)

Mirrors vision_agent.py. Discovers the Windows host's installed printers and
reports them to the Togen server on a short poll. The server (Togen Admin →
Agents → Print) decides which printers are exposed for Togen documents/labels
and what published name each carries. On each check-in the agent drains the FIFO
command queue: a ``print`` command downloads the job PDF and prints it silently
via the bundled SumatraPDF (DVI-669 — Print System P3), then acks the result.

Quick-enroll (admin token → active immediately):
    print_agent.exe --server https://togen.example.com --enroll <admin-token>

Self-enroll (open-enroll secret → pending, admin must approve in the Print pane):
    print_agent.exe --server https://togen.example.com --enroll <open-enroll-secret>

Run (after enrollment; token is stored in Windows Credential Manager):
    print_agent.exe --server https://togen.example.com

Enroll-only mode (for installers — registers and exits, does not start the loop):
    print_agent.exe --server https://togen.example.com --enroll <token> --enroll-only
"""

import argparse
import json
import logging
import os
import platform
import shutil
import socket
import subprocess
import sys
import tempfile
import time
import uuid
import winreg
from datetime import datetime, timezone

import requests

try:
    import keyring
    _HAS_KEYRING = True
except ImportError:
    _HAS_KEYRING = False

try:
    import win32print
    _HAS_WIN32PRINT = True
except ImportError:
    _HAS_WIN32PRINT = False

AGENT_VERSION = "0.1.0"
_REG_PATH = r"Software\Togen\Print"
_REG_MACHINE_ID = "MachineId"
_REG_TOKEN_KEY = "AgentToken"        # fallback when keyring unavailable
_KEYRING_SERVICE = "TogenPrintAgent"
_KEYRING_USERNAME = "agent_token"
_DEFAULT_INTERVAL_S = 10             # board default D: 10s poll
_LOG_FMT = "%(asctime)s %(levelname)s %(message)s"


# ---------------------------------------------------------------------------
# Registry helpers — machine_id
# ---------------------------------------------------------------------------

def _reg_read(hive, path, key):
    try:
        h = winreg.OpenKey(hive, path)
        val, _ = winreg.QueryValueEx(h, key)
        winreg.CloseKey(h)
        return val
    except FileNotFoundError:
        return None


def _reg_write(hive, path, key, value):
    h = winreg.CreateKeyEx(hive, path, access=winreg.KEY_WRITE)
    winreg.SetValueEx(h, key, 0, winreg.REG_SZ, value)
    winreg.CloseKey(h)


def get_machine_id() -> str:
    """Return the stable machine UUID, generating and persisting it on first call.

    Canonical location: HKLM\\Software\\Togen\\Print\\MachineId.
    Falls back to HKCU when HKLM is not writable (non-admin session).
    """
    for hive in (winreg.HKEY_LOCAL_MACHINE, winreg.HKEY_CURRENT_USER):
        val = _reg_read(hive, _REG_PATH, _REG_MACHINE_ID)
        if val:
            return val

    machine_id = str(uuid.uuid4())
    for hive in (winreg.HKEY_LOCAL_MACHINE, winreg.HKEY_CURRENT_USER):
        try:
            _reg_write(hive, _REG_PATH, _REG_MACHINE_ID, machine_id)
            return machine_id
        except PermissionError:
            continue

    logging.getLogger(__name__).warning(
        "Could not persist machine_id to registry; value is session-only."
    )
    return machine_id


# ---------------------------------------------------------------------------
# Credential storage — agent_token (Windows Credential Manager via keyring)
# ---------------------------------------------------------------------------

def _store_agent_token(token: str):
    if _HAS_KEYRING:
        try:
            keyring.set_password(_KEYRING_SERVICE, _KEYRING_USERNAME, token)
        except Exception:
            pass
    # Write to HKLM so SYSTEM (scheduled task at startup) can read it back.
    # Fall back to HKCU if HKLM write fails (non-admin context).
    for hive in (winreg.HKEY_LOCAL_MACHINE, winreg.HKEY_CURRENT_USER):
        try:
            _reg_write(hive, _REG_PATH, _REG_TOKEN_KEY, token)
            return
        except PermissionError:
            continue


def _load_agent_token() -> str | None:
    if _HAS_KEYRING:
        try:
            val = keyring.get_password(_KEYRING_SERVICE, _KEYRING_USERNAME)
            if val:
                return val
        except Exception:
            pass
    # Check HKLM first (readable by SYSTEM), then fall back to HKCU.
    for hive in (winreg.HKEY_LOCAL_MACHINE, winreg.HKEY_CURRENT_USER):
        val = _reg_read(hive, _REG_PATH, _REG_TOKEN_KEY)
        if val:
            return val
    return None


# ---------------------------------------------------------------------------
# Printer discovery
# ---------------------------------------------------------------------------

# win32print printer status bit → human label (best-effort; 0 == ready).
_PRINTER_STATUS_FLAGS = [
    (0x00000001, "paused"),
    (0x00000002, "error"),
    (0x00000080, "offline"),
    (0x00000100, "paper_out"),
    (0x00000400, "busy"),
    (0x00200000, "printing"),
]


def _printer_status_label(status_int) -> str:
    """Map a win32print PRINTER_INFO_2 Status bitmask to a short label."""
    try:
        status_int = int(status_int)
    except (TypeError, ValueError):
        return "unknown"
    if status_int == 0:
        return "ready"
    labels = [name for bit, name in _PRINTER_STATUS_FLAGS if status_int & bit]
    return ",".join(labels) if labels else "ready"


def discover_printers() -> list[dict]:
    """Enumerate local + connected printers via win32print.EnumPrinters (level 2).

    Returns a list of dicts with the raw fields the server upserts into
    print_printers: raw_name, driver, port, is_default, status.
    Returns [] when pywin32 is unavailable (non-Windows / dev environments).
    """
    log = logging.getLogger(__name__)
    if not _HAS_WIN32PRINT:
        log.warning("pywin32 (win32print) not available — printer discovery skipped.")
        return []
    try:
        default_name = win32print.GetDefaultPrinter()
    except Exception:
        default_name = None
    flags = win32print.PRINTER_ENUM_LOCAL | win32print.PRINTER_ENUM_CONNECTIONS
    out = []
    try:
        for p in win32print.EnumPrinters(flags, None, 2):
            name = p.get("pPrinterName") or ""
            if not name:
                continue
            out.append({
                "raw_name": name,
                "driver": p.get("pDriverName") or "",
                "port": p.get("pPortName") or "",
                "is_default": bool(default_name and name == default_name),
                "status": _printer_status_label(p.get("Status")),
            })
    except Exception as exc:
        log.warning("Printer enumeration failed: %s", exc)
    return out


# ---------------------------------------------------------------------------
# Enrollment
# ---------------------------------------------------------------------------

def enroll(base_url: str, enrollment_token: str, machine_id: str, hostname: str) -> tuple[str, str]:
    """POST /print/agents/register and persist the returned agent_token.

    Returns (agent_token, status) where status is 'active' or 'pending'.
    Admin-issued machine-bound tokens → 'active'. NULL-machine admin tokens or
    the open-enroll shared secret → 'pending' (admin must approve in the Print pane).
    """
    log = logging.getLogger(__name__)
    url = f"{base_url.rstrip('/')}/print/agents/register"
    payload = {
        "enrollment_token": enrollment_token,
        "machine_id": machine_id,
        "hostname": hostname,
    }
    log.info("Enrolling with %s (machine_id=%s)", url, machine_id)
    resp = requests.post(url, json=payload, timeout=30)
    resp.raise_for_status()
    data = resp.json()
    agent_token = data["agent_token"]
    status = data.get("status", "active")
    _store_agent_token(agent_token)
    if status == "pending":
        log.info(
            "Enrolled as PENDING (machine_id=%s). An admin must approve this "
            "agent in Togen Admin → Agents → Print before printers are usable. "
            "The agent will poll and activate automatically.",
            machine_id,
        )
    else:
        log.info("Enrolled as ACTIVE — agent_token stored in credential manager.")
    return agent_token, status


# ---------------------------------------------------------------------------
# Check-in payload
# ---------------------------------------------------------------------------

def _get_logged_in_user() -> str:
    """DOMAIN\\username of the current user."""
    domain = os.environ.get("USERDOMAIN") or platform.node()
    user = os.environ.get("USERNAME") or os.environ.get("USER") or "unknown"
    return f"{domain}\\{user}"


def _build_payload(machine_id: str, hostname: str, uptime_s: int) -> dict:
    return {
        "machine_id": machine_id,
        "hostname": hostname,
        "logged_in_user": _get_logged_in_user(),
        "os_version": platform.version(),
        "agent_version": AGENT_VERSION,
        "status": "active",
        "printers": discover_printers(),
        "agent_uptime_s": uptime_s,
        "ts": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
    }


# ---------------------------------------------------------------------------
# Printing (DVI-669 — Print System P3) — silent print via bundled SumatraPDF
# ---------------------------------------------------------------------------

_SUMATRA_EXE = "SumatraPDF.exe"


def _find_sumatra() -> str | None:
    """Locate SumatraPDF.exe. Prefers the copy bundled with the frozen agent
    (PyInstaller extracts data files to sys._MEIPASS), then the directory next to
    the running exe, then PATH, then the common per-user install location."""
    candidates = []
    meipass = getattr(sys, "_MEIPASS", None)
    if meipass:
        candidates.append(os.path.join(meipass, _SUMATRA_EXE))
    exe_dir = os.path.dirname(os.path.abspath(sys.argv[0]))
    candidates.append(os.path.join(exe_dir, _SUMATRA_EXE))
    on_path = shutil.which(_SUMATRA_EXE) or shutil.which("SumatraPDF")
    if on_path:
        candidates.append(on_path)
    local = os.environ.get("LOCALAPPDATA")
    if local:
        candidates.append(os.path.join(local, "SumatraPDF", _SUMATRA_EXE))
    for c in candidates:
        if c and os.path.isfile(c):
            return c
    return None


def _download_job_pdf(base_url: str, agent_token: str, job_id: str) -> str:
    """Download a job's PDF to a temp file and return its path. Raises on failure."""
    url = f"{base_url.rstrip('/')}/print/jobs/{job_id}/file"
    resp = requests.get(
        url, headers={"Authorization": f"Bearer {agent_token}"}, timeout=60, stream=True
    )
    resp.raise_for_status()
    fd, path = tempfile.mkstemp(prefix=f"togen_print_{job_id}_", suffix=".pdf")
    try:
        with os.fdopen(fd, "wb") as fh:
            for chunk in resp.iter_content(chunk_size=65536):
                if chunk:
                    fh.write(chunk)
    except Exception:
        try:
            os.remove(path)
        except OSError:
            pass
        raise
    return path


def _print_pdf(pdf_path: str, printer: str, print_settings: str, copies: int) -> None:
    """Silently print ``pdf_path`` to ``printer`` via SumatraPDF, once per copy.

    SumatraPDF has no copies flag, so copies are issued as repeated print jobs.
    Raises RuntimeError if SumatraPDF is missing or any invocation fails."""
    sumatra = _find_sumatra()
    if not sumatra:
        raise RuntimeError(
            "SumatraPDF.exe not found (expected bundled with the agent or on PATH)."
        )
    base_cmd = [sumatra, "-print-to", printer, "-silent"]
    if print_settings:
        base_cmd += ["-print-settings", print_settings]
    base_cmd.append(pdf_path)
    for i in range(max(1, copies)):
        proc = subprocess.run(base_cmd, capture_output=True, text=True, timeout=180)
        if proc.returncode != 0:
            raise RuntimeError(
                f"SumatraPDF exited {proc.returncode} on copy {i + 1}: "
                f"{(proc.stderr or proc.stdout or '').strip()[:300]}"
            )


def _handle_print_command(base_url: str, agent_token: str, cmd: dict) -> tuple[str, str | None]:
    """Execute a 'print' command. Returns (ack_status, error_message)."""
    log = logging.getLogger(__name__)
    params = cmd.get("params") or {}
    printer = (params.get("printer") or "").strip()
    if not printer:
        return "error", "Print command missing target printer."
    copies = params.get("copies") or 1
    try:
        copies = int(copies)
    except (TypeError, ValueError):
        copies = 1
    print_settings = (params.get("options") or {}).get("print_settings") or ""

    pdf_path = None
    try:
        pdf_path = _download_job_pdf(base_url, agent_token, cmd["id"])
        log.info("Printing job %s → %s (%dx, settings=%r)",
                 cmd["id"], printer, copies, print_settings)
        _print_pdf(pdf_path, printer, print_settings, copies)
        return "ok", None
    except requests.RequestException as exc:
        return "error", f"Job PDF download failed: {exc}"
    except (RuntimeError, subprocess.SubprocessError, OSError) as exc:
        return "error", str(exc)
    finally:
        if pdf_path:
            try:
                os.remove(pdf_path)
            except OSError:
                pass


# ---------------------------------------------------------------------------
# Command ack
# ---------------------------------------------------------------------------

def _ack_command(base_url: str, agent_token: str, cmd_id: str,
                 status: str, error_message: str | None = None):
    """POST /print/commands/<id>/ack. Best-effort; logs and swallows failures."""
    log = logging.getLogger(__name__)
    url = f"{base_url.rstrip('/')}/print/commands/{cmd_id}/ack"
    body = {"status": status}
    if error_message:
        body["error_message"] = error_message
    try:
        requests.post(
            url, json=body,
            headers={"Authorization": f"Bearer {agent_token}"}, timeout=15,
        )
    except requests.RequestException as exc:
        log.warning("Command ack failed (%s): %s", cmd_id, exc)


# ---------------------------------------------------------------------------
# Check-in loop
# ---------------------------------------------------------------------------

def run_checkin_loop(base_url: str, agent_token: str, machine_id: str,
                     hostname: str, initial_pending: bool = False):
    """Main check-in loop: report discovered printers, drain the command queue.

    Printing of dequeued commands lands in a later phase; for now the agent
    acks any command it cannot execute so the server queue does not stall.
    """
    log = logging.getLogger(__name__)
    url = f"{base_url.rstrip('/')}/print/checkin"
    start = time.monotonic()
    interval = _DEFAULT_INTERVAL_S

    if initial_pending:
        log.info(
            "Agent is PENDING approval. Check-ins running so an admin can "
            "approve this agent in the Print pane; printers report meanwhile."
        )

    while True:
        uptime_s = int(time.monotonic() - start)
        payload = _build_payload(machine_id, hostname, uptime_s)

        try:
            resp = requests.post(
                url,
                json=payload,
                headers={"Authorization": f"Bearer {agent_token}"},
                timeout=15,
            )
            if resp.status_code == 401:
                log.error("agent_token rejected (401). Re-run with --enroll <token>.")
                sys.exit(1)
            resp.raise_for_status()
            data = resp.json()

            # Respect server-issued polling interval.
            if "next_interval_s" in data:
                new_interval = int(data["next_interval_s"])
                if new_interval != interval:
                    log.info("Check-in interval updated: %ds → %ds", interval, new_interval)
                    interval = new_interval

            # Drain the command queue. Print commands download the job PDF and
            # silently print it via SumatraPDF; anything else is acked as an
            # explicit error so the server queue advances.
            cmd = data.get("command")
            if cmd and cmd.get("id"):
                if cmd.get("action") == "print":
                    status, err = _handle_print_command(base_url, agent_token, cmd)
                    if err:
                        log.warning("Print job %s failed: %s", cmd["id"], err)
                    _ack_command(base_url, agent_token, cmd["id"], status, err)
                else:
                    log.warning("Unsupported command action: %s", json.dumps(cmd))
                    _ack_command(base_url, agent_token, cmd["id"], "error",
                                 f"Unsupported command action: {cmd.get('action')}")

            log.debug("Check-in OK (uptime=%ds, interval=%ds)", uptime_s, interval)

        except requests.RequestException as exc:
            log.warning("Check-in failed: %s", exc)

        time.sleep(interval)


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="Togen Print Windows Agent")
    parser.add_argument(
        "--server", required=True,
        help="Base URL of the Togen server (e.g. https://togen.example.com)",
    )
    parser.add_argument(
        "--enroll", metavar="ENROLLMENT_TOKEN",
        help=(
            "Enroll this machine. Accepts either an admin-issued short-lived token "
            "(quick-enroll → active) or the open-enroll shared secret from the admin "
            "panel (self-enroll → pending, admin must approve in the Print pane)."
        ),
    )
    parser.add_argument(
        "--enroll-only", action="store_true",
        help=(
            "Enroll and exit without starting the check-in loop. "
            "Used by install-print.ps1; the Scheduled Task handles the ongoing run."
        ),
    )
    parser.add_argument("--debug", action="store_true", help="Enable debug logging.")
    args = parser.parse_args()

    if args.enroll_only and not args.enroll:
        parser.error("--enroll-only requires --enroll <token>")

    logging.basicConfig(
        level=logging.DEBUG if args.debug else logging.INFO,
        format=_LOG_FMT,
    )
    log = logging.getLogger(__name__)

    machine_id = get_machine_id()
    hostname = socket.gethostname()

    initial_pending = False
    if args.enroll:
        _, status = enroll(args.server, args.enroll, machine_id, hostname)
        initial_pending = (status == "pending")
        if args.enroll_only:
            # Exit code 0 = active, 2 = pending (install-print.ps1 prints guidance).
            sys.exit(2 if initial_pending else 0)

    agent_token = _load_agent_token()
    if not agent_token:
        log.error("No agent_token found. Enroll first: --enroll <token>")
        sys.exit(1)

    if not _HAS_WIN32PRINT:
        log.warning(
            "pywin32 not installed — printer discovery disabled. "
            "Install it: pip install pywin32"
        )

    log.info("Print agent starting (machine_id=%s, hostname=%s)", machine_id, hostname)
    run_checkin_loop(args.server, agent_token, machine_id, hostname,
                     initial_pending=initial_pending)


if __name__ == "__main__":
    main()