"""Persistent host coordinator. Control/telemetry I/O never runs in the simulator.

The Windows task owns the coordinator; a Job Object fences its native children.
Closing the viewer or the private control tunnel cannot stop the physics clock.
Runs are bounded locally; reconnect reconciles the same identity, without replay.
"""

import argparse
import ctypes
import hashlib
import json
import logging
import os
import subprocess
import threading
import time
from pathlib import Path
from uuid import uuid4

from core_client import CoreClient
from local_state import StateChannel, read_json, write_json
from model_stack import ROOT, ModelStack, docker, sha256
from worker import terminate_episode


class NativeJob:
    """Windows kills all native descendants if the coordinator unexpectedly exits."""

    def __init__(self):
        from ctypes import wintypes

        self.kernel = ctypes.WinDLL("kernel32", use_last_error=True)
        self.kernel.CreateJobObjectW.restype = wintypes.HANDLE
        self.kernel.CreateJobObjectW.argtypes = [ctypes.c_void_p, wintypes.LPCWSTR]
        self.kernel.SetInformationJobObject.argtypes = [
            wintypes.HANDLE,
            ctypes.c_int,
            ctypes.c_void_p,
            wintypes.DWORD,
        ]
        self.kernel.AssignProcessToJobObject.argtypes = [wintypes.HANDLE, wintypes.HANDLE]
        self.kernel.CloseHandle.argtypes = [wintypes.HANDLE]
        self.handle = self.kernel.CreateJobObjectW(None, None)
        # JOBOBJECT_EXTENDED_LIMIT_INFORMATION: 144 bytes on 64-bit Windows;
        # BasicLimitInformation.LimitFlags is DWORD at offset 16.
        limits = ctypes.create_string_buffer(144)
        ctypes.c_uint32.from_buffer(limits, 16).value = 0x2000  # KILL_ON_JOB_CLOSE
        if not self.handle or not self.kernel.SetInformationJobObject(self.handle, 9, limits, 144):
            raise ctypes.WinError(ctypes.get_last_error())

    def assign(self, child):
        if not self.kernel.AssignProcessToJobObject(self.handle, int(child._handle)):
            terminate_episode(child)
            raise ctypes.WinError(ctypes.get_last_error())

    def close(self):
        if self.handle:
            self.kernel.CloseHandle(self.handle)
            self.handle = None


