#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import socket
import subprocess
import sys
import tempfile
import time
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
from urllib import request


ENV_FILE = Path("/etc/activitywatch/aw-server.env")


def load_env_file(path: Path) -> None:
    if not path.exists():
        return
    try:
        lines = path.read_text(encoding="utf-8").splitlines()
    except PermissionError:
        # systemd EnvironmentFile has already injected the variables for service runs.
        return
    for raw_line in lines:
        line = raw_line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        key = key.strip()
        value = value.strip().strip("'").strip('"')
        os.environ.setdefault(key, value)


def env(name: str, default: str) -> str:
    value = os.environ.get(name)
    return value if value not in (None, "") else default


def now_utc() -> datetime:
    return datetime.now(UTC)


def parse_ts(value: str | None) -> datetime | None:
    if not value:
        return None
    try:
        return datetime.fromisoformat(value.replace("Z", "+00:00")).astimezone(UTC)
    except ValueError:
        return None


def age_seconds(ts: datetime | None, now: datetime) -> int | None:
    if ts is None:
        return None
    return max(0, int((now - ts).total_seconds()))


def http_json(url: str, timeout: int = 20, attempts: int = 2, backoff_seconds: float = 0.5) -> Any:
    last_exc: Exception | None = None
    for attempt in range(max(1, attempts)):
        try:
            with request.urlopen(url, timeout=timeout) as resp:
                return json.loads(resp.read().decode("utf-8"))
        except Exception as exc:
            last_exc = exc
            if attempt + 1 >= max(1, attempts):
                break
            time.sleep(backoff_seconds * (2**attempt))
    assert last_exc is not None
    raise last_exc


def run_command(cmd: list[str]) -> tuple[int, str]:
    proc = subprocess.run(
        cmd,
        check=False,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )
    return proc.returncode, proc.stdout.strip()


def tcp_connect(host: str, port: int, timeout: float) -> tuple[bool, str]:
    try:
        with socket.create_connection((host, port), timeout=timeout):
            return True, "connected"
    except OSError as exc:
        return False, str(exc)


@dataclass
class CheckResult:
    name: str
    status: str
    summary: str
    details: dict[str, Any]


class Report:
    def __init__(self) -> None:
        self.results: list[CheckResult] = []

    def add(self, name: str, status: str, summary: str, **details: Any) -> None:
        self.results.append(CheckResult(name=name, status=status, summary=summary, details=details))

    @property
    def ok(self) -> bool:
        return not any(item.status == "fail" for item in self.results)

    def as_dict(self) -> dict[str, Any]:
        counts = {"ok": 0, "warn": 0, "fail": 0}
        for item in self.results:
            counts[item.status] = counts.get(item.status, 0) + 1
        return {
            "generated_at_utc": now_utc().isoformat().replace("+00:00", "Z"),
            "ok": self.ok,
            "counts": counts,
            "results": [
                {
                    "name": item.name,
                    "status": item.status,
                    "summary": item.summary,
                    "details": item.details,
                }
                for item in self.results
            ],
        }

    def render_text(self) -> str:
        icon = {"ok": "OK", "warn": "WARN", "fail": "FAIL"}
        lines = ["=== AW-RUS Health ===", f"Timestamp: {now_utc().isoformat().replace('+00:00', 'Z')}", ""]
        for item in self.results:
            lines.append(f"[{icon.get(item.status, item.status.upper())}] {item.name}: {item.summary}")
        lines.append("")
        payload = self.as_dict()
        lines.append(
            "Counts: ok={ok} warn={warn} fail={fail}".format(
                ok=payload["counts"]["ok"],
                warn=payload["counts"]["warn"],
                fail=payload["counts"]["fail"],
            )
        )
        lines.append(f"Overall: {'OK' if payload['ok'] else 'FAIL'}")
        return "\n".join(lines)


def latest_bucket_event(api_base: str, bucket_id: str) -> dict[str, Any] | None:
    events = http_json(f"{api_base}/buckets/{bucket_id}/events?limit=1", timeout=25)
    if isinstance(events, list) and events:
        events = [item for item in events if isinstance(item, dict)]
        if not events:
            return None
        events.sort(key=lambda item: item.get("timestamp") or "", reverse=True)
        return events[0]
    return None


def bucket_metadata_ts(buckets: dict[str, Any], bucket_id: str) -> datetime | None:
    bucket = buckets.get(bucket_id) if isinstance(buckets, dict) else None
    if not isinstance(bucket, dict):
        return None
    metadata = bucket.get("metadata")
    if not isinstance(metadata, dict):
        return None
    return parse_ts(metadata.get("end"))


