"""Reversibly disable Windows telemetry / data collection.

Targets, in order of leak severity:

  • DiagTrack service     — primary telemetry pipeline ("Connected User
                            Experiences and Telemetry"). Phones home every
                            few minutes.
  • dmwappushservice      — WAP Push routing (DSMS), historically tied to
                            DiagTrack. Disabled when present.
  • AllowTelemetry (HKLM) — DataCollection policy. 0 = lowest tier.
  • Activity History      — EnableActivityFeed / Upload / PublishUserActivities.
  • Advertising ID        — HKCU\\…\\AdvertisingInfo\\Enabled.
  • Cortana Bing search   — HKCU\\…\\Search\\BingSearchEnabled.
  • Scheduled tasks       — Compatibility Appraiser + CEIP Consolidator
                            + UsbCeip.

Every change is snapshotted before mutation; ``restore_telemetry`` rolls back.
"""
from __future__ import annotations

import json
import logging
import subprocess
import winreg
from dataclasses import dataclass, field

from .firewall import is_admin

log = logging.getLogger(__name__)

_SERVICES = ("DiagTrack", "dmwappushservice")

_SCHED_TASKS = (
    r"\Microsoft\Windows\Application Experience\Microsoft Compatibility Appraiser",
    r"\Microsoft\Windows\Customer Experience Improvement Program\Consolidator",
    r"\Microsoft\Windows\Customer Experience Improvement Program\UsbCeip",
)

# (hive, path, value_name, dword_when_killed)
_REG_KILLS = (
    (winreg.HKEY_LOCAL_MACHINE,
     r"SOFTWARE\Policies\Microsoft\Windows\DataCollection",
     "AllowTelemetry", 0),
    (winreg.HKEY_LOCAL_MACHINE,
     r"SOFTWARE\Policies\Microsoft\Windows\System",
     "EnableActivityFeed", 0),
    (winreg.HKEY_LOCAL_MACHINE,
     r"SOFTWARE\Policies\Microsoft\Windows\System",
     "PublishUserActivities", 0),
    (winreg.HKEY_LOCAL_MACHINE,
     r"SOFTWARE\Policies\Microsoft\Windows\System",
     "UploadUserActivities", 0),
    (winreg.HKEY_LOCAL_MACHINE,
     r"SOFTWARE\Policies\Microsoft\Windows\CloudContent",
     "DisableWindowsConsumerFeatures", 1),
    (winreg.HKEY_CURRENT_USER,
     r"Software\Microsoft\Windows\CurrentVersion\AdvertisingInfo",
     "Enabled", 0),
    (winreg.HKEY_CURRENT_USER,
     r"Software\Microsoft\Windows\CurrentVersion\Search",
     "BingSearchEnabled", 0),
    (winreg.HKEY_CURRENT_USER,
     r"Software\Microsoft\Windows\CurrentVersion\Search",
     "CortanaConsent", 0),
)


@dataclass
class TelemetrySnapshot:
    service_state: dict[str, str] = field(default_factory=dict)   # name -> 'auto'|'manual'|'disabled'|'absent'
    sched_tasks_state: dict[str, bool] = field(default_factory=dict)  # task -> was_enabled
    reg_values: list[dict] = field(default_factory=list)


def _run(args: list[str], timeout: float = 12.0) -> tuple[int, str, str]:
    try:
        r = subprocess.run(
            args, capture_output=True, text=True, timeout=timeout,
            creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
        )
        return r.returncode, r.stdout or "", r.stderr or ""
    except Exception as e:
        return 1, "", str(e)


# ── service helpers ──────────────────────────────────────────────────────────

def _service_start_type(name: str) -> str:
    code, out, _ = _run(["sc", "qc", name])
    if code != 0:
        return "absent"
    for ln in out.splitlines():
        s = ln.strip().lower()
        if "start_type" in s:
            if "2" in s or "auto" in s:
                return "auto"
            if "3" in s or "demand" in s or "manual" in s:
                return "manual"
            if "4" in s or "disabled" in s:
                return "disabled"
    return "absent"


def _set_service_disabled(name: str) -> bool:
    _run(["sc", "stop", name], timeout=15)
    code, _, _ = _run(["sc", "config", name, "start=", "disabled"])
    return code == 0


def _set_service_start(name: str, want: str) -> bool:
    start_value = {"auto": "auto", "manual": "demand", "disabled": "disabled"}.get(want)
    if not start_value:
        return False
    _run(["sc", "stop", name], timeout=10)
    code, _, _ = _run(["sc", "config", name, "start=", start_value])
    if want == "auto":
        _run(["sc", "start", name], timeout=15)
    return code == 0


# ── scheduled-task helpers ──────────────────────────────────────────────────

def _task_is_enabled(name: str) -> bool | None:
    code, out, _ = _run(["schtasks", "/Query", "/TN", name])
    if code != 0:
        return None
    for ln in out.splitlines():
        s = ln.strip().lower()
        if s.startswith("status:") or "ready" in s or "disabled" in s:
            if "disabled" in s:
                return False
            if "ready" in s or "running" in s:
                return True
    return True  # assume ready when query succeeded


def _set_task_state(name: str, enable: bool) -> bool:
    code, _, _ = _run([
        "schtasks", "/Change", "/TN", name,
        "/Enable" if enable else "/Disable",
    ])
    return code == 0