def acquire_service_lock(path):
    import msvcrt

    stream = path.open("a+b")
    stream.seek(0)
    if not stream.read(1):
        stream.write(b"0")
        stream.flush()
    stream.seek(0)
    msvcrt.locking(stream.fileno(), msvcrt.LK_NBLCK, 1)
    return stream


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--core", default="http://127.0.0.1:18080")
    parser.add_argument("--token-file", type=Path, required=True)
    parser.add_argument("--state", type=Path, required=True)
    parser.add_argument("--isaac", type=Path, required=True)
    parser.add_argument("--stream-address", required=True)
    args = parser.parse_args()
    args.state.mkdir(parents=True, exist_ok=True)
    lock = acquire_service_lock(args.state / "realtime-service.lock")
    logging.basicConfig(
        filename=args.state / "realtime-service.log",
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
    )
    identity_path = args.state / "realtime-identity.json"
    identity = read_json(identity_path)
    release = sha256(ROOT / "realtime_worker.py")
    active_path = args.state / "active.json"
    if identity is None or (identity["release"] != release and not active_path.exists()):
        identity = {"instance_id": uuid4().hex, "release": release}
        write_json(identity_path, identity)
    instance = identity["instance_id"]
    client = CoreClient(args.core, args.token_file, instance)
    stack = ModelStack()
    stack.preflight()
    sources = {
        "worker": release,
        "scene": hashlib.sha256(
            "".join(
                sha256(ROOT / n)
                for n in (
                    "run_realtime.py",
                    "motion_control.py",
                    "realtime_ai.py",
                    "local_state.py",
                    "terrain.py",
                    "spawn_clearance.py",
                    "prepare_terrain.py",
                    "navigation_client.py",
                    "navigation/server.py",
                    "navigation/footprint.py",
                    "navigation/fastdds.xml",
                    "navigation/Prepare-Terrain.ps1",
                )
            ).encode()
            + "".join(
                sha256(ROOT.parents[1] / "src/k1link/simulation/ai_polygon" / n)
                for n in (
                    "inference.py",
                    "policy.py",
                    "contracts.py",
                    "composition.py",
                    "terrain_contract.py",
                )
            ).encode()
            + "".join(
                sha256(ROOT.parents[1] / "src/k1link" / n)
                for n in (
                    "artifacts.py",
                    "observatory/__init__.py",
                    "observatory/modular_composition.py",
                    "simulation/__init__.py",
                )
            ).encode()
        ).hexdigest(),
        "models": hashlib.sha256(
            "".join(
                sha256(ROOT / name)
                for name in (
                    "models.worker-006.json",
                    "model_stack.py",
                    "compose.models.yaml",
                    "ddrnet_server.py",
                    "segformer/server.py",
                )
            ).encode()
        ).hexdigest(),
        "robot": sha256(ROOT / "rover_profile.py"),
    }
    hello = {
        "worker_id": "worker-006-ai-polygon",
        "instance_id": instance,
        "runtime": "isaac-sim-6.1",
        "execution_modes": ["realtime"],
        "model_ids": [m["id"] for m in stack.profile["models"]]
        + [stack.profile["navigation"]["id"]],
        "profile_sha256": hashlib.sha256(json.dumps(sources, sort_keys=True).encode()).hexdigest(),
        "runtime_sources": sources,
        "stream": {
            "server": args.stream_address,
            "signaling_port": 49100,
            "media_port": 47998,
            "width": 1280,
            "height": 720,
            "fps": 30,
        },
    }
    write_json(
        args.state / "service-status.json",
        {"pid": os.getpid(), "instance_id": instance, "hello": hello},
    )

    def connect():
        while True:
            try:
                client.request("/worker/register", hello)
                return
            except Exception as exc:
                logging.warning("Control unavailable: %s", type(exc).__name__)
                time.sleep(3)

    # A previous process died. Its Job Object has killed native children.
    # Reconcile only exact persisted container IDs; never infer release from a lease.
    orphan = read_json(active_path)
    if orphan:
        if orphan.get("instance_id") != instance:
            raise RuntimeError("Reconciliation identity mismatch")
        for cid in orphan.get("model_container_ids", []):
            record = docker("container", "inspect", cid, check=False)
            if record.returncode == 0:
                data = json.loads(record.stdout)[0]
                if data["Config"]["Labels"].get("com.nodedc.stack") != "ai-polygon":
                    raise RuntimeError("Reconciliation resource owner changed")
                stack.ids.append(cid)
        stack.stop()
        connect()
        client.request(
            "/worker/runs/" + orphan["run_id"] + "/finish",
            {
                "instance_id": instance,
                "outcome": "failed",
                "resources_released": True,
                "message": "Worker перезапустился. Прогон завершён без повторного запуска.",
            },
        )
        active_path.unlink()
    connect()
    while True:
        try:
            polled = client.request("/worker/poll", {"instance_id": instance})
        except Exception:
            connect()
            continue
        if polled["action"] == "idle":
            time.sleep(0.5)
            continue
        run = polled["run"]
        if run.get("clock") != "realtime":
            raise RuntimeError("A realtime worker cannot execute lockstep work")
        episode = args.state / run["run_id"]
        if episode.exists():
            raise RuntimeError("An existing episode must never be replayed")
        episode.mkdir()
        run_file = episode / "run.json"
        write_json(run_file, run)
        channel = StateChannel(episode)
        channel.write("control", run)
        active = {"run_id": run["run_id"], "instance_id": instance, "model_container_ids": []}
        write_json(active_path, active)
        child = job = starter = None
        cancelled = threading.Event()
        model_failure = []
        result = {"outcome": "failed", "message": "Прогон прерван. Журнал сохранён на Worker."}
        try:
            cache = args.state / "worlds"
            cache.mkdir(exist_ok=True)
            source = cache / (run["world"]["sha256"] + ".ply")
            client.download(run["world"], source)
            client.request("/worker/runs/" + run["run_id"] + "/progress", {"phase": "scene"})
            from prepare_terrain import prepare_terrain

            job = NativeJob()

            def preparation_pulse(run=run):
                desired = client.request(
                    "/worker/poll", {"instance_id": instance, "run_id": run["run_id"]}
                )
                if desired["run"]["control"] == "stop":
                    raise InterruptedError("Terrain preparation cancelled")

            terrain_path = prepare_terrain(
                args.state.parent, episode, run["world"], job, preparation_pulse
            )
            run["terrain_manifest"] = str(terrain_path)
            write_json(run_file, run)
            with (episode / "isaac.log").open("wb") as output:
                child = subprocess.Popen(
                    [
                        str(args.isaac / "python.bat"),
                        str(ROOT / "run_realtime.py"),
                        "--run",
                        str(run_file),
                        "--source",
                        str(source),
                        "--stream-address",
                        args.stream_address,
                    ],
                    stdout=output,
                    stderr=subprocess.STDOUT,
                )
                job.assign(child)

                def acquired(ids, active=active):
                    active["model_container_ids"] = ids
                    write_json(active_path, active)

                def start_models(
                    cancelled=cancelled, model_failure=model_failure, acquired=acquired, run=run
                ):
                    try:
                        stack.start(
                            cancelled=cancelled.is_set,
                            on_acquired=acquired,
                            selection=run["request"].get("composition"),
                        )
                    except InterruptedError:
                        pass
                    except Exception as exc:
                        model_failure.append(type(exc).__name__)
                        logging.exception("Model startup failed")

                stop_deadline = None
                deadline = time.monotonic() + 300 + run["request"]["duration_seconds"]
                last_snapshot = -1
                registered = True
                while child.poll() is None:
                    now = time.monotonic()
                    if now >= deadline or (stop_deadline and now >= stop_deadline):
                        terminate_episode(child)
                        break
                    if model_failure:
                        raise RuntimeError("Local model startup failed")
                    try:
                        if not registered:
                            client.request("/worker/register", hello)
                            registered = True
                        polled = client.request(
                            "/worker/poll", {"instance_id": instance, "run_id": run["run_id"]}
                        )
                        desired = polled["run"]
                        channel.write("control", desired)
                        if desired["control"] == "stop":
                            cancelled.set()
                            stop_deadline = stop_deadline or now + 15
                        if desired["control"] == "play" and starter is None:
                            starter = threading.Thread(
                                target=start_models, name="polygon-model-start", daemon=True
                            )
                            starter.start()
                        snapshot = channel.read("snapshot")
                        if snapshot and snapshot["sequence"] > last_snapshot:
                            client.request("/worker/runs/" + run["run_id"] + "/snapshot", snapshot)
                            last_snapshot = snapshot["sequence"]
                    except Exception as exc:
                        registered = False
                        logging.warning(
                            "Episode control/telemetry disconnected: %s", type(exc).__name__
                        )
                    time.sleep(0.25)
                result = read_json(episode / "result.json", result)
                if cancelled.is_set():
                    result = {"outcome": "stopped", "message": "Движение и inference остановлены."}
        except InterruptedError:
            result = {"outcome": "stopped", "message": "Подготовка симуляции остановлена."}
        except Exception:
            logging.exception("Episode failed")
        finally:
            channel.close()
            cancelled.set()
            if child is not None and child.poll() is None:
                terminate_episode(child)
            if job is not None:
                job.close()
            if starter is not None:
                starter.join(timeout=150)
                if starter.is_alive():
                    raise RuntimeError("Model startup has not released ownership")
            stack.stop()
            active["resources_released"] = True
            write_json(active_path, active)
            write_json(episode / "worker-result.json", result)
        while True:
            try:
                client.request("/worker/register", hello)
                client.request(
                    "/worker/runs/" + run["run_id"] + "/finish",
                    {**result, "instance_id": instance, "resources_released": True},
                )
                active_path.unlink()
                break
            except Exception:
                logging.warning("Finish pending reconciliation")
                time.sleep(3)
    lock.close()


if __name__ == "__main__":
    try:
        main()
    except Exception:
        logging.exception("Simulation coordinator failed")
        raise