def host_activity_from_worktime(event: dict[str, Any] | None, max_age_seconds: int) -> dict[str, Any]:
    now = now_utc()
    if not event:
        return {"fresh": False, "active": False, "age_seconds": None, "timestamp": None}
    ts = parse_ts(event.get("timestamp"))
    age = age_seconds(ts, now)
    data = event.get("data") or {}
    is_fresh = age is not None and age <= max_age_seconds
    is_active = bool(is_fresh and data.get("active"))
    return {
        "fresh": bool(is_fresh),
        "active": bool(is_active),
        "age_seconds": age,
        "timestamp": event.get("timestamp"),
        "data": data,
    }


def bucket_health(
    api_base: str,
    bucket_id: str,
    max_age_seconds: int,
    missing_status: str,
    stale_status: str,
) -> tuple[str, str, dict[str, Any]]:
    try:
        event = latest_bucket_event(api_base, bucket_id)
    except Exception as exc:
        return "fail", f"bucket query failed: {exc}", {"bucket": bucket_id}

    if not event:
        return missing_status, "no events", {"bucket": bucket_id}

    ts = parse_ts(event.get("timestamp"))
    age = age_seconds(ts, now_utc())
    details = {"bucket": bucket_id, "timestamp": event.get("timestamp"), "age_seconds": age}
    if age is None:
        return "warn", "timestamp parse failed", details
    if age > max_age_seconds:
        return stale_status, f"stale ({age}s)", details
    return "ok", f"fresh ({age}s)", details


def bucket_timestamp_health(
    api_base: str,
    buckets: dict[str, Any],
    bucket_id: str,
    max_age_seconds: int,
    missing_status: str,
    stale_status: str,
) -> tuple[str, str, dict[str, Any]]:
    metadata_ts = bucket_metadata_ts(buckets, bucket_id)
    if metadata_ts is not None:
        age = age_seconds(metadata_ts, now_utc())
        timestamp = metadata_ts.isoformat().replace("+00:00", "Z")
        details = {
            "bucket": bucket_id,
            "timestamp": timestamp,
            "age_seconds": age,
            "timestamp_source": "bucket_metadata.end",
        }
        if age is None:
            return "warn", "timestamp parse failed", details
        if age > max_age_seconds:
            return stale_status, f"stale ({age}s)", details
        return "ok", f"fresh ({age}s)", details

    return bucket_health(api_base, bucket_id, max_age_seconds, missing_status, stale_status)


def guard_bucket_health(
    api_base: str,
    host: str,
    max_age_seconds: int,
    required: bool,
) -> tuple[str, str, dict[str, Any]]:
    bucket_id = f"aw-rus-collector-guard_{host}"
    try:
        event = latest_bucket_event(api_base, bucket_id)
    except Exception as exc:
        status = "fail" if required else "warn"
        return status, f"guard bucket query failed: {exc}", {"bucket": bucket_id}

    if not event:
        status = "fail" if required else "warn"
        return status, "no guard heartbeat", {"bucket": bucket_id, "required": required}

    ts = parse_ts(event.get("timestamp"))
    age = age_seconds(ts, now_utc())
    data = event.get("data") if isinstance(event.get("data"), dict) else {}
    guard_status = str(data.get("status") or "unknown").lower()
    details = {
        "bucket": bucket_id,
        "timestamp": event.get("timestamp"),
        "age_seconds": age,
        "required": required,
        "guard_status": guard_status,
        "mode": data.get("mode"),
        "live_session_count": data.get("liveSessionCount"),
        "problems": data.get("problems", []),
        "actions": data.get("actions", []),
    }

    stale_status = "fail" if required else "warn"
    if age is None:
        return "warn", "guard timestamp parse failed", details
    if age > max_age_seconds:
        return stale_status, f"guard stale ({age}s)", details
    if guard_status in ("fail", "error"):
        return "fail" if required else "warn", f"guard reports {guard_status}", details
    if guard_status == "warn":
        return "warn", "guard reports warn", details
    return "ok", f"guard fresh ({age}s)", details


def latest_validation_report(validation_dir: Path) -> Path | None:
    candidates = sorted(
        (path for path in validation_dir.glob("*-aw_validate_ansible.json") if path.is_file()),
        key=lambda item: item.stat().st_mtime,
        reverse=True,
    )
    return candidates[0] if candidates else None


def write_atomic(path: Path, content: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=path.parent, delete=False) as handle:
        handle.write(content)
        tmp_name = handle.name
    os.replace(tmp_name, path)


def chmod_if_possible(path: Path, mode: int) -> None:
    try:
        path.chmod(mode)
    except OSError:
        pass