# ── registry helpers ────────────────────────────────────────────────────────

def _read_reg_dword(hive: int, path: str, name: str) -> tuple[bool, int | None]:
    """Returns (present, value)."""
    try:
        with winreg.OpenKey(hive, path, 0, winreg.KEY_QUERY_VALUE) as key:
            v, _ = winreg.QueryValueEx(key, name)
            return True, int(v)
    except OSError:
        return False, None


def _write_reg_dword(hive: int, path: str, name: str, value: int) -> bool:
    try:
        with winreg.CreateKeyEx(hive, path, 0, winreg.KEY_SET_VALUE) as key:
            winreg.SetValueEx(key, name, 0, winreg.REG_DWORD, value)
        return True
    except OSError:
        return False


def _delete_reg_value(hive: int, path: str, name: str) -> None:
    try:
        with winreg.OpenKey(hive, path, 0, winreg.KEY_SET_VALUE) as key:
            winreg.DeleteValue(key, name)
    except OSError:
        pass


# ── public API ──────────────────────────────────────────────────────────────

def telemetry_status() -> dict[str, str]:
    """Read-only snapshot for the audit panel."""
    out: dict[str, str] = {}
    for svc in _SERVICES:
        out[f"svc.{svc}"] = _service_start_type(svc)
    pres, val = _read_reg_dword(
        winreg.HKEY_LOCAL_MACHINE,
        r"SOFTWARE\Policies\Microsoft\Windows\DataCollection",
        "AllowTelemetry",
    )
    out["AllowTelemetry"] = "0 (killed)" if pres and val == 0 else (
        f"{val}" if pres else "unset"
    )
    pres, val = _read_reg_dword(
        winreg.HKEY_CURRENT_USER,
        r"Software\Microsoft\Windows\CurrentVersion\AdvertisingInfo",
        "Enabled",
    )
    out["AdvertisingID"] = "off" if pres and val == 0 else (
        "on" if pres and val == 1 else "unset"
    )
    return out


def engage_telemetry_kill() -> tuple[TelemetrySnapshot | None, list[str]]:
    logs: list[str] = []
    if not is_admin():
        return None, ["Telemetry kill skipped — needs Administrator."]
    snap = TelemetrySnapshot()

    for svc in _SERVICES:
        st = _service_start_type(svc)
        snap.service_state[svc] = st
        if st == "absent":
            continue
        if _set_service_disabled(svc):
            logs.append(f"Service '{svc}' stopped + disabled (was {st}).")
        else:
            logs.append(f"Service '{svc}' could not be disabled.")

    for task in _SCHED_TASKS:
        was = _task_is_enabled(task)
        if was is None:
            continue
        snap.sched_tasks_state[task] = was
        if was:
            if _set_task_state(task, enable=False):
                logs.append(f"Scheduled task disabled: {task}")

    for hive, path, name, kill_val in _REG_KILLS:
        present, prev = _read_reg_dword(hive, path, name)
        snap.reg_values.append({
            "hive": "HKLM" if hive == winreg.HKEY_LOCAL_MACHINE else "HKCU",
            "path": path,
            "name": name,
            "present": present,
            "prev": prev,
        })
        if _write_reg_dword(hive, path, name, kill_val):
            logs.append(f"reg: {('HKLM' if hive == winreg.HKEY_LOCAL_MACHINE else 'HKCU')}\\{path}\\{name} → {kill_val}")
        else:
            logs.append(f"reg write failed: {path}\\{name}")

    return snap, logs


def restore_telemetry(snap: TelemetrySnapshot | None) -> list[str]:
    if snap is None or not is_admin():
        return []
    logs: list[str] = []

    for svc, prev in snap.service_state.items():
        if prev == "absent":
            continue
        if _set_service_start(svc, prev):
            logs.append(f"Service '{svc}' restored → {prev}.")

    for task, was in snap.sched_tasks_state.items():
        if was:
            _set_task_state(task, enable=True)
            logs.append(f"Scheduled task re-enabled: {task}")

    for entry in snap.reg_values:
        hive = winreg.HKEY_LOCAL_MACHINE if entry["hive"] == "HKLM" else winreg.HKEY_CURRENT_USER
        path = entry["path"]
        name = entry["name"]
        if entry["present"]:
            _write_reg_dword(hive, path, name, int(entry["prev"]))
            logs.append(f"reg restore: {entry['hive']}\\{path}\\{name} ← {entry['prev']}")
        else:
            _delete_reg_value(hive, path, name)
            logs.append(f"reg removed (was absent): {entry['hive']}\\{path}\\{name}")
    return logs


def snapshot_to_json(snap: TelemetrySnapshot | None) -> str:
    if snap is None:
        return ""
    return json.dumps({
        "services": snap.service_state,
        "tasks": snap.sched_tasks_state,
        "reg": snap.reg_values,
    })


def snapshot_from_json(raw: str) -> TelemetrySnapshot | None:
    if not raw:
        return None
    try:
        d = json.loads(raw)
        return TelemetrySnapshot(
            service_state=dict(d.get("services") or {}),
            sched_tasks_state={k: bool(v) for k, v in (d.get("tasks") or {}).items()},
            reg_values=list(d.get("reg") or []),
        )
    except Exception:
        return None