#!/usr/bin/env python3
"""SHiRE distributed worker capability registry and heartbeat provider."""

from __future__ import annotations

import json
import os
import socket
import subprocess
import threading
import time
import urllib.error
import urllib.request
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


CAPABILITY_SCHEMA = "shire.worker.capabilities.v1"
HEARTBEAT_SCHEMA = "shire.worker.heartbeat.v1"
NODE_ROLE = "distributed_compute_worker"
WORKER_VERSION = "0.1.0-capability-registry"

PROJECT_OWNERSHIP = {
    "owns_projects": False,
    "owns_approvals": False,
    "owns_memory": False,
    "owns_shirevault": False,
    "makes_safety_decisions": False,
    "accepts_jobs_from_command_core": True,
    "returns_results_for_command_core_validation": True,
}


class WorkerManager:
    """Report worker identity, runtime health and registered capabilities."""

    def __init__(
        self,
        *,
        node_id: str,
        ollama_url: str,
        fast_model: str,
        deep_model: str,
        parametric_python: str,
        cache_seconds: int = 20,
    ) -> None:
        self.node_id = node_id
        self.ollama_url = ollama_url.rstrip("/")
        self.fast_model = fast_model
        self.deep_model = deep_model
        self.parametric_python = str(Path(parametric_python))
        self.cache_seconds = max(5, int(cache_seconds))

        self.boot_id = str(uuid.uuid4())
        self.started_monotonic = time.monotonic()
        self.started_utc = self._utc_now()
        self._heartbeat_sequence = 0

        self._lock = threading.Lock()
        self._cached_probe: dict[str, Any] | None = None
        self._cached_at = 0.0

    @staticmethod
    def _utc_now() -> str:
        return datetime.now(timezone.utc).isoformat()

    @staticmethod
    def _tailscale_ip() -> str:
        try:
            result = subprocess.run(
                ["tailscale", "ip", "-4"],
                check=False,
                capture_output=True,
                text=True,
                timeout=5,
            )
        except Exception:
            return ""

        for line in result.stdout.splitlines():
            candidate = line.strip()
            if candidate:
                return candidate

        return ""

    def _probe_ollama(self) -> dict[str, Any]:
        request = urllib.request.Request(
            self.ollama_url + "/api/tags",
            method="GET",
        )

        started = time.monotonic()

        try:
            with urllib.request.urlopen(request, timeout=4) as response:
                payload = json.loads(response.read().decode("utf-8"))

            models = sorted(
                item.get("name", "")
                for item in payload.get("models", [])
                if item.get("name")
            )

            return {
                "available": True,
                "latency_ms": round(
                    (time.monotonic() - started) * 1000,
                    2,
                ),
                "models": models,
                "fast_model_present": self.fast_model in models,
                "deep_model_present": self.deep_model in models,
                "error": None,
            }
        except (
            OSError,
            ValueError,
            urllib.error.HTTPError,
            urllib.error.URLError,
        ) as exc:
            return {
                "available": False,
                "latency_ms": round(
                    (time.monotonic() - started) * 1000,
                    2,
                ),
                "models": [],
                "fast_model_present": False,
                "deep_model_present": False,
                "error": f"{type(exc).__name__}: {exc}",
            }

    def _probe_parametric_runtime(self) -> dict[str, Any]:
        probe_script = r'''
import importlib
import importlib.metadata
import json
import sys

groups = {
    "cadquery": ["cadquery", "OCP"],
    "image_analysis": ["PIL", "cv2", "skimage", "numpy"],
    "mesh_processing": [
        "trimesh",
        "meshio",
        "manifold3d",
        "shapely",
        "networkx",
        "numpy",
        "scipy",
    ],
}

result = {
    "python": sys.executable,
    "groups": {},
    "versions": {},
}

for group_name, modules in groups.items():
    group = {
        "available": True,
        "modules": {},
    }

    for module_name in modules:
        try:
            module = importlib.import_module(module_name)
            group["modules"][module_name] = {
                "available": True,
                "version": getattr(module, "__version__", None),
                "error": None,
            }
        except Exception as exc:
            group["available"] = False
            group["modules"][module_name] = {
                "available": False,
                "version": None,
                "error": f"{type(exc).__name__}: {exc}",
            }

    result["groups"][group_name] = group

for distribution in (
    "cadquery",
    "cadquery-ocp",
    "trimesh",
    "meshio",
    "manifold3d",
    "shapely",
    "opencv-python-headless",
    "scikit-image",
):
    try:
        result["versions"][distribution] = (
            importlib.metadata.version(distribution)
        )
    except importlib.metadata.PackageNotFoundError:
        result["versions"][distribution] = None

print(json.dumps(result))
'''

        started = time.monotonic()

        try:
            completed = subprocess.run(
                [self.parametric_python, "-c", probe_script],
                check=False,
                capture_output=True,
                text=True,
                timeout=30,
            )

            if completed.returncode != 0:
                raise RuntimeError(
                    completed.stderr.strip()
                    or f"probe exited {completed.returncode}"
                )

            payload = json.loads(completed.stdout)

            payload["available"] = all(
                group.get("available", False)
                for group in payload.get("groups", {}).values()
            )
            payload["latency_ms"] = round(
                (time.monotonic() - started) * 1000,
                2,
            )
            payload["error"] = None
            return payload

        except Exception as exc:
            return {
                "available": False,
                "python": self.parametric_python,
                "groups": {},
                "versions": {},
                "latency_ms": round(
                    (time.monotonic() - started) * 1000,
                    2,
                ),
                "error": f"{type(exc).__name__}: {exc}",
            }

    def _runtime_probe(self, *, force: bool = False) -> dict[str, Any]:
        now = time.monotonic()

        with self._lock:
            if (
                not force
                and self._cached_probe is not None
                and now - self._cached_at < self.cache_seconds
            ):
                return self._cached_probe

            probe = {
                "checked_at": self._utc_now(),
                "ollama": self._probe_ollama(),
                "parametric_runtime": self._probe_parametric_runtime(),
            }

            self._cached_probe = probe
            self._cached_at = time.monotonic()
            return probe

    @staticmethod
    def _capability(
        *,
        capability_id: str,
        available: bool,
        runtime: str,
        description: str,
        execution: str,
        dependencies: list[str],
    ) -> dict[str, Any]:
        return {
            "id": capability_id,
            "registered": True,
            "available": bool(available),
            "state": "available" if available else "unavailable",
            "runtime": runtime,
            "execution": execution,
            "dependencies": dependencies,
            "description": description,
        }

    def capabilities(self, *, force_probe: bool = False) -> dict[str, Any]:
        probe = self._runtime_probe(force=force_probe)

        ollama = probe["ollama"]
        parametric = probe["parametric_runtime"]
        groups = parametric.get("groups", {})

        deep_available = bool(
            ollama.get("available")
            and ollama.get("deep_model_present")
        )
        cadquery_available = bool(
            groups.get("cadquery", {}).get("available")
        )
        image_available = bool(
            groups.get("image_analysis", {}).get("available")
        )
        mesh_available = bool(
            groups.get("mesh_processing", {}).get("available")
        )

        capabilities = [
            self._capability(
                capability_id="deep_reasoning",
                available=deep_available,
                runtime="ollama",
                description="Heavy SHiRE reasoning using the deep model.",
                execution="legacy_ask_endpoint",
                dependencies=[self.deep_model],
            ),
            self._capability(
                capability_id="cadquery_generate",
                available=cadquery_available,
                runtime="shire-parametric-forge",
                description="Generate parametric CadQuery geometry.",
                execution="registration_ready",
                dependencies=["cadquery", "cadquery-ocp"],
            ),
            self._capability(
                capability_id="cadquery_validate",
                available=cadquery_available,
                runtime="shire-parametric-forge",
                description="Validate solids, geometry and export readiness.",
                execution="registration_ready",
                dependencies=["cadquery", "cadquery-ocp"],
            ),
            self._capability(
                capability_id="academy_learning",
                available=deep_available,
                runtime="ollama",
                description="Perform compute-heavy Academy learning work.",
                execution="registration_ready",
                dependencies=[self.deep_model],
            ),
            self._capability(
                capability_id="covercanvas_theme_generation",
                available=deep_available and cadquery_available,
                runtime="ollama+shire-parametric-forge",
                description="Generate CoverCanvas themes and parametric templates.",
                execution="registration_ready",
                dependencies=[self.deep_model, "cadquery"],
            ),
            self._capability(
                capability_id="image_analysis",
                available=image_available,
                runtime="shire-parametric-forge",
                description="Analyse supplied images for engineering workflows.",
                execution="registration_ready",
                dependencies=["Pillow", "OpenCV", "scikit-image"],
            ),
            self._capability(
                capability_id="mesh_processing",
                available=mesh_available,
                runtime="shire-parametric-forge",
                description="Inspect and process mesh geometry.",
                execution="registration_ready",
                dependencies=[
                    "trimesh",
                    "meshio",
                    "manifold3d",
                    "shapely",
                ],
            ),
            self._capability(
                capability_id="evidence_packaging",
                available=True,
                runtime="stdlib",
                description="Package worker outputs and integrity evidence.",
                execution="registration_ready",
                dependencies=["hashlib", "json", "tarfile"],
            ),
        ]

        available_count = sum(
            1 for capability in capabilities
            if capability["available"]
        )

        return {
            "ok": True,
            "schema": CAPABILITY_SCHEMA,
            "worker_manager_version": WORKER_VERSION,
            "node": {
                "node_id": self.node_id,
                "host": socket.gethostname(),
                "tailnet_ip": self._tailscale_ip(),
                "role": NODE_ROLE,
                "boot_id": self.boot_id,
            },
            "ownership": dict(PROJECT_OWNERSHIP),
            "summary": {
                "registered": len(capabilities),
                "available": available_count,
                "unavailable": len(capabilities) - available_count,
            },
            "capabilities": capabilities,
            "runtime_probe": probe,
            "generated_at": self._utc_now(),
        }

    def heartbeat(self) -> dict[str, Any]:
        capability_payload = self.capabilities()

        with self._lock:
            self._heartbeat_sequence += 1
            sequence = self._heartbeat_sequence

        summary = capability_payload["summary"]

        if summary["available"] == 0:
            state = "unavailable"
        elif summary["unavailable"] > 0:
            state = "degraded"
        else:
            state = "online"

        return {
            "ok": state != "unavailable",
            "schema": HEARTBEAT_SCHEMA,
            "worker_manager_version": WORKER_VERSION,
            "node_id": self.node_id,
            "host": socket.gethostname(),
            "tailnet_ip": self._tailscale_ip(),
            "role": NODE_ROLE,
            "state": state,
            "sequence": sequence,
            "boot_id": self.boot_id,
            "started_at": self.started_utc,
            "timestamp": self._utc_now(),
            "uptime_seconds": round(
                time.monotonic() - self.started_monotonic,
                3,
            ),
            "capability_summary": summary,
            "ownership": dict(PROJECT_OWNERSHIP),
        }