def check_wrapper(
    report: Report,
    name: str,
    cmd: list[str],
    json_mode: bool = False,
    failure_status: str = "fail",
) -> None:
    if not Path(cmd[0]).exists():
        report.add(name, "warn", "binary missing", command=cmd)
        return
    rc, output = run_command(cmd)
    details: dict[str, Any] = {"command": cmd, "returncode": rc}
    if json_mode:
        try:
            details["payload"] = json.loads(output) if output else {}
        except json.JSONDecodeError:
            details["raw_output"] = output
            report.add(name, "fail", "invalid JSON output", **details)
            return
    else:
        details["output"] = output
    report.add(name, "ok" if rc == 0 else failure_status, "passed" if rc == 0 else "failed", **details)


def main() -> int:
    load_env_file(ENV_FILE)

    parser = argparse.ArgumentParser(description="Unified AW-RUS health orchestrator")
    parser.add_argument("--aw-server", default=env("AW_SERVER_URL", "http://127.0.0.1:5600"))
    parser.add_argument(
        "--worktime-api",
        default=env("AW_RUS_HEALTH_WORKTIME_API", env("AW_WORKTIME_REPORT_BASE", "http://127.0.0.1:5610")),
    )
    parser.add_argument("--rdp-host", default=env("AW_MONITORED_WINDOWS_HOST", "192.168.100.18"))
    parser.add_argument("--rdp-hostname", default=env("AW_MONITORED_WINDOWS_HOSTNAME", "SHARKON2025"))
    parser.add_argument("--state-dir", default=env("AW_RUS_HEALTH_STATE_DIR", "/var/lib/activitywatch/health"))
    parser.add_argument("--validation-dir", default=env("AW_RUS_HEALTH_VALIDATION_DIR", "/var/lib/activitywatch/health/windows-validation"))
    parser.add_argument("--session-max-age-seconds", type=int, default=int(env("AW_RUS_HEALTH_SESSION_MAX_AGE_SECONDS", "900")))
    parser.add_argument("--interactive-max-age-seconds", type=int, default=int(env("AW_RUS_HEALTH_INTERACTIVE_MAX_AGE_SECONDS", "900")))
    parser.add_argument("--session-events-max-age-seconds", type=int, default=int(env("AW_RUS_HEALTH_SESSION_EVENTS_MAX_AGE_SECONDS", "86400")))
    parser.add_argument("--guard-max-age-seconds", type=int, default=int(env("AW_RUS_HEALTH_GUARD_MAX_AGE_SECONDS", "300")))
    parser.add_argument("--guard-required", action="store_true", default=env("AW_RUS_HEALTH_GUARD_REQUIRED", "0").lower() in ("1", "true", "yes"))
    parser.add_argument("--validation-max-age-seconds", type=int, default=int(env("AW_RUS_HEALTH_VALIDATION_MAX_AGE_SECONDS", "259200")))
    parser.add_argument("--tcp-timeout-seconds", type=float, default=float(env("AW_RUS_HEALTH_TCP_TIMEOUT_SECONDS", "3")))
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    report = Report()
    aw_api_base = args.aw_server.rstrip("/")
    if not aw_api_base.endswith("/api/0"):
        aw_api_base = aw_api_base.rstrip("/") + "/api/0"

    check_wrapper(report, "wrapper:aw-health-check", ["/usr/local/bin/aw-health-check"], failure_status="warn")
    check_wrapper(
        report,
        "wrapper:dlp-health-check",
        ["/usr/local/bin/dlp-health-check", "--json"],
        json_mode=True,
        failure_status="warn",
    )

    try:
        info = http_json(f"{aw_api_base}/info")
        report.add("http:aw-server", "ok", "activitywatch API responded", version=info.get("version"))
    except Exception as exc:
        report.add("http:aw-server", "fail", f"activitywatch API failed: {exc}", url=f"{aw_api_base}/info")

    try:
        payload = http_json(args.worktime_api.rstrip("/") + "/health")
        report.add("http:worktime-api", "ok", "worktime API responded", payload=payload if isinstance(payload, dict) else {})
    except Exception as exc:
        report.add("http:worktime-api", "fail", f"worktime API failed: {exc}", url=args.worktime_api)

    for port, label in ((5985, "winrm"), (3389, "rdp")):
        ok, message = tcp_connect(args.rdp_host, port, args.tcp_timeout_seconds)
        report.add(f"tcp:{label}", "ok" if ok else "fail", message if ok else f"unreachable: {message}", host=args.rdp_host, port=port)

    try:
        buckets = http_json(f"{aw_api_base}/buckets")
        if not isinstance(buckets, dict):
            raise RuntimeError("bucket index is not a dict")
        report.add("aw:buckets-index", "ok", "bucket index loaded", total=len(buckets))
    except Exception as exc:
        report.add("aw:buckets-index", "fail", f"failed to load bucket index: {exc}")
        buckets = {}

    host = args.rdp_hostname
    guard_status, guard_summary, guard_details = guard_bucket_health(
        aw_api_base,
        host,
        args.guard_max_age_seconds,
        args.guard_required,
    )
    report.add("bucket:collector-guard", guard_status, guard_summary, **guard_details)

    worktime_bucket = f"aw-worktime-sessions_{host}"
    worktime_event = None
    if buckets:
        try:
            worktime_event = latest_bucket_event(aw_api_base, worktime_bucket)
        except Exception:
            worktime_event = None
    activity = host_activity_from_worktime(worktime_event, args.session_max_age_seconds)
    if worktime_event:
        status, summary, details = bucket_health(
            aw_api_base,
            worktime_bucket,
            args.session_max_age_seconds,
            missing_status="fail",
            stale_status="fail",
        )
        details["host_activity"] = activity
        report.add("bucket:worktime-sessions", status, summary, **details)
    else:
        report.add("bucket:worktime-sessions", "fail", "no events", bucket=worktime_bucket, host_activity=activity)

    interactive_required = bool(activity["active"])
    for bucket_name, label in (
        ("aw-watcher-afk", "bucket:afk"),
        ("aw-watcher-window", "bucket:window"),
        ("aw-dlp-endpoint-signals", "bucket:endpoint-signals"),
    ):
        status, summary, details = bucket_health(
            aw_api_base,
            f"{bucket_name}_{host}",
            args.interactive_max_age_seconds,
            missing_status="fail" if interactive_required else "warn",
            stale_status="fail" if interactive_required else "warn",
        )
        details["interactive_required"] = interactive_required
        details["host_activity"] = activity
        if not interactive_required and status != "ok":
            details["inactive_summary"] = summary
            status = "ok"
            summary = "inactive: no active interactive users"
        report.add(label, status, summary, **details)

    session_status, session_summary, session_details = bucket_timestamp_health(
        aw_api_base,
        buckets,
        f"aw-session-events_{host}",
        args.session_events_max_age_seconds,
        missing_status="fail",
        stale_status="warn",
    )
    if session_status == "warn" and session_details.get("age_seconds") is not None:
        session_status = "ok"
        session_summary = f"event-driven ({session_details['age_seconds']}s since last logon marker)"
    report.add("bucket:session-events", session_status, session_summary, **session_details)

    validation_dir = Path(args.validation_dir)
    validation_report = latest_validation_report(validation_dir)
    if validation_report is None:
        report.add("validation:windows", "warn", "no validation report snapshot", directory=str(validation_dir))
    else:
        try:
            payload = json.loads(validation_report.read_text(encoding="utf-8-sig"))
            age = age_seconds(datetime.fromtimestamp(validation_report.stat().st_mtime, tz=UTC), now_utc())
            if age is not None and age > args.validation_max_age_seconds:
                report.add(
                    "validation:windows",
                    "warn",
                    f"validation snapshot is stale ({age}s)",
                    path=str(validation_report),
                    overall_ok=payload.get("overallOk"),
                    failed_sections=payload.get("summary", {}).get("failedSections", []),
                )
            elif payload.get("overallOk") is True:
                report.add("validation:windows", "ok", "validation snapshot OK", path=str(validation_report), age_seconds=age)
            else:
                report.add(
                    "validation:windows",
                    "fail",
                    "validation snapshot reports failure",
                    path=str(validation_report),
                    age_seconds=age,
                    failed_sections=payload.get("summary", {}).get("failedSections", []),
                )
        except Exception as exc:
            report.add("validation:windows", "fail", f"invalid validation snapshot: {exc}", path=str(validation_report))

    payload = report.as_dict()
    state_dir = Path(args.state_dir)
    health_json_path = state_dir / "aw-rus-health.json"
    health_txt_path = state_dir / "aw-rus-health.txt"
    write_atomic(health_json_path, json.dumps(payload, ensure_ascii=False, indent=2) + "\n")
    write_atomic(health_txt_path, report.render_text() + "\n")
    chmod_if_possible(health_json_path, 0o644)
    chmod_if_possible(health_txt_path, 0o644)

    if args.json:
        print(json.dumps(payload, ensure_ascii=False, indent=2))
    else:
        print(report.render_text())
    return 0 if payload["ok"] else 1


if __name__ == "__main__":
    sys.exit(main())