File

from __future__ import annotations

import argparse

import datetime as dt

import hashlib

import json

import time

from dataclasses import dataclass, field

from typing import Dict, List, Tuple

REVISION = "20260902-01-SCHEDV2-SHADOW"

TERMINAL = {"done", "error", "cancelled", "expired", "superseded"}

COALESCABLE = {"health", "status", "job_status", "jobs", "recent", "git_status"}

def utc_now() -> dt.datetime:

    return dt.datetime.now(dt.timezone.utc)

@dataclass

class Client:

    client_id: str

    weight: int = 1

    max_inflight: int = 1

    inflight: int = 0

    last_served_tick: int = 0

@dataclass

class Command:

    command_id: str

    client_id: str

    seq: int

    created_at: dt.datetime

    priority: int = 50

    resource_keys: Tuple[str, ...] = field(default_factory=tuple)

    action: str = "health"

    payload_json: str = "{}"

    status: str = "pending"

    ttl_seconds: int = 300

    coalesce_key: str = ""

    supersedes_command_id: str = ""

@dataclass

class Lock:

    resource_key: str

    command_id: str

    client_id: str

    fence_token: str

    lease_until: dt.datetime

class Scheduler:

    def __init__(self, clients: Dict[str, Client], lease_seconds: int = 60, aging_seconds: int = 30):

        self.clients = clients

        self.lease_seconds = lease_seconds

        self.aging_seconds = aging_seconds

        self.tick = 0

        self.locks: Dict[str, Lock] = {}

    def expire_and_coalesce(self, cmds: List[Command], now: dt.datetime | None = None) -> None:

        now = now or utc_now()

        by_id = {c.command_id: c for c in cmds}

        latest: Dict[Tuple[str, str], Command] = {}

        for c in cmds:

            if c.status in TERMINAL:

                continue

            if c.ttl_seconds and (now - c.created_at).total_seconds() > c.ttl_seconds:

                c.status = "expired"

                continue

            if c.supersedes_command_id:

                superseded = by_id.get(c.supersedes_command_id)

                if superseded and superseded.status == "pending":

                    superseded.status = "superseded"

            if c.action in COALESCABLE and c.coalesce_key:

                key = (c.client_id, c.coalesce_key)

                prev = latest.get(key)

                if prev is None:

                    latest[key] = c

                elif (c.seq, c.created_at) > (prev.seq, prev.created_at):

                    prev.status = "superseded"

                    latest[key] = c

                else:

                    c.status = "superseded"

    def resources_free(self, c: Command, now: dt.datetime) -> bool:

        for resource in c.resource_keys:

            lock = self.locks.get(resource)

            if lock and lock.lease_until > now and lock.command_id != c.command_id:

                return False

        return True

    def effective_priority(self, c: Command, now: dt.datetime) -> int:

        age = max(0.0, (now - c.created_at).total_seconds())

        return c.priority + int(age / self.aging_seconds)

    def choose(self, cmds: List[Command], now: dt.datetime | None = None) -> Command | None:

        now = now or utc_now()

        self.tick += 1

        self.expire_and_coalesce(cmds, now)

        by_client: Dict[str, List[Command]] = {}

        for c in cmds:

            if c.status != "pending":

                continue

            client = self.clients.get(c.client_id)

            if client is None or client.inflight >= client.max_inflight:

                continue

            if not self.resources_free(c, now):

                continue

            by_client.setdefault(c.client_id, []).append(c)

        if not by_client:

            return None

        candidates = []

        for client_id, arr in by_client.items():

            client = self.clients[client_id]

            best = max(arr, key=lambda c: (self.effective_priority(c, now), -c.seq))

            fairness_age = self.tick - client.last_served_tick

            score = (self.effective_priority(best, now), fairness_age * max(1, client.weight))

            candidates.append((score, client_id, best))

        _, client_id, chosen = max(candidates, key=lambda item: (item[0][0], item[0][1]))

        self.clients[client_id].last_served_tick = self.tick

        return chosen

    def acquire(self, c: Command, now: dt.datetime | None = None) -> Tuple[str, dt.datetime]:

        now = now or utc_now()

        if not self.resources_free(c, now):

            raise RuntimeError("resource_locked")

        token = hashlib.sha256(

            f"{c.command_id}|{c.client_id}|{time.time_ns()}".encode()

        ).hexdigest()[:16]

        lease_until = now + dt.timedelta(seconds=self.lease_seconds)

        for resource in c.resource_keys:

            self.locks[resource] = Lock(resource, c.command_id, c.client_id, token, lease_until)

        self.clients[c.client_id].inflight += 1

        return token, lease_until

    def release(self, command_id: str, client_id: str, fence_token: str) -> None:

        for resource, lock in list(self.locks.items()):

            if lock.command_id == command_id and lock.fence_token == fence_token:

                self.locks.pop(resource, None)

        client = self.clients.get(client_id)

        if client:

            client.inflight = max(0, client.inflight - 1)

def self_test() -> None:

    now = utc_now()

    clients = {

        "chatgpt-main": Client("chatgpt-main", weight=2, max_inflight=2),

        "chatgpt-work": Client("chatgpt-work", weight=1, max_inflight=1),

        "codex": Client("codex", weight=1, max_inflight=2),

    }

    scheduler = Scheduler(clients, lease_seconds=60, aging_seconds=10)

    cmds = [

        Command("a1", "chatgpt-main", 1, now - dt.timedelta(seconds=5), 50, ("apps-script",), "health", coalesce_key="status"),

        Command("a2", "chatgpt-main", 2, now, 50, ("apps-script",), "health", coalesce_key="status"),

        Command("b1", "chatgpt-work", 1, now, 100, ("android-build",), "health"),

        Command("c1", "codex", 1, now, 20, ("apps-script",), "job_status", ttl_seconds=1),

        Command("c2", "codex", 2, now - dt.timedelta(seconds=20), 20, ("file:project/x",), "git_status", ttl_seconds=5),

    ]

    scheduler.expire_and_coalesce(cmds, now)

    assert cmds[0].status == "superseded", cmds[0].status

    assert cmds[4].status == "expired", cmds[4].status

    first = scheduler.choose(cmds, now)

    assert first and first.command_id == "b1", first

    token, _ = scheduler.acquire(first, now)

    second = scheduler.choose(cmds, now)

    assert second and second.command_id == "a2", second

    scheduler.release(first.command_id, first.client_id, token)

    assert clients["chatgpt-work"].inflight == 0

    print(json.dumps({

        "ok": True,

        "revision": REVISION,

        "first": first.command_id,

        "second": second.command_id,

        "states": {c.command_id: c.status for c in cmds},

        "execution_enabled": False,

    }, separators=(",", ":")))

def main() -> None:

    parser = argparse.ArgumentParser()

    parser.add_argument("--self-test", action="store_true")

    parser.add_argument("--mode", default="shadow", choices=["shadow"])

    args = parser.parse_args()

    if args.self_test:

        self_test()

        return

    print(json.dumps({

        "ok": True,

        "revision": REVISION,

        "mode": args.mode,

        "execution_enabled": False,

        "note": "Engine only. Live Sheet adapter not enabled yet.",

    }, separators=(",", ":")))

if __name__ == "__main__":

    main()